lang0.ml

open OCanren
open Format

type binop =
  | Conj
  | Disj
  | Impl
[@@deriving gt ~options:{ gmap; show }]

type ('a, 'b, 'op) t =
  | True
  | Not of 'a
  | Binop of 'op * 'a * 'a
  | Var of 'b
[@@deriving gt ~options:{ gmap; show }]

type ground = (ground, GT.string, binop) t
[@@distrib] [@@deriving gt ~options:{ gmap; show }]

type logic = (logic, GT.string OCanren.logic, binop OCanren.logic) t OCanren.logic
[@@deriving gt ~options:{ gmap; show }]

module M = Fmap3 (struct
  type nonrec ('a, 'b, 'c) t = ('a, 'b, 'c) t

  let fmap x = GT.gmap t x
end)

type injected = (ground, logic) OCanren.injected

let trueo : injected = inj @@ M.distrib True
let noto x : injected = inj @@ M.distrib (Not x)
let falso = noto trueo
let conj x y : injected = inj @@ M.distrib (Binop (!!Conj, x, y))
let disj x y : injected = inj @@ M.distrib (Binop (!!Disj, x, y))
let impl x y : injected = inj @@ M.distrib (Binop (!!Impl, x, y))
let var n : injected = inj @@ M.distrib (Var n)
let rec reify eta = M.reify reify OCanren.reify OCanren.reify eta

module Env = struct
  let rec lookupo
      env
      (key : (string, _) OCanren.injected)
      (rez : (bool, _) OCanren.injected)
    =
    let open Std in
    conde
      [ env === Std.List.nil () &&& failure
      ; fresh
          (k v tl)
          (env === Std.Pair.pair k v % tl)
          (conde [ k === key &&& (rez === v); k =/= key &&& lookupo tl key rez ])
      ]
  ;;

  let of_list (xs : (string * bool) list) =
    Std.list (fun (k, v) -> Std.Pair.pair !!k !!v) xs
  ;;

  let pp_ground ppf =
    Format.pp_print_list
      ~pp_sep:(fun ppf () -> Format.fprintf ppf " ")
      (fun ppf (s, b) -> Format.fprintf ppf "(%s,%b) " s b)
      ppf
  ;;
end

let implo a b rez =
  conde [ a === !!false &&& (rez === !!true); a === !!true &&& (b === rez) ]
;;

let rec evalo env ph rez =
  conde
    [ fresh
        (l r ansl ansr)
        (ph === conj l r)
        (Std.Bool.ando ansl ansr rez)
        (evalo env l ansl)
        (evalo env r ansr)
    ; fresh
        (l r ansl ansr)
        (ph === disj l r)
        (Std.Bool.oro ansl ansr rez)
        (evalo env l ansl)
        (evalo env r ansr)
    ; fresh (p1 ans1) (ph === noto p1) (Std.Bool.noto ans1 rez) (evalo env p1 ans1)
    ; fresh
        vname
        (vname === !!"y")
        (ph === var vname)
        (defer
           (defer (defer (defer (defer (defer (defer (Env.lookupo env vname rez))))))))
    ; fresh vname (vname =/= !!"y") (ph === var vname) (defer (Env.lookupo env vname rez))
    ; ph === trueo &&& (rez === Std.Bool.truo)
    ; fresh
        (l r ansl ansr)
        (defer
           (defer
              (defer
                 (defer
                    (defer
                       (defer
                          (defer
                             (defer
                                (defer
                                   (defer
                                      (defer
                                         (defer
                                            (defer
                                               (defer
                                                  (defer
                                                     (defer
                                                        (defer
                                                           (defer
                                                              (defer
                                                                 (defer
                                                                    (defer
                                                                       (defer
                                                                          (defer
                                                                             (defer
                                                                                (defer
                                                                                   (defer
                                                                                      (defer
                                                                                         (defer
                                                                                          (
                                                                                          defer
                                                                                          (
                                                                                          defer
                                                                                          (
                                                                                          defer
                                                                                          (
                                                                                          defer
                                                                                          (
                                                                                          defer
                                                                                          (
                                                                                          ph
                                                                                          === 
                                                                                          impl
                                                                                          l
                                                                                          r
                                                                                          ))))))))))))))))))))))))))))))))))
        (defer (defer (defer (defer (defer (defer (implo ansl ansr rez)))))))
        (evalo env l ansl)
        (evalo env r ansr)
    ]
;;

let flip f a b = f b a

let get_vars_rel acc ph =
  let rec iter = function
    | Value (Not x) -> iter x
    | Value (Binop (_, l, r)) ->
      iter l;
      iter r
    | Value (Var (Value s)) -> acc := s :: !acc
    | Value (Var (OCanren.Var _)) | Var _ -> ()
    | Value True -> ()
  in
  debug_var ph (flip reify) (function
      | [ h ] ->
        iter h;
        acc := List.sort_uniq compare !acc;
        (* Format.printf
          "get_vars_rel finished with %a\n%!"
          (pp_print_list pp_print_string)
          !acc; *)
        success
      | _ -> failwith "should not happen")
;;

let test_get_vars_rel ph answer =
  let ans = ref [] in
  let _ =
    OCanren.(run q) (fun _ -> get_vars_rel ans ph) (fun x -> x#reify reify) |> Stream.take
  in
  answer = !ans
;;

let%test _ =
  test_get_vars_rel (conj (conj (var !!"b") (var !!"a")) (var !!"b")) [ "a"; "b" ]
;;

let%test _ = test_get_vars_rel (var !!"b") [ "b" ]

let rec cart xs =
  match xs with
  | [] -> [ [] ]
  | h :: tl -> List.concat_map (fun xs -> [ (h, true) :: xs; (h, false) :: xs ]) (cart tl)
;;

let simplify : injected -> injected -> goal =
 fun f1 f2 ->
  let names = ref [] in
  (* let names = List.sort_uniq compare @@ List.append (get_vars f1) (get_vars f2) in *)
  (* let e = cart names in *)
  let e = ref None in
  let rec loop i =
    (* printf "loop %d\n%!" i; *)
    let env_count =
      match !e with
      | None ->
        let () = printf "names = %a\n%!" (pp_print_list pp_print_string) !names in
        let g = cart !names in
        let env_count = List.length g in
        printf
          "envs  = %a\n%!"
          (pp_print_list ~pp_sep:(fun ppf () -> fprintf ppf "; ") Env.pp_ground)
          g;
        e := Some ((g, List.map Env.of_list g), env_count);
        env_count
      | Some ((_, _), count) -> count
    in
    if i >= env_count
    then success
    else (
      let __ () =
        Format.printf "testing env %d:" i;
        Format.printf " %a\n%!" Env.pp_ground (List.nth (fst @@ fst @@ Option.get !e) i)
      in
      let e0 = List.nth (snd @@ fst @@ Option.get !e) i in
      fresh rez (evalo e0 f1 rez) (evalo e0 f2 rez) (delay (fun () -> loop (i + 1))))
  in
  fresh () (get_vars_rel names f1) (get_vars_rel names f2) (delay (fun () -> loop 0))
;;



(**************************** *)
let run_simpl ?(n = -1) ph =
  runR Types.reify (GT.show ground) (GT.show logic) n q qh (REPR (fun q -> simplify ph q))
;;

(* let () = run_simpl ~n:20 (conj (var !!"x") (var !!"x")) *)
(* let () = run_simpl ~n:20 (noto (disj (noto (var !!"x")) (noto trueo))) *)
let () = run_simpl ~n:20 (impl (var !!"x") (var !!"y"))