mirror of
https://github.com/c-cube/sidekick.git
synced 2025-12-06 03:05:31 -05:00
384 lines
10 KiB
OCaml
384 lines
10 KiB
OCaml
(** {1 simple sudoku solver} *)
|
||
|
||
module Fmt = CCFormat
|
||
module Vec = Sidekick_util.Vec
|
||
module Log = Sidekick_util.Log
|
||
module Profile = Sidekick_util.Profile
|
||
|
||
let errorf msg = Fmt.kasprintf failwith msg
|
||
|
||
module Cell : sig
|
||
type t = private int
|
||
|
||
val equal : t -> t -> bool
|
||
val neq : t -> t -> bool
|
||
val hash : t -> int
|
||
val empty : t
|
||
val is_empty : t -> bool
|
||
val is_full : t -> bool
|
||
val make : int -> t
|
||
val pp : t Fmt.printer
|
||
end = struct
|
||
type t = int
|
||
|
||
let empty = 0
|
||
|
||
let[@inline] make i =
|
||
assert (i >= 0 && i <= 9);
|
||
i
|
||
|
||
let[@inline] is_empty x = x = 0
|
||
let[@inline] is_full x = x > 0
|
||
let hash = CCHash.int
|
||
let[@inline] equal (a : t) b = a = b
|
||
let[@inline] neq (a : t) b = a <> b
|
||
|
||
let pp out i =
|
||
if i = 0 then
|
||
Fmt.char out '.'
|
||
else
|
||
Fmt.int out i
|
||
end
|
||
|
||
module Grid : sig
|
||
type t
|
||
|
||
val get : t -> int -> int -> Cell.t
|
||
val set : t -> int -> int -> Cell.t -> t
|
||
|
||
type set = (int * int * Cell.t) Iter.t
|
||
(** A set of related cells *)
|
||
|
||
val rows : t -> set Iter.t
|
||
val cols : t -> set Iter.t
|
||
val squares : t -> set Iter.t
|
||
val all_cells : t -> (int * int * Cell.t) Iter.t
|
||
val parse : string -> t
|
||
val is_full : t -> bool
|
||
val is_valid : t -> bool
|
||
val matches : pat:t -> t -> bool
|
||
val pp : t Fmt.printer
|
||
end = struct
|
||
type t = Cell.t array
|
||
|
||
let[@inline] get (s : t) i j = s.((i * 9) + j)
|
||
|
||
let[@inline] set (s : t) i j n =
|
||
let s' = Array.copy s in
|
||
s'.((i * 9) + j) <- n;
|
||
s'
|
||
|
||
type set = (int * int * Cell.t) Iter.t
|
||
(** A set of related cells *)
|
||
|
||
open Iter.Infix
|
||
|
||
let all_cells (g : t) =
|
||
0 -- 8 >>= fun i ->
|
||
0 -- 8 >|= fun j -> i, j, get g i j
|
||
|
||
let rows (g : t) =
|
||
0 -- 8 >|= fun i ->
|
||
0 -- 8 >|= fun j -> i, j, get g i j
|
||
|
||
let cols g =
|
||
0 -- 8 >|= fun j ->
|
||
0 -- 8 >|= fun i -> i, j, get g i j
|
||
|
||
let squares g =
|
||
0 -- 2 >>= fun sq_i ->
|
||
0 -- 2 >|= fun sq_j ->
|
||
0 -- 2 >>= fun off_i ->
|
||
0 -- 2 >|= fun off_j ->
|
||
let i = (3 * sq_i) + off_i in
|
||
let j = (3 * sq_j) + off_j in
|
||
i, j, get g i j
|
||
|
||
let is_full g = Array.for_all Cell.is_full g
|
||
|
||
let is_valid g =
|
||
let all_distinct (s : set) =
|
||
s
|
||
>|= (fun (_, _, c) -> c)
|
||
|> Iter.diagonal
|
||
|> Iter.for_all (fun (c1, c2) -> Cell.neq c1 c2)
|
||
in
|
||
(Iter.for_all all_distinct @@ rows g)
|
||
&& (Iter.for_all all_distinct @@ cols g)
|
||
&& (Iter.for_all all_distinct @@ squares g)
|
||
|
||
let matches ~pat:g1 g2 : bool =
|
||
all_cells g1
|
||
|> Iter.filter (fun (_, _, c) -> Cell.is_full c)
|
||
|> Iter.for_all (fun (x, y, c) -> Cell.equal c @@ get g2 x y)
|
||
|
||
let pp out g =
|
||
Fmt.fprintf out "@[<v>";
|
||
Array.iteri
|
||
(fun i n ->
|
||
Cell.pp out n;
|
||
if i mod 9 = 8 then Fmt.fprintf out "@,")
|
||
g;
|
||
Fmt.fprintf out "@]"
|
||
|
||
let parse (s : string) : t =
|
||
if String.length s < 81 then
|
||
errorf "line is too short, expected 81 chars, not %d" (String.length s);
|
||
let a = Array.make 81 Cell.empty in
|
||
for i = 0 to 80 do
|
||
let c = String.get s i in
|
||
let n =
|
||
if c = '.' then
|
||
0
|
||
else
|
||
Char.code c - Char.code '0'
|
||
in
|
||
if n < 0 || n > 9 then errorf "invalid char %c" c;
|
||
a.(i) <- Cell.make n
|
||
done;
|
||
a
|
||
end
|
||
|
||
module B_ref = Sidekick_util.Backtrackable_ref
|
||
|
||
module Solver : sig
|
||
type t
|
||
|
||
val create : Grid.t -> t
|
||
val solve : t -> Grid.t option
|
||
end = struct
|
||
open Sidekick_sat.Solver_intf
|
||
|
||
(* formulas *)
|
||
module F = struct
|
||
type t = bool * int * int * Cell.t
|
||
|
||
let equal (sign1, x1, y1, c1) (sign2, x2, y2, c2) =
|
||
sign1 = sign2 && x1 = x2 && y1 = y2 && Cell.equal c1 c2
|
||
|
||
let hash (sign, x, y, c) =
|
||
CCHash.(combine4 (bool sign) (int x) (int y) (Cell.hash c))
|
||
|
||
let pp out (sign, x, y, c) =
|
||
Fmt.fprintf out "[@[(%d,%d) %s %a@]]" x y
|
||
(if sign then
|
||
"="
|
||
else
|
||
"!=")
|
||
Cell.pp c
|
||
|
||
let neg (sign, x, y, c) = not sign, x, y, c
|
||
|
||
let norm_sign ((sign, _, _, _) as f) =
|
||
if sign then
|
||
f, true
|
||
else
|
||
neg f, false
|
||
|
||
let make sign x y (c : Cell.t) : t = sign, x, y, c
|
||
end
|
||
|
||
type lit = F.t
|
||
|
||
module Theory = struct
|
||
type proof = unit
|
||
type proof_step = unit
|
||
|
||
module Lit = F
|
||
|
||
type lit = Lit.t
|
||
|
||
module Proof = Sidekick_sat.Proof_dummy.Make (Lit)
|
||
|
||
type t = { grid: Grid.t B_ref.t }
|
||
|
||
let create g : t = { grid = B_ref.create g }
|
||
let[@inline] grid self : Grid.t = B_ref.get self.grid
|
||
let[@inline] set_grid self g : unit = B_ref.set self.grid g
|
||
let push_level self = B_ref.push_level self.grid
|
||
let pop_levels self n = B_ref.pop_levels self.grid n
|
||
let pp_c_ = Fmt.(list ~sep:(return "@ ∨ ")) F.pp
|
||
|
||
let[@inline] logs_conflict kind c : unit =
|
||
Log.debugf 4 (fun k -> k "(@[conflict.%s@ %a@])" kind pp_c_ c)
|
||
|
||
(* check that all cells are full *)
|
||
let check_full_ (self : t) (acts : (Lit.t, proof, proof_step) acts) : unit =
|
||
Profile.with_ "check-full" @@ fun () ->
|
||
let (module A) = acts in
|
||
Grid.all_cells (grid self) (fun (x, y, c) ->
|
||
if Cell.is_empty c then (
|
||
let c =
|
||
CCList.init 9 (fun c -> F.make true x y (Cell.make (c + 1)))
|
||
in
|
||
Log.debugf 4 (fun k -> k "(@[add-clause@ %a@])" pp_c_ c);
|
||
A.add_clause ~keep:true c ()
|
||
))
|
||
|
||
(* check constraints *)
|
||
let check_ (self : t) (acts : (Lit.t, proof, proof_step) acts) : unit =
|
||
Profile.with_ "check-constraints" @@ fun () ->
|
||
Log.debugf 4 (fun k ->
|
||
k "(@[sudoku.check@ @[:g %a@]@])" Grid.pp (B_ref.get self.grid));
|
||
let (module A) = acts in
|
||
let[@inline] all_diff kind f =
|
||
let pairs =
|
||
f (grid self)
|
||
|> Iter.flat_map (fun set ->
|
||
set
|
||
|> Iter.filter (fun (_, _, c) -> Cell.is_full c)
|
||
|> Iter.diagonal)
|
||
in
|
||
pairs (fun ((x1, y1, c1), (x2, y2, c2)) ->
|
||
if Cell.equal c1 c2 then (
|
||
assert (x1 <> x2 || y1 <> y2);
|
||
let c = [ F.make false x1 y1 c1; F.make false x2 y2 c2 ] in
|
||
logs_conflict ("all-diff." ^ kind) c;
|
||
A.raise_conflict c ()
|
||
))
|
||
in
|
||
all_diff "rows" Grid.rows;
|
||
all_diff "cols" Grid.cols;
|
||
all_diff "squares" Grid.squares;
|
||
()
|
||
|
||
let trail_ (acts : (Lit.t, proof, proof_step) acts) =
|
||
let (module A) = acts in
|
||
A.iter_assumptions
|
||
|
||
(* update current grid with the given slice *)
|
||
let add_slice (self : t) (acts : (Lit.t, proof, proof_step) acts) : unit =
|
||
let (module A) = acts in
|
||
trail_ acts (function
|
||
| false, _, _, _ -> ()
|
||
| true, x, y, c ->
|
||
assert (Cell.is_full c);
|
||
let grid = grid self in
|
||
let c' = Grid.get grid x y in
|
||
if Cell.is_empty c' then
|
||
set_grid self (Grid.set grid x y c)
|
||
else if Cell.neq c c' then (
|
||
(* conflict: at most one value *)
|
||
let c = [ F.make false x y c; F.make false x y c' ] in
|
||
logs_conflict "at-most-one" c;
|
||
A.raise_conflict c ()
|
||
))
|
||
|
||
let partial_check (self : t) acts : unit =
|
||
Profile.with_ "partial-check" @@ fun () ->
|
||
Log.debugf 4 (fun k ->
|
||
k "(@[sudoku.partial-check@ :trail [@[%a@]]@])" (Fmt.list F.pp)
|
||
(trail_ acts |> Iter.to_list));
|
||
add_slice self acts;
|
||
check_ self acts
|
||
|
||
let final_check (self : t) acts : unit =
|
||
Profile.with_ "final-check" @@ fun () ->
|
||
Log.debugf 4 (fun k -> k "(@[sudoku.final-check@])");
|
||
check_full_ self acts;
|
||
check_ self acts
|
||
end
|
||
|
||
module S = Sidekick_sat.Make_cdcl_t (Theory)
|
||
|
||
type t = { grid0: Grid.t; solver: S.t }
|
||
|
||
let solve (self : t) : _ option =
|
||
Profile.with_ "sudoku.solve" @@ fun () ->
|
||
let assumptions =
|
||
Grid.all_cells self.grid0
|
||
|> Iter.filter (fun (_, _, c) -> Cell.is_full c)
|
||
|> Iter.map (fun (x, y, c) -> F.make true x y c)
|
||
|> Iter.to_rev_list
|
||
in
|
||
Log.debugf 2 (fun k ->
|
||
k "(@[sudoku.solve@ :assumptions %a@])" (Fmt.Dump.list F.pp) assumptions);
|
||
let r =
|
||
match S.solve self.solver ~assumptions with
|
||
| S.Sat _ -> Some (Theory.grid (S.theory self.solver))
|
||
| S.Unsat _ -> None
|
||
in
|
||
(* TODO: print some stats *)
|
||
r
|
||
|
||
let create g : t =
|
||
{ solver = S.create ~proof:() (Theory.create g); grid0 = g }
|
||
end
|
||
|
||
let solve_grid (g : Grid.t) : Grid.t option =
|
||
let s = Solver.create g in
|
||
Solver.solve s
|
||
|
||
module type CHRONO = sig
|
||
val pp_elapsed : Fmt.formatter -> unit
|
||
end
|
||
|
||
let chrono ~pp_time : (module CHRONO) =
|
||
let module M = struct
|
||
let start = Sys.time ()
|
||
|
||
let pp_elapsed out =
|
||
if pp_time then Fmt.fprintf out " (in %.3fs)" (Sys.time () -. start)
|
||
end in
|
||
(module M)
|
||
|
||
let solve_file ~pp_time file =
|
||
Profile.with_ "solve-file" @@ fun () ->
|
||
let open (val chrono ~pp_time) in
|
||
Format.printf "solve grids in file %S@." file;
|
||
|
||
let grids =
|
||
CCIO.with_in file CCIO.read_lines_l
|
||
|> CCList.filter_map (fun s ->
|
||
let s = String.trim s in
|
||
if s = "" then
|
||
None
|
||
else (
|
||
match Grid.parse s with
|
||
| g -> Some g
|
||
| exception e ->
|
||
errorf "cannot parse sudoku %S: %s@." s (Printexc.to_string e)
|
||
))
|
||
in
|
||
Format.printf "parsed %d grids%t@." (List.length grids) pp_elapsed;
|
||
List.iter
|
||
(fun g ->
|
||
Format.printf
|
||
"@[<v>@,#########################@,@[<2>solve grid:@ %a@]@]@." Grid.pp g;
|
||
let open (val chrono ~pp_time) in
|
||
match solve_grid g with
|
||
| None -> Format.printf "no solution%t@." pp_elapsed
|
||
| Some g' when not @@ Grid.is_full g' ->
|
||
errorf "grid %a@ is not full" Grid.pp g'
|
||
| Some g' when not @@ Grid.is_valid g' ->
|
||
errorf "grid %a@ is not valid" Grid.pp g'
|
||
| Some g' when not @@ Grid.matches ~pat:g g' ->
|
||
errorf "grid %a@ @[<2>does not match original@ %a@]" Grid.pp g' Grid.pp
|
||
g
|
||
| Some g' ->
|
||
Format.printf "@[<v>@[<2>solution%t:@ %a@]@,###################@]@."
|
||
pp_elapsed Grid.pp g')
|
||
grids;
|
||
Format.printf "@.solved %d grids%t@." (List.length grids) pp_elapsed;
|
||
()
|
||
|
||
let () =
|
||
Sidekick_tef.with_setup @@ fun () ->
|
||
Fmt.set_color_default true;
|
||
let files = ref [] in
|
||
let debug = ref 0 in
|
||
let pp_time = ref true in
|
||
let opts =
|
||
[
|
||
"--debug", Arg.Set_int debug, " debug";
|
||
"-d", Arg.Set_int debug, " debug";
|
||
"--no-time", Arg.Clear pp_time, " do not print solve time";
|
||
]
|
||
|> Arg.align
|
||
in
|
||
Arg.parse opts (fun f -> files := f :: !files) "sudoku_solve [options] <file>";
|
||
Log.set_debug !debug;
|
||
try List.iter (fun f -> solve_file ~pp_time:!pp_time f) !files
|
||
with Failure msg | Invalid_argument msg ->
|
||
Format.printf "@{<Red>Error@}:@.%s@." msg;
|
||
exit 1
|