split code into smaller libraries and make a better repl

1 files changed, 161 insertions, 0 deletions

diff --git a/src/elaborator/Elaborator.ml b/src/elaborator/Elaborator.ml
new file mode 100644
index 0000000..0f72f24
--- /dev/null
+++ b/src/elaborator/Elaborator.ml
@@ -0,0 +1,161 @@
+open Bwd
+open Bwd.Infix
+
+module A = Ast
+module S = NbE.Syntax
+module D = NbE.Domain
+
+(* TODO: REALLY improve error messages *)
+
+(* invariant: `tps`, `tms` and `names` all have length `size` *)
+type env = {
+  tps : D.env;
+  tms : D.env;
+  names : Ident.local bwd;
+  size : int;
+}
+
+module Eff = Algaeff.Reader.Make (struct type nonrec t = env end)
+
+(** General helpers *)
+
+let lookup (id : Ident.t) : D.t * S.t = match Ident.to_local id with
+  | Some name -> begin
+      let env = Eff.read() in
+      match Bwd.find_index ((=) name) env.names with
+      | Some ix ->
+        let tp = Bwd.nth env.tps ix in
+        (tp, S.Var ix)
+      | None -> Reporter.unbound_variable id
+    end
+  | None -> Reporter.unbound_variable id
+
+let bind ~(name : Ident.local) ~(tp : D.t) f =
+  let arg = D.var (Eff.read()).size in
+  let update env = {
+    tps = env.tps <: tp;
+    tms = env.tms <: arg;
+    names = env.names <: name;
+    size = env.size + 1;
+  } in
+  Eff.scope update (fun () -> f arg)
+
+(** NbE helpers *)
+
+let eval tm = NbE.eval ~env:(Eff.read()).tms tm
+
+(** Evaluate under the current environment augmented by `arg` *)
+(* TODO: this is kind of inelegant, can we do better? *)
+let eval_at arg = NbE.eval ~env:((Eff.read()).tms <: arg)
+
+let quote v = NbE.quote ~size:(Eff.read()).size v
+
+(** Pretty-printing helpers *)
+
+let pp () = let names = (Eff.read()).names in fun fmt tm -> 
+  Pretty.pp ~names fmt tm
+
+(** Main algorithm *)
+
+let rec check ~(tm : A.expr) ~(tp : D.t) : S.t =
+  Reporter.tracef ?loc:tm.loc "when checking against the type @[%a@]" (pp()) (quote tp) @@ fun () ->
+  match tm.value with
+  | A.Pi ((name, base), fam) -> begin match tp with
+      | D.Type ->
+        let base = check ~tm:base ~tp in
+        let fam = bind ~name:name.value ~tp:(eval base) @@ fun _ -> check ~tm:fam ~tp in
+        S.Pi (name.value, base, fam)
+      | _ -> Reporter.fatalf Reporter.Message.IllTyped "expected an element of type @[%a@], but found an element of a universe" (pp()) (quote tp) end
+  | A.Lam (name, body) -> begin match tp with
+      | D.Pi (_, base, fam) -> 
+        let body = bind ~name:name.value ~tp:base @@ fun arg ->
+          let fib = NbE.inst_clo fam arg in
+          check ~tm:body ~tp:fib in
+        S.Lam (name.value, body)
+      | _ -> Reporter.fatalf Reporter.Message.IllTyped "expected an element of type @[%a@], but found an element of a Pi type" (pp()) (quote tp)
+    end
+  | A.Sg ((name, base), fam) -> begin match tp with
+      | D.Type ->
+        let base = check ~tm:base ~tp in
+        let fam = bind ~name:name.value ~tp:(eval base) @@ fun _ -> check ~tm:fam ~tp in
+        S.Sg (name.value, base, fam)
+      | _ -> Reporter.fatalf Reporter.Message.IllTyped "expected an element of type @[%a@], but found an element of a universe" (pp()) (quote tp) end
+  | A.Pair (fst, snd) -> begin match tp with
+      | D.Sg (_, base, fam) -> 
+        let fst = check ~tm:fst ~tp:base in
+        let fib = NbE.inst_clo fam (eval fst) in
+        let snd = check ~tm:snd ~tp:fib in
+        S.Pair (fst, snd)
+      | _ -> Reporter.fatalf Reporter.Message.IllTyped "expected an element of type @[%a@], but found an element of a Sigma type" (pp()) (quote tp)
+    end
+  | A.Type -> begin match tp with (* TODO type-in-type *)
+      | D.Type -> S.Type
+      | _ -> Reporter.fatalf Reporter.Message.IllTyped "expected an element of type @[%a@], but found an element of a universe" (pp()) (quote tp)
+    end
+  | A.Bool -> begin match tp with
+      | D.Type -> S.Bool
+      | _ -> Reporter.fatalf Reporter.Message.IllTyped "expected an element of type @[%a@], but found an element of a universe" (pp()) (quote tp)
+    end
+  | A.True -> begin match tp with
+      | D.Bool -> S.True
+      | _ -> Reporter.fatalf Reporter.Message.IllTyped "expected an element of type @[%a@], but found an element of type Bool" (pp()) (quote tp)
+    end
+  | A.False -> begin match tp with
+      | D.Bool -> S.False
+      | _ -> Reporter.fatalf Reporter.Message.IllTyped "expected an element of type @[%a@], but found an element of type Bool" (pp()) (quote tp)
+    end
+  | _ -> let (inferred_tp, tm) = infer tm in begin
+      try NbE.equate ~size:((Eff.read()).size) inferred_tp tp with
+      | NbE.Unequal ->
+        Reporter.fatalf Reporter.Message.IllTyped "expected an element of type @[%a@], but found an element of type @[%a@]" (pp()) (quote tp) (pp()) (quote inferred_tp)
+    end;
+    tm
+
+and infer (tm : A.expr) : D.t * S.t =
+  Reporter.tracef ?loc:tm.loc "when inferring the type" @@ fun () ->
+  match tm.value with
+  | A.Var name -> lookup name
+  | A.Check (tm, tp) ->
+    let tp = eval @@ check ~tp:D.Type ~tm:tp in
+    let tm = check ~tp ~tm in
+    (tp, tm)
+  | A.App (fn, arg) -> begin match infer fn with
+      | (D.Pi (_, base, fam), fn) ->
+        let arg = check ~tm:arg ~tp:base in
+        let tp = NbE.inst_clo fam (eval arg) in
+        let tm = S.App (fn, arg) in
+        (tp, tm)
+      | _ -> Reporter.fatalf Reporter.Message.IllTyped "cannot apply this term because it is not a function"
+    end
+  | A.Fst p -> begin match infer p with
+      | (D.Sg (_, base, _), p) -> (base, S.Fst p)
+      | _ -> Reporter.fatalf Reporter.Message.IllTyped "cannot apply first projection to this term because it is not of sigma type"
+    end
+  | A.Snd p -> begin match infer p with
+      | (D.Sg (_, _, fam), p) ->
+        let tp = NbE.inst_clo fam (eval (S.Fst p)) in
+        let tm = S.Snd p in
+        (tp, tm)
+      | _ -> Reporter.fatalf Reporter.Message.IllTyped "cannot apply second projection to this term because it is of sigma type"
+    end
+  | A.BoolElim { motive_var; motive_body; true_case; false_case; scrut } ->
+    let scrut = check ~tm:scrut ~tp:D.Bool in
+    let motive = bind ~name:motive_var.value ~tp:D.Bool @@ fun _ ->
+      check ~tm:motive_body ~tp:D.Type in
+    let motive_true = eval_at D.True motive in
+    let motive_false = eval_at D.False motive in
+    let motive_scrut = eval_at (eval scrut) motive in
+    let true_case = check ~tm:true_case ~tp:motive_true in
+    let false_case = check ~tm:false_case ~tp:motive_false in
+    let tm = S.BoolElim { motive_var = motive_var.value; motive; true_case; false_case; scrut } in
+    (motive_scrut, tm)
+  | _ -> Reporter.fatalf Reporter.Message.CannotInferType "cannot infer type"
+
+let empty_env : env = {
+  tps = Emp;
+  tms = Emp;
+  names = Emp;
+  size = 0;
+}
+
+let infer_toplevel (tm : A.expr) = Eff.run ~env:empty_env @@ fun () -> infer tm
\ No newline at end of file