split code into smaller libraries and make a better repl

1 files changed, 93 insertions, 0 deletions

diff --git a/src/pretty/Pretty.ml b/src/pretty/Pretty.ml
new file mode 100644
index 0000000..66d6159
--- /dev/null
+++ b/src/pretty/Pretty.ml
@@ -0,0 +1,93 @@
+open Bwd
+open Bwd.Infix
+
+module S = NbE.Syntax
+
+module Internal =
+struct
+  type env = {
+    names : Ident.local bwd;
+    prec : int;
+  }
+
+  module Eff = Algaeff.Reader.Make (struct type nonrec t = env end)
+
+  let bind (name : Ident.local) f fmt x =
+    let update env = { env with names = env.names <: name } in
+    Eff.scope update (fun () -> f fmt x)
+
+  let with_prec (prec : int) (inner : 'a Fmt.t) (fmt : Format.formatter) (v : 'a) =
+    Eff.scope (fun env -> { env with prec; }) (fun () -> inner fmt v)
+
+  let delimited inner = with_prec 0 inner
+
+  let parenthesize (prec : int) : 'a Fmt.t -> 'a Fmt.t =
+    if (Eff.read()).prec > prec then
+      Fmt.parens
+    else
+      Fun.id
+
+  let lassoc (prec : int)
+      (pp_left : 'l Fmt.t)
+      ?(sep = " ")
+      (pp_right : 'r Fmt.t) : ('l * 'r) Fmt.t = 
+    parenthesize prec @@ fun fmt v -> 
+    Fmt.pf fmt "@[%a@]%s@[%a@]" 
+      (with_prec prec pp_left) (fst v)
+      sep
+      (with_prec (prec + 1) pp_right) (snd v)
+
+  let rassoc (prec : int)
+      (pp_left : 'l Fmt.t)
+      ?(sep = " ")
+      (pp_right : 'r Fmt.t) : ('l * 'r) Fmt.t = 
+    parenthesize prec @@ fun fmt v -> 
+    Fmt.pf fmt "@[%a@]%s@[%a@]" 
+      (with_prec (prec + 1) pp_left) (fst v)
+      sep
+      (with_prec prec pp_right) (snd v)
+
+  let pp_local_name fmt name =
+    Fmt.string fmt @@ Option.value ~default:"_" name
+
+  let pp_var fmt ix =
+    let names = (Eff.read()).names in
+    match Bwd.nth_opt names ix with
+    (* the variable has a bound name which has not been shadowed *)
+    | Some ((Some name) as id) when Bwd.find_index ((=) id) names = Some ix -> Fmt.string fmt name
+    (* the variable does not have a bound name or its bound name has been shadowed *)
+    | Some _ -> Fmt.pf fmt "%@%d" ix
+    | None -> Reporter.bug "index out of range in pp_var"
+
+  let rec pp_pair fmt = (delimited @@ fun fmt (a, b) -> Fmt.pf fmt "@[(@[%a@], @[%a@])@]" pp a pp b) fmt
+
+  and pp_bool_elim fmt =
+    (delimited @@ fun fmt (motive_var, motive, true_case, false_case, scrut) -> Fmt.pf fmt
+       "@[bool-elim @[%a@] at @[%a@] -> @[%a@] [ true => @[%a@], false => @[%a@] ]@]"
+       pp scrut pp_local_name motive_var (bind motive_var pp) motive pp true_case pp false_case
+    ) fmt
+
+  and pp_arg fmt (name, tp) =
+    Fmt.pf fmt "(@[%a@] : @[%a@])" pp_local_name name pp tp
+
+  (* TODO: improve indentation *)
+  and pp fmt = function
+    | S.Var ix -> pp_var fmt ix
+    | S.Pi (name, base, fam) -> rassoc 1 pp_arg ~sep:" -> " (bind name pp) fmt ((name, base), fam)
+    | S.Lam (name, body) -> Fmt.pf fmt "@[\\@[%a@] -> @[%a@]@]" pp_local_name name (bind name pp) body
+    | S.Sg (name, base, fam) -> rassoc 2 pp_arg ~sep:" * " (bind name pp) fmt ((name, base), fam)
+    | S.App (a, b) -> lassoc 3 pp pp fmt (a, b)
+    | S.Pair (a, b) -> pp_pair fmt (a, b)
+    | S.Fst a -> lassoc 3 Fmt.string pp fmt ("fst", a)
+    | S.Snd a -> lassoc 3 Fmt.string pp fmt ("snd", a)
+    | S.Type -> Fmt.string fmt "type"
+    | S.Bool -> Fmt.string fmt "bool"
+    | S.True -> Fmt.string fmt "true"
+    | S.False -> Fmt.string fmt "false"
+    | S.BoolElim { motive_var; motive; true_case; false_case; scrut }
+      -> pp_bool_elim fmt (motive_var, motive, true_case, false_case, scrut)
+end
+
+let pp ~names fmt tm =
+  let env : Internal.env = { names; prec = 0; } in
+  Internal.Eff.run ~env (fun () -> Internal.pp fmt tm)