(* A translation of Hemann and Friedman’s μKanren into OCaml. *)
module Index = Int32          (* The type of logic variable indices *)
module Env = Map.Make(Index)       (* maps logic variables to terms *)
type var = Var of Index.t (* the index is a counter typically from call_fresh *)
type term = Vart of var | Const of int | Pair of term * term
type env = term Env.t
let rec walk (s : env) = function (* Compute the value of the given term in s *)
  | Vart (Var x) when Env.mem x s -> walk s (Env.find x s)
  | u -> u
and ext_s (Var x) (v : term) (s : env) = Env.add x v s
and unify (u : term) (v : term) (s : env) = match walk s u, walk s v with
  | u', v' when u' = v' -> Some s       (* u and v unify when they are equal, *)
  | Vart u', v' -> Some (ext_s u' v' s) (* or when either is a variable, *)
  | u', Vart v' -> Some (ext_s v' u' s) (* which we can bind to the other, or *)
  | Pair(ua, ud), Pair(va, vd) -> (match unify ua va s with
    | Some s' -> unify ud vd s'         (* when they are pairs *)
    | None -> None)                     (* whose respective members unify *)
  | _ -> None
type state = State of env * Index.t (* index of the next variable to create *)
type stream = Cons of state * stream | Thunk of (unit -> stream) | Mzero
type goal = state -> stream
let rec mplus (t : stream) = function (* Monad plus: merge two streams *) 
  | Mzero -> t (* If the second stream is empty, return the first stream *)
  | Thunk f -> Thunk (fun () -> mplus (f()) t) (* η⁻¹-delay and swap *)
  | Cons(e, s) -> Cons(e, mplus t s)  (* or simply recurse if eager *)
and bind (g : goal) = function (* Monad bind: transform stream with g *)
  | Mzero -> Mzero
  | Thunk t -> Thunk (fun () -> bind g (t()))
  | Cons(e, s) -> mplus (g e) (bind g s)
and just (s : state) = Cons(s, Mzero) (* Lift a state into the stream monad *)
and disj (g1 : goal) (g2 : goal) (t : state) = mplus (g1 t) (g2 t)
and conj (g1 : goal) (g2 : goal) (t : state) = bind g2 (g1 t)
(* Supply a fresh logic variable to a goal-returning function *)
and call_fresh (f : term -> goal) (State(e, c)) = (* <Scheme> I’m monomorphic *)
  (f (Vart (Var c))) (State (e, Index.succ c))    (* <OCaml>  hold my beer. *)
let match_goal (u : term) (v : term) (State (s, c)) = match unify u v s with
  | Some s' -> just (State (s', c))  (* Create an equality goal (≡) *)
  | None -> Mzero

(* full version at <http://canonical.org/~kragen/sw/dev3/mukanren.ml> *)

(*
# conj (call_fresh (fun a -> match_goal a (Const 7)))
       (call_fresh (fun b -> disj (match_goal b (Const 5))
                                  (match_goal b (Const 6)))) (State(Env.empty, Index.zero)) ;;
- : Exkanren.stream =
Cons (State (<abstr>, 2l), Cons (State (<abstr>, 2l), Mzero))
# conj (call_fresh (fun a -> match_goal (Const 3) (Const 7)))
       (call_fresh (fun b -> disj (match_goal b (Const 5))
                                  (match_goal b (Const 6)))) (State(Env.empty, Index.zero)) ;;
- : Exkanren.stream = Mzero
*)