Solve and memo for Project 2 of FUN Introduction to Functional Programming in OCaml

fun_ocaml_project2.ml

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 exists p l =
  List.fold_left (fun a b -> a || p b) false l
;;

let find p l =
  try
    List.find p l
  with
  | _ -> raise NotFound
;;

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

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

(* flat_map : 'e rel -> ('e list -> 'e list) *)
let tr_map f l = List.rev @@ List.rev_map f l ;;

let flat_map r = fun l -> List.concat @@ tr_map r l ;;

let rec iter_rel rel n =
  if n > 0 then
    fun x -> flat_map rel @@ iter_rel rel (n - 1) x
  else
    fun x -> [x]
;;

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

let solve_path r p x =
  let l = loop (exists (exists p)) (flat_map (fun ys -> tr_map (fun y -> y :: ys) @@ r @@ List.hd ys)) [[x]] in
  List.rev @@ find (exists p) l
;;

(**
  9. Write a function archive_map of type
  ('a, 'set) set_operations -> 'a rel -> ('set * 'a list) -> ('set * 'a list)
  such that: archive_map opset rel (s, l) = (s', l'), where:

  - l' is the list of elements that are reachable using rel from the elements of l and which are not already in the set s.
  - s' is the union of s and the elements of l'.

  10. Use archive_map to program a new function solve' (don't forget the quote after the name) of type ('a, 'set) set_operations -> 'a rel -> 'a prop -> 'a -> 'a that explores the search space with no redundancy.
  11. Same question for solve_path' of type ('a list, 'set) set_operations -> 'a rel -> 'a prop -> 'a -> 'a list.
*)

(* archive_map : ('a, 'set) set_operations -> 'a rel -> ('set * 'a list) -> ('set * 'a list) *)
let archive_map opset r (s, l) =
  loop (fun (_,xs,_) -> xs = [])
    (fun (s,(x::xs),l') ->
      if opset.mem x s then (s,xs,l') else (opset.add x s,xs,x::l'))
    (s,flat_map r l,[]) |> (fun (s',_,l') -> (s',List.rev l'))
;;

(* solve' : ('a, 'set) set_operations -> 'a rel -> 'a prop -> 'a -> 'a *)
let solve' opset r p x =
  loop (fun (_, l) -> exists p l) (archive_map opset r) (opset.empty, [x]) |> (fun (s, l) -> find p l)
;;


(* solve_path' : ('a list, 'set) set_operations -> 'a rel -> 'a prop -> 'a -> 'a list *)
let solve_path' opset r p x =
  loop (fun (_, l) -> print_int (List.length l) ; exists (fun x -> p @@ List.hd x) l)
    (archive_map opset (fun xs -> List.hd xs |> r |> tr_map (fun x -> x :: xs)))
      (opset.empty,[[x]]) |> fun (_, l) -> find (fun x -> p @@ List.hd x) l |> List.rev
;;
(*
let solve_path' opset r p x =
  loop (fun (_, l) -> exists (fun x -> p @@ List.hd x) l)
    (archive_map opset (fun xs -> List.hd xs |> r |> tr_map (fun x -> x :: xs)))
    (opset.empty,[[x]]) |> fun (_, l) -> find (fun x -> p @@ List.hd x) l |> List.rev
;;
*)

(* solve_puzzle : ('c, 'm) puzzle -> ('c list, 's) set_operations -> 'c -> 'c list *)
let solve_puzzle p opset c =
  solve_path' opset (fun c -> p.possible_moves c |> tr_map (p.move c)) p.final c
;;

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

let final board =
  List.for_all (fun p -> p = (S, 0)) [board.(3).(1); board.(3).(2); board.(4).(1); board.(4).(1)]
;;

(* bind : 'a option -> ('a -> 'b option) -> 'b option *)
let (>>=) a f =
  match a with
  | None -> None
  | Some(x) -> f x
;;

(* let rec range s e = if s >= e then [] else s :: (range (s + 1) e) ;; *)
let range s e =
  let rec loop i acc = if i < s then acc else loop (i - 1) (i :: acc)
  in loop e []
;;

let indexes =
  List.fold_left (fun l row -> List.fold_left
    (fun l col -> (row, col) :: l) l (range 0 3)) [] (range 0 4)
;;

let find_index board piece =
  List.filter (fun (row,col) -> board.(row).(col) = piece) indexes
;;

let lookup board (row, col) =
  try
    Some (board.(row).(col))
  with
    _ -> None
;;

let update board (row, col) piece = Array.set board.(row) col piece

let flat_option l =
  let loop acc = function
        | Some(a) -> a :: acc
        | None    -> acc
  in List.rev @@ List.fold_left loop [] l
;;

(* move_piece : board -> piece -> direction -> board option *)
let move_piece board piece { drow; dcol } =
  let index = find_index board piece in
  let index' = tr_map (fun (i,j) -> (i + drow, j + dcol)) index in
  if index' <> [] &&
    (List.for_all (fun ij -> let p = lookup board ij in
                p = Some (X, 0) || p = Some piece) index') then
    let board' = Array.map Array.copy board in
    (List.iter (fun ij -> update board' ij (X,0)) index ;
      List.iter (fun ij -> update board' ij piece) index' ;
      Some board')
  else
    None
;;

(* possible_moves : board -> move list *)
let possible_moves board =
  let pd = List.concat @@ tr_map (fun p ->
      [(p,{dcol=0;drow=1});(p,{dcol=0;drow=(-1)});
       (p,{dcol=1;drow=0});(p,{dcol=(-1);drow=0})]) all_pieces in
  tr_map (fun (p,d) -> move_piece board p d >>=
             (fun b -> Some (Move (p,d,b)))) pd |> flat_option
;;

let compare_piece p1 p2 =
  match (p1,p2) with
  | ((k1,n1),(k2,n2)) when k1 = k2 -> compare n1 n2
  | ((S,_), _) -> 1
  | ((H,_), (k,n)) when k <> S -> 1
  | ((C,_), (k,n)) when k <> S && k <> H -> 1
  | ((V,_), (k,n)) when k <> S && k <> H && k <> C -> 1
  | _ -> -1
;;

module BoardSet = Set.Make (struct
    type t = board
    let compare b1 b2 =
      loop (fun (_,_,_,_,_,fin) -> fin)
        (fun (i,j,r1,r2,ans,fin) -> match (i,j) with
          | (x,y) when (x+1) > 4 && y > 3 -> (i,j,r1,r2,ans,true)
          | (_,y) when y > 3              -> (i+1,0,b1.(i+1),b2.(i+1),ans,fin)
          | _ -> let n = compare_piece r1.(j) r2.(j) in
                   if n <> 0 then (i,j,r1,r2,n,true) else (i,j+1,r1,r2,ans,fin)
        ) (0,0,b1.(0),b2.(0),0,false) |> (fun (_,_,_,_,ans,_) -> ans)
    ;;
  end)

let all_x =
  [| [| x  ; x  ; x  ; x  |] ;
     [| x  ; x  ; x  ; x  |] ;
     [| x  ; x  ; x  ; x  |] ;
     [| x  ; x  ; x  ; x  |] ;
     [| x  ; x  ; x  ; x  |] |]
;;

let make_mirror board = "aaa"

(* solve_klotski : board -> board list *)
let solve_klotski initial_board =
  let compare_piece p1 p2 =
    match (p1,p2) with
    | ((k1,n1),(k2,n2)) when k1 = k2 -> 0
    | ((S,_), _) -> 1
    | ((H,_), (k,n)) when k <> S -> 1
    | ((C,_), (k,n)) when k <> S && k <> H -> 1
    | ((V,_), (k,n)) when k <> S && k <> H && k <> C -> 1
    | _ -> -1 in
  let module BoardSet = Set.Make (struct
      type t = board
      let compare b1 b2 =
        loop (fun (_,_,_,_,_,fin) -> fin)
          (fun (i,j,r1,r2,ans,fin) -> match (i,j) with
            | (x,y) when (x+1) > 4 && y > 3 -> (i,j,r1,r2,ans,true)
            | (_,y) when y > 3              -> (i+1,0,b1.(i+1),b2.(i+1),ans,fin)
            | _ -> let n = compare_piece r1.(j) r2.(j) in
                     if n <> 0 then (i,j,r1,r2,n,true) else (i,j+1,r1,r2,ans,fin)
          ) (0,0,b1.(0),b2.(0),0,false) |> (fun (_,_,_,_,ans,_) -> ans)
      ;;
    end) in
  let p = { move = (fun _ (Move (_,_,b)) -> b) ;
            possible_moves = possible_moves ;
            final = final } in
  let opset = { empty = BoardSet.add initial_board BoardSet.empty ;
                mem = (fun al s -> BoardSet.mem (List.hd al) s) ;
                add = (fun al s -> BoardSet.add (List.hd al) s) } in
  solve_puzzle p opset initial_board
;;

let initial_board_simpler1 =
  [| [| c2 ; s  ; s  ; c1 |] ;
     [| c0 ; s  ; s  ; c3 |] ;
     [| x  ; v1 ; x  ; v0 |] ;
     [| x  ; v1 ; x  ; v0 |] ;
     [| x  ; x  ; x  ; x  |] |]
;;

let initial_board_simpler2 =
  [| [| x  ; s  ; s  ; v3 |] ;
     [| x  ; s  ; s  ; v3 |] ;
     [| v2 ; v1 ; x  ; v0 |] ;
     [| v2 ; v1 ; x  ; v0 |] ;
     [| c0 ; x  ; x  ; c1 |] |]
;;

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

let initial_board_simpler4 =
  [| [| x  ; s  ; s  ; x  |] ;
     [| x  ; s  ; s  ; x  |] ;
     [| v2 ; x  ; x  ; v0 |] ;
     [| v2 ; x  ; x  ; v0 |] ;
     [| c0 ; c1 ; c2 ; c3 |] |]
;;