Solving the klotski puzzle in OCaml

klotski.ml

(* For the original problem statement, check http://mourjo.me/klotski_problem.pdf *)

exception NotFound

type 'e rel = 'e -> 'e list
type 'e prop = 'e -> bool

type ('a, 'set) set_operations = {
  empty : 'set;              (* The empty set. *)
  mem : 'a -> 'set -> bool;  (* [mem x s = true] iff [x] is in [s]. *)
  add : 'a -> 'set -> 'set;  (* [add s x] is the set [s] union {x}. *)
}

type ('configuration, 'move) puzzle = {
  move : 'configuration -> 'move -> 'configuration;
  possible_moves : 'configuration -> 'move list;
  final : 'configuration -> bool
}

type piece_kind = S | H | V | C | X
type piece = piece_kind * int
let x = (X, 0) and s = (S, 0) and h = (H, 0)
let (c0, c1, c2, c3) = ((C, 0), (C, 1), (C, 2), (C, 3))
let (v0, v1, v2, v3) = ((V, 0), (V, 1), (V, 2), (V, 3))
let all_pieces : piece list = [ s; h; c0; c1; c2; c3; v0; v1; v2; v3 ]

type board = piece array array
let initial_board =
  [| [| v0 ; s  ; s  ; v1 |];
     [| v0 ; s  ; s  ; v1 |];
     [| v2 ; h  ; h  ; v3 |];
     [| v2 ; c0 ; c1 ; v3 |];
     [| c2 ; x  ; x  ; c3 |] |]

let initial_board_simpler =
  [| [| c2 ; s  ; s  ; c1 |] ;
     [| c0 ; s  ; s  ; c3 |] ;
     [| v1 ; v2 ; v3 ; v0 |] ;
     [| v1 ; v2 ; v3 ; v0 |] ;
     [| x  ; x  ; x  ; x  |] |]

let initial_board_trivial =
  [| [| x  ; s  ; s  ; x  |] ;
     [| x  ; s  ; s  ; x  |] ;
     [| x  ; x  ; x  ; x  |] ;
     [| x  ; x  ; x  ; x  |] ;
     [| x  ; x  ; x  ; x  |] |]

type direction = { dcol : int; drow : int; }
type move = Move of piece * direction * board
let move _ (Move (_, _, b)) = b

let rec loop p f x =
  if (p x) then x else (loop p f (f x));;

let rec exists p l =
  match l with 
  | [] -> false
  | a :: xs -> if (p a) then true else (exists p xs);; 

let rec find p l =
  match l with
  | [] -> raise NotFound
  | a :: xs -> if (p a) then a else (find p xs) ;; 

(* --- Part A: A Generic Problem Solver --- *)

let near x =
  [x-2; x-1; x; x+1; x+2];;

let rec flat_map r arg =
  List.fold_left (fun acc a -> 
      (r a)@acc) 
    [] arg;;

let rec iter_rel rel n = 
  let rec aux i acc = if i=0 then acc else (aux (i-1) (flat_map rel acc)) in
  fun item -> aux n [item];;

let rec solve r p x =
  let res = loop (exists p) (flat_map r) [x]
  in
  (find p res);;

let solve_path r p x =
  let new_r a = match a with
    | a1 :: b :: xs -> let nvs= (r (List.nth b 0)) and prec=(a1 @ b) in
        List.map (fun nv -> [prec;  [nv]]) nvs
    | _ -> []
  and
    new_p b = match b with 
    | a :: b1 :: xs -> (p (List.nth b1 0))
    | _ -> false
  in
  let r1 = (solve new_r new_p [[];[x]]) in
  match r1 with
  | a :: b :: xs -> a @ b
  | _ -> [];;


let archive_map opset r (s, l) = 
  let all_possible = (flat_map r l) in
  List.fold_left
    (fun (union_s_abnis, all_but_not_in_s) v ->
       if (opset.mem v union_s_abnis)
       then (union_s_abnis, all_but_not_in_s)
       else (opset.add v union_s_abnis),(v::all_but_not_in_s))
    (s,[]) 
    all_possible;;


let solve' opset r p x =
  let rec aux seen unseen =
    let (new_seen, new_unseen) = (archive_map opset r (seen,unseen)) in 
    begin
      try (find p new_unseen) 
      with NotFound -> (aux new_seen new_unseen)
    end
  in
  if (p x) then x else
    (aux opset.empty [x]);; 


let solve_path' opset r p x =
  let new_r a = match a with
    | a1 :: b :: xs -> let nvs= (r (List.nth b 0)) and prec=(a1 @ b) in
        List.map (fun nv -> [prec;  [nv]]) nvs
    | _ -> []
  and
    new_p b = match b with 
    | a :: b1 :: xs -> (p (List.nth b1 0))
    | _ -> false
  and
    new_opset = {empty = opset.empty; 
                 mem = (fun v s -> (opset.mem [(List.hd (List.nth v 1))] s)); 
                 add = (fun v s -> (opset.add [(List.hd (List.nth v 1))] s))}
  in let r1 = (solve' new_opset new_r new_p [[];[x]]) in
  match r1 with
  | a :: b :: xs -> a @ b
  | _ -> [];;


let solve_puzzle p opset c = 
  solve_path' opset   
    (fun cfg -> 
       List.map (fun mv -> p.move cfg mv)
         (p.possible_moves cfg))
    p.final
    c;;

(* --- Part B: A Solver for Klotski --- *)

let final board =
  let second_last_row = (Array.get board 3) and
  last_row = (Array.get board 4) in
  (Array.get second_last_row 1) = (S,0) &&
  (Array.get second_last_row 2) = (S,0) &&
  (Array.get last_row 1) = (S,0) &&
  (Array.get last_row 2) = (S,0);;



let deep_copy a =
  Array.map (fun row -> (Array.copy row)) a;;

let move_piece board1 piece { drow; dcol } =
  let board = (deep_copy board1) in 
  let get_in i j = (Array.get (Array.get board i) j) 
  and kind (a,_) = a
  and set_in (i, j) v = (Array.set (Array.get board i) j v)
  in
  let rec find_piece i j = 
    if ((i > 4) || (i < 0) || (j > 3) || (j < 0)) then (-100,-100) else
    if (get_in i j) = piece 
    then (i,j) 
    else 
      (if (j < 3) 
       then (find_piece i (j + 1)) 
       else (find_piece (i + 1) 0))
  in
  let (pos_i, pos_j) = (find_piece 0 0) 
  and any coll item = List.fold_left
      (fun acc itm -> if (itm = item) then true else acc) 
      false coll
  in 
  let my_positions = 
    if (kind piece) = S then [(pos_i, pos_j);
                              (pos_i+1, pos_j); 
                              (pos_i, pos_j+1); 
                              (pos_i+1, pos_j+1)]
    else if (kind piece) = V then [(pos_i, pos_j);
                                   (pos_i+1, pos_j)]
    else if (kind piece) = H then [(pos_i, pos_j);
                                   (pos_i, pos_j+1)]
    else [(pos_i, pos_j)]
  and my_new_positions = 
    if (kind piece) = S then [(pos_i+drow, pos_j+dcol);
                              (pos_i+1+drow, pos_j+dcol); 
                              (pos_i+drow, pos_j+1+dcol); 
                              (pos_i+1+drow, pos_j+1+dcol)]
    else if (kind piece) = V then [(pos_i+drow, pos_j+dcol);
                                   (pos_i+1+drow, pos_j+dcol)]
    else if (kind piece) = H then [(pos_i+drow, pos_j+dcol);
                                   (pos_i+drow, pos_j+1+dcol)]
    else [(pos_i+drow, pos_j+dcol)]
  in
  let new_minus_old = 
    List.fold_left (fun acc pos -> 
        if (any my_positions pos) then acc else pos::acc) 
      []
      my_new_positions
  and old_minus_new = 
    List.fold_left (fun acc pos -> 
        if (any my_new_positions pos) then acc else pos::acc) 
      []
      my_positions
  in
  let correct_bound = 
    List.fold_left (fun acc (pi,pj) -> 
        if ((pi <= 4) && (pi >= 0) && (pj <= 3) && (pj >= 0)) 
        then acc else false) 
      true 
      my_new_positions
  in
  let correct = correct_bound && 
                (List.fold_left 
                   (fun acc (pi, pj) -> 
                      if (kind(get_in pi pj))=X 
                      then acc 
                      else false)  
                   true 
                   new_minus_old)
  in 
  if (not correct) then None
  else 
    begin
      if (kind piece) = S then 
        begin
          let p1 = (get_in pos_i pos_j)
          and p2 = (get_in (pos_i+1) pos_j)
          and p3 = (get_in pos_i (pos_j+1))
          and p4 = (get_in (pos_i+1) (pos_j+1)) in
          begin 
            set_in (pos_i+drow, pos_j+dcol) p1;
            set_in (pos_i+1+drow, pos_j+dcol) p2; 
            set_in (pos_i+drow, pos_j+1+dcol) p3; 
            set_in (pos_i+1+drow, pos_j+1+dcol) p4;
          end
        end ;


      if (kind piece) = V then 
        begin
          let p1 = (get_in pos_i pos_j)
          and p2 = (get_in (pos_i+1) pos_j) in
          set_in (pos_i+drow, pos_j+dcol) p1;
          set_in (pos_i+1+drow, pos_j+dcol) p2;
        end;

      if (kind piece) = H then
        begin
          let p1 = (get_in pos_i pos_j)
          and p2 = (get_in pos_i (pos_j+1)) in
          set_in (pos_i+drow, pos_j+dcol) p1;
          set_in (pos_i+drow, pos_j+1+dcol) p2;
        end;

      if (kind piece) = C then
        begin
          let p1 = (get_in pos_i pos_j) in
          set_in (pos_i+drow, pos_j+dcol) p1;
        end;

      for i=0 to ((List.length old_minus_new)-1) 
      do
        set_in (List.nth old_minus_new i) (X,0)
      done;
      Some board
    end;;


let possible_moves b = 
  let board=(deep_copy b) in
  List.fold_left (fun acc pc -> 
      let candidates = [ { drow=1; dcol=0 };
                         { drow=(-1); dcol=0 };
                         { drow=0; dcol=1 };
                         { drow=0; dcol=(-1) }; ] in
      List.fold_left (fun acc1 dir -> 
          let res = move_piece board pc dir in
          match res with 
          | None -> acc1
          | Some b -> Move(pc, dir, b)::acc1)
        acc 
        candidates)

    [] all_pieces ;;



exception Foo of int

module BoardSet = Set.Make (struct
    type t = board
    let compare b1 b2 =
      let get_value p = 
        if p = S then 9
        else if p = H then 8
        else if p = C then 7 
        else if p = V then 6
        else 5 (* x *)
      in
      let rec aux_outer i = let row1 = b1.(i) and row2 = b2.(i) in
        let rec aux_inner j = let (p1,c1) = row1.(j) and (p2,c2) = row2.(j) in 
          let num_rows = (Array.length b1) and 
            num_cols = (Array.length row1) in
          if p1 = p2 then 
            if c1 = c2 then if (j < (num_cols-1) ) then (aux_inner (j+1)) else if (i < (num_rows-1)) then (aux_outer (i+1)) else 0
            else (raise (Foo(c1-c2)))
          else (raise (Foo ((get_value p1) - (get_value p2))))
        in 
        if ((Array.length row1)>0)
        then (aux_inner 0)
        else 0
      in
      try
        if ((Array.length b1)>0)then
          (aux_outer 0)
        else 0
      with | Foo x -> x | Invalid_argument _ -> 0;;
  end)


module BoardSet1 = Set.Make (struct
    type t = board
    let compare b1 b2 =
      let get_value p = 
        if p = S then 9
        else if p = H then 8
        else if p = C then 7 
        else if p = V then 6
        else 5 (* x *)
      in
      let num_rows=(Array.length b1) in
      let rec aux_outer i = let row1 = b1.(i) and row2 = b2.(i) in
        let num_cols=(Array.length row1) in
        let rec aux_inner j = let (p1,c1) = row1.(j) and (p2,c2) = row2.(j) in 
          if p1 = p2 then 
            if (j < (num_cols-1) ) 
            then (aux_inner (j+1)) 
            else if (i < (num_rows-1)) then (aux_outer (i+1))
            else 0
          else (raise (Foo ((get_value p1) - (get_value p2))))
        in 
        if ((Array.length row1)>0)
        then (aux_inner 0)
        else 0
      in
      try
        if ((Array.length b1)>0)then
          (aux_outer 0)
        else 0
      with | Foo x -> x | Invalid_argument _ -> 0;;
  end)


let solve_klotski initial_board =
  solve_puzzle
    { move; 
      possible_moves; 
      final; } 
    {empty = BoardSet1.empty; 
     mem = (fun llist sset -> (BoardSet1.mem (List.nth llist 0) sset)); 
     add = (fun llist sset -> (BoardSet1.add (List.nth llist 0) sset));
    } 
    initial_board ;;