#!/usr/bin/env python3 """Optional OR-Tools CP-SAT reference solver for the Partridge puzzle.""" from __future__ import annotations import argparse import importlib.metadata import json import os import resource import sys import tempfile import time from dataclasses import dataclass from pathlib import Path from typing import Any sys.path.insert(0, str(Path(__file__).resolve().parent)) from partridge_validator import validate # noqa: E402 @dataclass class Piece: side: int copy: int x: Any y: Any def maximum_rss_bytes() -> int: rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss # Darwin reports bytes; Linux and the other supported Unix wheels use KiB. return int(rss if sys.platform == "darwin" else rss * 1024) def minimum_edge_distance(side: int) -> int: for distance in range(1, side + 1): available = (distance * (distance + 1) // 2) ** 2 if distance * side > available: return distance raise AssertionError("every non-unit side has a limiting edge distance") def add_edge_exclusions(model: Any, start: Any, side: int, board_side: int) -> int: if side == 1: return 0 distance = minimum_edge_distance(side) forbidden = list(range(1, distance + 1)) forbidden += list( range(board_side - side - distance, board_side - side) ) for position in sorted(set(forbidden)): model.add(start != position) return len(set(forbidden)) def build_model(order: int, cp_model: Any) -> tuple[Any, list[Piece], dict[str, int]]: from ortools.util.python.sorted_interval_list import Domain board_side = order * (order + 1) // 2 model = cp_model.CpModel() pieces: list[Piece] = [] x_intervals = [] y_intervals = [] edge_exclusions = 0 for side in range(order, 0, -1): equal_size: list[Piece] = [] for copy in range(side): name = f"s{side}_{copy}" x = model.new_int_var(0, board_side - side, f"x_{name}") y = model.new_int_var(0, board_side - side, f"y_{name}") piece = Piece(side, copy, x, y) pieces.append(piece) equal_size.append(piece) x_intervals.append( model.new_fixed_size_interval_var(x, side, f"xi_{name}") ) y_intervals.append( model.new_fixed_size_interval_var(y, side, f"yi_{name}") ) edge_exclusions += add_edge_exclusions(model, x, side, board_side) edge_exclusions += add_edge_exclusions(model, y, side, board_side) # Strict lexicographic ordering of the (x, y) coordinate pairs. for before, after in zip(equal_size, equal_size[1:]): model.add( before.x * board_side + before.y < after.x * board_side + after.y ) model.add_no_overlap_2d(x_intervals, y_intervals) exact_fill_literals = 0 for axis in ("x", "y"): for line in range(board_side): contributions = [] for index, piece in enumerate(pieces): start = getattr(piece, axis) overlaps = model.new_bool_var(f"{axis}_{line}_covers_{index}") overlap_domain = Domain.from_intervals( [[line - piece.side + 1, line]] ) model.add_linear_expression_in_domain( start, overlap_domain ).only_enforce_if(overlaps) model.add_linear_expression_in_domain( start, overlap_domain.complement() ).only_enforce_if(overlaps.negated()) contributions.append(piece.side * overlaps) exact_fill_literals += 1 model.add(sum(contributions) == board_side) return model, pieces, { "pieces": len(pieces), "edge_exclusions": edge_exclusions, "exact_fill_literals": exact_fill_literals, } def model_proto(model: Any) -> Any: if hasattr(model, "proto"): return model.proto return model.Proto() def serialized_model_size(model: Any) -> int: """Return the binary proto size across old protobuf and new pybind APIs.""" proto = model_proto(model) if hasattr(proto, "SerializeToString"): return len(proto.SerializeToString()) descriptor, path = tempfile.mkstemp(suffix=".bin") os.close(descriptor) try: if not model.export_to_file(path): raise RuntimeError("OR-Tools could not export the model proto") return os.path.getsize(path) finally: os.unlink(path) def solve(order: int, workers: int, time_limit: float | None) -> dict[str, Any]: try: from ortools.sat.python import cp_model except ImportError as error: raise RuntimeError( "OR-Tools is optional and is not installed; see " "CP_SAT_REFERENCE.md for virtual-environment instructions" ) from error initial_rss = maximum_rss_bytes() build_begin = time.perf_counter() model, pieces, counts = build_model(order, cp_model) build_seconds = time.perf_counter() - build_begin proto = model_proto(model) model_stats = { **counts, "variables": len(proto.variables), "constraints": len(proto.constraints), "serialized_bytes": serialized_model_size(model), } solver = cp_model.CpSolver() solver.parameters.num_search_workers = workers solver.parameters.stop_after_first_solution = True if time_limit is not None: solver.parameters.max_time_in_seconds = time_limit solve_begin = time.perf_counter() status = solver.solve(model) solve_seconds = time.perf_counter() - solve_begin feasible = status in (cp_model.FEASIBLE, cp_model.OPTIMAL) board_side = order * (order + 1) // 2 placements = ( [ { "x": solver.value(piece.x), "y": solver.value(piece.y), "side": piece.side, } for piece in pieces ] if feasible else [] ) solution = { "order": order, "board_side": board_side, "placements": placements, } validation_begin = time.perf_counter() diagnostics = validate(solution) if feasible else ["no solution produced"] validation_seconds = time.perf_counter() - validation_begin peak_rss = maximum_rss_bytes() return { "schema": "partridge-cp-sat-reference-v1", "dependency": { "name": "ortools", "version": importlib.metadata.version("ortools"), "optional": True, }, "configuration": { "workers": workers, "available_cpus": os.cpu_count(), "first_solution": True, "time_limit_seconds": time_limit, }, "model": model_stats, "result": { "status": solver.status_name(status), "feasible": feasible, "valid": feasible and not diagnostics, "diagnostics": diagnostics, }, "timing_seconds": { "build": build_seconds, "solve": solve_seconds, "solver_wall": solver.wall_time, "validation": validation_seconds, }, "memory": { "initial_rss_bytes": initial_rss, "peak_rss_bytes": peak_rss, "peak_rss_increase_bytes": max(0, peak_rss - initial_rss), }, "solution": solution, "solver_stats": solver.response_stats(), } def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("order", type=int) parser.add_argument( "--workers", type=int, default=max(1, os.cpu_count() or 1), help="CP-SAT search workers (default: all available logical CPUs)", ) parser.add_argument("--time-limit", type=float) args = parser.parse_args() if args.order < 1: parser.error("order must be positive") if args.workers < 1: parser.error("workers must be positive") if args.time_limit is not None and args.time_limit <= 0: parser.error("time limit must be positive") try: report = solve(args.order, args.workers, args.time_limit) except RuntimeError as error: parser.exit(2, f"{error}\n") json.dump(report, sys.stdout, indent=2, sort_keys=True) sys.stdout.write("\n") return 0 if report["result"]["valid"] else 1 if __name__ == "__main__": raise SystemExit(main())