bindlib_cata.ml

[@@@warnerror "-unused-value-declaration"]

open Bindlib

let ( <.> ) f g x = f (g x)
let ( <..> ) f g x y = f (g x y)

module Mu (M : sig
  type 'fix t
end) =
struct
  (*
     Inductive type is only allowed if wrapped in a constructor.
     Haskell users can use newtype to avoid this.
  *)
  type mu = In of mu M.t

  let inj (In x) = x (* mu -> f mu *)
  let prj x = In x (* f mu -> mu *)

  (* induction principle:
     P t := t -> a
     (P mu -> P (f mu)) -> P mu
  *)
  let rec ind f (In x) = f (ind f) x

  (* ... for two arguments *)
  let rec ind_oo f (In x) (In y) = f (ind_oo f) x y

  (* ... not all arguments must be of type mu *)
  let rec ind_xo f x (In y) = f (ind_xo f) x y
end

module type BindFunctor = sig
  type 'fix t

  (* Functor fmap, but the codomain must have box lift and mkfree impls *)
  val fmap_bind : ('a -> 'b) -> ('b var -> 'b) -> ('b -> 'b box) -> 'a t -> 'b t

  (* Functor fmap with ctx support. This implementation updates ctx while traversing binders. *)
  val fmap_bind_ctx :
    (ctxt -> 'a -> 'b) ->
    ('b var -> 'b) ->
    ('b -> 'b box) ->
    ctxt ->
    'a t ->
    'b t
end

module type BindLift = sig
  type t

  val mkfree : t var -> t
  val lift : t -> t box
end

module BindCata
    (M : BindFunctor)
    (Mf : module type of Mu (M)) (F : BindLift) =
struct
  (* catamorphism *)
  let rec cata f x = f (M.fmap_bind (cata f) F.mkfree F.lift (Mf.inj x))

  (* catamorphism that updates ctx while traversing binders
     There's a performance penalty for updating ctx in each step.
  *)
  let rec cata_ctx f ctx x =
    f ctx (M.fmap_bind_ctx (cata_ctx f) F.mkfree F.lift ctx (Mf.inj x))
end

module IR = struct
  module T = struct
    type 'e t =
      | Var of 'e Bindlib.var
      | App of 'e * 'e
      | Abs of ('e, 'e) Bindlib.binder

    let fmap_bind f free lift = function
      | Var v -> Var (copy_var v free (name_of v))
      | App (e1, e2) -> App (f e1, f e2)
      | Abs b ->
          let x, e = unbind b in
          let x' = copy_var x free (name_of x) in
          let box_binder = bind_var x' (f e |> lift) in
          Abs (unbox box_binder)

    let fmap_bind_ctx f free lift ctx = function
      | Abs b ->
          let x, e, ctx = unbind_in ctx b in
          let x' = copy_var x free (name_of x) in
          let box_binder = bind_var x' (f ctx e |> lift) in
          Abs (unbox box_binder)
      | e -> fmap_bind (f ctx) free lift e (* no ctx update for other cases *)
  end

  module Mu = Mu (T)
  include T
  include Mu
  module Cata (F : BindLift) = BindCata (T) (Mu) (F)

  (* smart constructors *)
  let var = box_var
  let app = box_apply2 (fun t u -> App (t, u) |> prj)
  let abs_raw b = box_apply (fun b -> Abs b |> prj) b
  let abs = abs_raw <..> bind_var

  module Lift = struct
    type t = mu

    let lift =
      ind (fun hypo -> function
        | Var v -> var v
        | App (e1, e2) -> app (hypo e1) (hypo e2)
        | Abs b -> abs_raw (box_binder hypo b))

    let mkfree x = Var x |> prj
  end

  let lift = Lift.lift
  let mkfree = Lift.mkfree

  module FCata = Cata (Lift)

  let map f = FCata.cata (f <.> prj)

  (* map that update ctx while traversing binders
     Prefer `map` over this one if ctx is not needed for better performance.
  *)
  let map_ctx f = FCata.cata_ctx (fun ctx x -> f ctx (prj x))
end

(* convenience module for output type of cata *)
module BindData (M : sig
  type t
end) =
struct
  type t = Data of M.t | Var of t var

  let map f = function Data i -> Data (f i) | Var x -> Var x
  let unwrap_or default = function Data i -> i | Var _ -> default
  let mkfree x = Var x
  let lift = function Data i -> box (Data i) | Var x -> box_var x
end

module BindInt = BindData (Int)
module IntCata = IR.Cata (BindInt)

(* calculate rank of a term *)
let calc_rank =
  BindInt.unwrap_or 0
  <.> IntCata.cata (function
        | Var _ -> Data 0
        | Abs b ->
            let _, body = unbind b in
            BindInt.map (fun x -> x + 1) body
        | App (Data e1, Data e2) -> Data (max e1 e2)
        | _ -> failwith "impossible")

module BindString = BindData (String)
module StringCata = IR.Cata (BindString)

(* show a term as a string
   use ctx version because we want up-to-date names
   NOTE: I believe there's no need to use `unbind_in` in `show` because
   1) no subst is done in show
   2) `cata` calls fmap first which uses `unbind_in` to update ctx, so the ctx
      is already up-to-date when `show` is called
*)
let show =
  BindString.unwrap_or ""
  <.> StringCata.cata_ctx
        (fun _ -> function
          | Var x -> Data (name_of x)
          | Abs b ->
              let x, t = unbind b in
              BindString.map (fun t -> "λ" ^ name_of x ^ ". " ^ t) t
          | App (Data e1, Data e2) -> Data ("(" ^ e1 ^ " " ^ e2 ^ ")")
          | _ -> failwith "impossible")
        empty_ctxt

open IR

let () =
  let x = new_var mkfree "x" in
  let y = new_var mkfree "y" in
  let z = new_var mkfree "z" in
  let e = abs x (abs y (abs z (app (app (var x) (var y)) (var z)))) in
  Printf.printf "show e: %s\n" (show (unbox e));
  Printf.printf "rank of e: %d\n" (calc_rank (unbox e))