1 files changed, 43 insertions, 21 deletions

diff --git a/src/elaborator/Elaborator.ml b/src/elaborator/Elaborator.ml
index f40076d..c366ff3 100644
--- a/src/elaborator/Elaborator.ml
+++ b/src/elaborator/Elaborator.ml
@@ -5,32 +5,42 @@ module A = Ast
 module S = NbE.Syntax
 module D = NbE.Domain
 
-(* invariant: `tps`, `tms` and `names` all have length `size` *)
 type env = {
+  (* local context *)
+  (* invariant: `tps`, `tms` and `names` all have length `size` *)
   tps : D.env;
   tms : D.env;
   names : Ident.local bwd;
   size : int;
+
+  (* top-level context *)
+  toplvl : TopLevel.t;
 }
 
 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 -> Error.unbound_variable id
-    end
-  | None -> Error.unbound_variable id
+(* general helpers *)
+
+let lookup (id : Ident.t) : D.t * S.t =
+  let env = Eff.read() in
+  (* search through local context *)
+  match Option.bind 
+          (Ident.to_local id) 
+          (fun name -> Bwd.find_index ((=) name) env.names) with
+  | Some ix ->
+    let tp = Bwd.nth env.tps ix in
+    (tp, S.Var ix)
+  | None ->
+    (* look up in top-level context *)
+    match Yuujinchou.Trie.find_singleton id env.toplvl with
+    | Some (Def { tp; tm }, _) ->
+      (tp, S.Def (id, tm))
+    | None -> Error.unbound_variable id
 
 let bind ~(name : Ident.local) ~(tp : D.t) f =
   let arg = D.var (Eff.read()).size in
-  let update env = {
+  let update env = { 
+    env with
     tps = env.tps <: tp;
     tms = env.tms <: arg;
     names = env.names <: name;
@@ -38,15 +48,15 @@ let bind ~(name : Ident.local) ~(tp : D.t) f =
   } in
   Eff.scope update (fun () -> f arg)
 
-(** NbE helpers *)
+(* NbE helpers *)
 
 let eval tm = NbE.eval ~env:(Eff.read()).tms tm
 
-(** Evaluate under the current environment augmented by `arg` *)
+(* 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)
 
-(** Pretty-printing helpers *)
+(* pretty-printing helpers *)
 
 let pp_tm () =
   let names = (Eff.read()).names in
@@ -57,7 +67,7 @@ let pp_val () =
   let pp_tm = pp_tm () in
   fun fmt v -> pp_tm fmt (NbE.quote ~size v)
 
-(** Main algorithm *)
+(* main algorithm *)
 
 type connective = [ `Pi | `Sigma ]
 
@@ -118,12 +128,14 @@ and check_connective connective ~(name : Ident.local) ~(base : A.expr) ~(fam : A
     end
   | _ -> Error.type_mismatch (pp_val()) tp Fmt.string "type"
 
+and check_tp (tp : A.expr) = eval @@ check ~tp:D.Type ~tm:tp
+
 and infer (tm : A.expr) : D.t * S.t =
   Error.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 tp = check_tp tp in
     let tm = check ~tp ~tm in
     (tp, tm)
   | A.App (fn, arg) -> begin match infer fn with
@@ -158,11 +170,21 @@ and infer (tm : A.expr) : D.t * S.t =
     (motive_scrut, tm)
   | _ -> Error.not_inferable ()
 
-let empty_env : env = {
+(* interface *)
+
+let initial_env toplvl : env = {
   tps = Emp;
   tms = Emp;
   names = Emp;
   size = 0;
+  toplvl;
 }
 
-let infer_toplevel (tm : A.expr) = Eff.run ~env:empty_env @@ fun () -> infer tm
\ No newline at end of file
+let check_toplevel ~(toplvl : TopLevel.t) ~(tm : A.expr) ~(tp : D.t) =
+  Eff.run ~env:(initial_env toplvl) @@ fun () -> check ~tm ~tp
+
+let check_tp_toplevel ~(toplvl : TopLevel.t) (tp : A.expr) =
+  Eff.run ~env:(initial_env toplvl) @@ fun () -> check_tp tp
+
+let infer_toplevel ~(toplvl : TopLevel.t) (tm : A.expr) =
+  Eff.run ~env:(initial_env toplvl) @@ fun () -> infer tm
\ No newline at end of file