1 files changed, 18 insertions, 13 deletions
diff --git a/src/elaborator/Elaborator.ml b/src/elaborator/Elaborator.ml
index df2a358..ba1da02 100644
--- a/src/elaborator/Elaborator.ml
+++ b/src/elaborator/Elaborator.ml
@@ -53,19 +53,17 @@ 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
+    Pretty.pp ~names fmt tm
 
 (** Main algorithm *)
 
+type connective = [ `Pi | `Sigma ]
+
 let rec check ~(tm : A.expr) ~(tp : D.t) : S.t =
   Error.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)
-      | _ -> Error.fatalf Error.Message.IllTyped "expected an element of type @[%a@], but found an element of a universe" (pp()) (quote tp) end
+  | A.Pi ((name, base), fam) -> check_connective `Pi ~name:name.value ~base ~fam ~tp
+  | A.Fun (base, fam) -> check_connective `Pi ~name:None ~base ~fam ~tp
   | A.Lam (name, body) -> begin match tp with
       | D.Pi (_, base, fam) -> 
         let body = bind ~name:name.value ~tp:base @@ fun arg ->
@@ -74,12 +72,8 @@ let rec check ~(tm : A.expr) ~(tp : D.t) : S.t =
         S.Lam (name.value, body)
       | _ -> Error.fatalf Error.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)
-      | _ -> Error.fatalf Error.Message.IllTyped "expected an element of type @[%a@], but found an element of a universe" (pp()) (quote tp) end
+  | A.Sg ((name, base), fam) -> check_connective `Sigma ~name:name.value ~base ~fam ~tp
+  | A.Prod (base, fam) -> check_connective `Sigma ~name:None ~base ~fam ~tp
   | A.Pair (fst, snd) -> begin match tp with
       | D.Sg (_, base, fam) -> 
         let fst = check ~tm:fst ~tp:base in
@@ -111,6 +105,17 @@ let rec check ~(tm : A.expr) ~(tp : D.t) : S.t =
     end;
     tm
 
+and check_connective connective ~(name : Ident.local) ~(base : A.expr) ~(fam : A.expr) ~(tp : D.t) =
+  match tp with
+  | D.Type ->
+    let base = check ~tm:base ~tp in
+    let fam = bind ~name:name ~tp:(eval base) @@ fun _ -> check ~tm:fam ~tp in
+    begin match connective with
+      | `Pi -> S.Pi (name, base, fam)
+      | `Sigma -> S.Sg (name, base, fam)
+    end
+  | _ -> Error.fatalf Error.Message.IllTyped "expected an element of type @[%a@], but found an element of a universe" (pp()) (quote tp)
+
 and infer (tm : A.expr) : D.t * S.t =
   Error.tracef ?loc:tm.loc "when inferring the type" @@ fun () ->
   match tm.value with