Typecheck more cases

5 files changed, 165 insertions, 62 deletions

diff --git a/dhall/src/error/mod.rs b/dhall/src/error/mod.rs
index 5595c53..850a792 100644
--- a/dhall/src/error/mod.rs
+++ b/dhall/src/error/mod.rs
@@ -83,6 +83,7 @@ pub(crate) enum TypeMessage {
     EquivalenceTypeMismatch(Value, Value),
     AssertMismatch(Value, Value),
     AssertMustTakeEquivalence,
+    Custom(String),
 }
 
 impl TypeError {
diff --git a/dhall/src/semantics/phase/normalize.rs b/dhall/src/semantics/phase/normalize.rs
index e848601..066a004 100644
--- a/dhall/src/semantics/phase/normalize.rs
+++ b/dhall/src/semantics/phase/normalize.rs
@@ -625,12 +625,14 @@ pub(crate) fn normalize_one_layer(
         // `RetWhole`), so they won't appear here.
         ExprKind::Lam(_, _, _)
         | ExprKind::Pi(_, _, _)
-        | ExprKind::Let(_, _, _, _)
         | ExprKind::Embed(_)
-        | ExprKind::Var(_)
-        | ExprKind::Annot(_, _) => {
+        | ExprKind::Var(_) => {
             unreachable!("This case should have been handled in typecheck")
         }
+        ExprKind::Let(_, _, val, body) => {
+            Ret::Value(body.subst_shift(&AlphaVar::default(), &val))
+        }
+        ExprKind::Annot(x, _) => Ret::Value(x),
         ExprKind::Const(c) => Ret::Value(const_to_value(c)),
         ExprKind::Builtin(b) => Ret::Value(builtin_to_value(b)),
         ExprKind::Assert(_) => Ret::Expr(expr),
@@ -895,11 +897,16 @@ pub(crate) fn normalize_tyexpr_whnf(tye: &TyExpr, env: &NzEnv) -> Value {
                 closure: Closure::new(env, body.clone()),
             }
         }
+        TyExprKind::Expr(ExprKind::Let(_, None, val, body)) => {
+            let val = val.normalize_whnf(env);
+            return body.normalize_whnf(&env.insert_value(val));
+        }
         TyExprKind::Expr(e) => {
             let e = e.map_ref(|tye| tye.normalize_whnf(env));
             normalize_one_layer(e, &ty)
         }
     };
 
+    // dbg!(tye.kind(), env, &kind);
     Value::from_kind_and_type_whnf(kind, ty)
 }
diff --git a/dhall/src/semantics/phase/typecheck.rs b/dhall/src/semantics/phase/typecheck.rs
index f2d1bf2..0f3754e 100644
--- a/dhall/src/semantics/phase/typecheck.rs
+++ b/dhall/src/semantics/phase/typecheck.rs
@@ -189,7 +189,7 @@ macro_rules! make_type {
     };
 }
 
-fn type_of_builtin<E>(b: Builtin) -> Expr<E> {
+pub(crate) fn type_of_builtin<E>(b: Builtin) -> Expr<E> {
     use syntax::Builtin::*;
     rc(match b {
         Bool | Natural | Integer | Double | Text => make_type!(Type),
@@ -321,7 +321,14 @@ fn type_with(ctx: &TyCtx, e: Expr<Normalized>) -> Result<Value, TypeError> {
 
             let v = type_with(ctx, v)?;
             let binder = ctx.new_binder(x);
-            type_with(&ctx.insert_value(&binder, v.clone())?, e.clone())
+            let e =
+                type_with(&ctx.insert_value(&binder, v.clone())?, e.clone())?;
+            // let e_ty = e.get_type()?;
+            // Ok(Value::from_kind_and_type(
+            //     ValueKind::PartialExpr(ExprKind::Let(x.clone(), None, v, e)),
+            //     e_ty,
+            // ))
+            Ok(e)
         }
         Embed(p) => Ok(p.clone().into_typed().into_value()),
         Var(var) => match ctx.lookup(&var) {
diff --git a/dhall/src/semantics/tck/typecheck.rs b/dhall/src/semantics/tck/typecheck.rs
index aab5f83..046c999 100644
--- a/dhall/src/semantics/tck/typecheck.rs
+++ b/dhall/src/semantics/tck/typecheck.rs
@@ -10,6 +10,9 @@ use crate::error::{TypeError, TypeMessage};
 use crate::semantics::core::context::TyCtx;
 use crate::semantics::nze::{NameEnv, QuoteEnv};
 use crate::semantics::phase::normalize::{merge_maps, NzEnv};
+use crate::semantics::phase::typecheck::{
+    builtin_to_value, const_to_value, type_of_builtin,
+};
 use crate::semantics::phase::Normalized;
 use crate::semantics::{
     AlphaVar, Binder, Closure, TyExpr, TyExprKind, Type, Value, ValueKind,
@@ -55,93 +58,171 @@ impl TyEnv {
             items: self.items.insert_value(e),
         }
     }
-    pub fn lookup_var(&self, var: &V<Label>) -> Option<(TyExprKind, Value)> {
+    pub fn lookup(&self, var: &V<Label>) -> Option<(TyExprKind, Value)> {
         let var = self.names.unlabel_var(var)?;
-        let ty = self.lookup_val(&var).get_type().unwrap();
+        let ty = self.items.lookup_val(&var).get_type().unwrap();
         Some((TyExprKind::Var(var), ty))
     }
-    pub fn lookup_val(&self, var: &AlphaVar) -> Value {
-        self.items.lookup_val(var)
-    }
     pub fn size(&self) -> usize {
         self.names.size()
     }
 }
 
+fn mkerr<T>(x: impl ToString) -> Result<T, TypeError> {
+    Err(TypeError::new(TypeMessage::Custom(x.to_string())))
+}
+
+fn type_last_layer(
+    env: &TyEnv,
+    kind: &ExprKind<TyExpr, Normalized>,
+) -> Result<Value, TypeError> {
+    Ok(match &kind {
+        ExprKind::Import(..) => unreachable!(
+            "There should remain no imports in a resolved expression"
+        ),
+        ExprKind::Var(..)
+        | ExprKind::Lam(..)
+        | ExprKind::Pi(..)
+        | ExprKind::Let(..)
+        | ExprKind::Const(Const::Sort)
+        | ExprKind::Embed(..) => unreachable!(),
+
+        ExprKind::Const(Const::Type) => const_to_value(Const::Kind),
+        ExprKind::Const(Const::Kind) => const_to_value(Const::Sort),
+        ExprKind::Builtin(b) => {
+            type_with(env, &type_of_builtin(*b))?.normalize_whnf(env.as_nzenv())
+        }
+        ExprKind::BoolLit(_) => builtin_to_value(Builtin::Bool),
+        ExprKind::NaturalLit(_) => builtin_to_value(Builtin::Natural),
+        ExprKind::IntegerLit(_) => builtin_to_value(Builtin::Integer),
+        ExprKind::DoubleLit(_) => builtin_to_value(Builtin::Double),
+
+        ExprKind::Annot(x, t) => {
+            if x.get_type()? != t.normalize_whnf(env.as_nzenv()) {
+                return mkerr("annot mismatch");
+            }
+            x.normalize_whnf(env.as_nzenv())
+        }
+        ExprKind::App(f, arg) => {
+            let tf = f.get_type()?;
+            let tf_borrow = tf.as_whnf();
+            let (expected_arg_ty, ty_closure) = match &*tf_borrow {
+                ValueKind::PiClosure { annot, closure, .. } => (annot, closure),
+                _ => return mkerr(format!("apply to not Pi: {:?}", tf_borrow)),
+            };
+            // if arg.get_type()? != *expected_arg_ty {
+            //     return mkerr(format!(
+            //         "function annot mismatch: {:?}, {:?}",
+            //         arg.get_type()?,
+            //         expected_arg_ty
+            //     ));
+            // }
+
+            let arg_nf = arg.normalize_whnf(env.as_nzenv());
+            ty_closure.apply(arg_nf)
+        }
+        ExprKind::BoolIf(x, y, z) => {
+            if *x.get_type()?.as_whnf()
+                != ValueKind::from_builtin(Builtin::Bool)
+            {
+                return mkerr("InvalidPredicate");
+            }
+            if y.get_type()?.get_type()?.as_const() != Some(Const::Type) {
+                return mkerr("IfBranchMustBeTerm");
+            }
+            if z.get_type()?.get_type()?.as_const() != Some(Const::Type) {
+                return mkerr("IfBranchMustBeTerm");
+            }
+            if y.get_type()? != z.get_type()? {
+                return mkerr("IfBranchMismatch");
+            }
+
+            y.get_type()?
+        }
+
+        _ => Value::from_const(Const::Type), // TODO
+    })
+}
+
 /// Type-check an expression and return the expression alongside its type if type-checking
 /// succeeded, or an error if type-checking failed.
-fn type_with(env: &TyEnv, expr: Expr<Normalized>) -> Result<TyExpr, TypeError> {
-    let mkerr = || TypeError::new(TypeMessage::Sort);
-
-    match expr.as_ref() {
-        ExprKind::Var(var) => match env.lookup_var(&var) {
-            Some((e, ty)) => Ok(TyExpr::new(e, Some(ty), expr.span())),
-            None => Err(mkerr()),
+fn type_with(
+    env: &TyEnv,
+    expr: &Expr<Normalized>,
+) -> Result<TyExpr, TypeError> {
+    let (tyekind, ty) = match expr.as_ref() {
+        ExprKind::Var(var) => match env.lookup(&var) {
+            Some((k, ty)) => (k, Some(ty)),
+            None => return mkerr("unbound variable"),
         },
         ExprKind::Lam(binder, annot, body) => {
-            let annot = type_with(env, annot.clone())?;
+            let annot = type_with(env, annot)?;
             let annot_nf = annot.normalize_whnf(env.as_nzenv());
-            let body = type_with(
-                &env.insert_type(&binder, annot_nf.clone()),
-                body.clone(),
-            )?;
+            let body =
+                type_with(&env.insert_type(&binder, annot_nf.clone()), body)?;
             let ty = Value::from_kind_and_type(
                 ValueKind::PiClosure {
                     binder: Binder::new(binder.clone()),
                     annot: annot_nf,
                     closure: Closure::new(
                         env.as_nzenv(),
-                        body.get_type()?.to_tyexpr(env.as_quoteenv()),
+                        body.get_type()?.to_tyexpr(env.as_quoteenv().insert()),
                     ),
                 },
                 Value::from_const(Const::Type), // TODO: function_check
             );
-            Ok(TyExpr::new(
+            (
                 TyExprKind::Expr(ExprKind::Lam(binder.clone(), annot, body)),
                 Some(ty),
-                expr.span(),
-            ))
+            )
         }
         ExprKind::Pi(binder, annot, body) => {
-            let annot = type_with(env, annot.clone())?;
+            let annot = type_with(env, annot)?;
             let annot_nf = annot.normalize_whnf(env.as_nzenv());
-            let body = type_with(
-                &env.insert_type(binder, annot_nf.clone()),
-                body.clone(),
-            )?;
-            Ok(TyExpr::new(
+            let body =
+                type_with(&env.insert_type(binder, annot_nf.clone()), body)?;
+            (
                 TyExprKind::Expr(ExprKind::Pi(binder.clone(), annot, body)),
-                None, // TODO: function_check
-                expr.span(),
-            ))
+                Some(Value::from_const(Const::Type)), // TODO: function_check
+            )
         }
-        kind => {
-            let kind = kind.traverse_ref(|e| type_with(env, e.clone()))?;
-            let ty = match &kind {
-                ExprKind::App(f, arg) => {
-                    let tf = f.get_type()?;
-                    let tf_borrow = tf.as_whnf();
-                    let (_expected_arg_ty, ty_closure) = match &*tf_borrow {
-                        ValueKind::PiClosure { annot, closure, .. } => {
-                            (annot, closure)
-                        }
-                        _ => return Err(mkerr()),
-                    };
-                    // if arg.get_type()? != *expected_arg_ty {
-                    //     return Err(mkerr());
-                    // }
-
-                    let arg_nf = arg.normalize_whnf(env.as_nzenv());
-                    ty_closure.apply(arg_nf)
-                }
-                _ => Value::from_const(Const::Type), // TODO
+        ExprKind::Let(binder, annot, val, body) => {
+            let v = if let Some(t) = annot {
+                t.rewrap(ExprKind::Annot(val.clone(), t.clone()))
+            } else {
+                val.clone()
             };
 
-            Ok(TyExpr::new(TyExprKind::Expr(kind), Some(ty), expr.span()))
+            let val = type_with(env, val)?;
+            let val_nf = val.normalize_whnf(&env.as_nzenv());
+            let body = type_with(&env.insert_value(&binder, val_nf), body)?;
+            let body_ty = body.get_type().ok();
+            (
+                TyExprKind::Expr(ExprKind::Let(
+                    binder.clone(),
+                    None,
+                    val,
+                    body,
+                )),
+                body_ty,
+            )
         }
-    }
+        ExprKind::Const(Const::Sort) => {
+            (TyExprKind::Expr(ExprKind::Const(Const::Sort)), None)
+        }
+        ExprKind::Embed(p) => {
+            return Ok(p.clone().into_typed().into_value().to_tyexpr_noenv())
+        }
+        ekind => {
+            let ekind = ekind.traverse_ref(|e| type_with(env, e))?;
+            let ty = type_last_layer(env, &ekind)?;
+            (TyExprKind::Expr(ekind), Some(ty))
+        }
+    };
+
+    Ok(TyExpr::new(tyekind, ty, expr.span()))
 }
 
-pub(crate) fn typecheck(e: Expr<Normalized>) -> Result<TyExpr, TypeError> {
+pub(crate) fn typecheck(e: &Expr<Normalized>) -> Result<TyExpr, TypeError> {
     type_with(&TyEnv::new(), e)
 }
diff --git a/dhall/src/tests.rs b/dhall/src/tests.rs
index 994a134..f51f6a8 100644
--- a/dhall/src/tests.rs
+++ b/dhall/src/tests.rs
@@ -158,12 +158,19 @@ pub fn run_test(test: Test<'_>) -> Result<()> {
             parse_file_str(&file_path)?.resolve().unwrap_err();
         }
         TypeInferenceSuccess(expr_file_path, expected_file_path) => {
-            let expr =
-                parse_file_str(&expr_file_path)?.resolve()?.typecheck()?;
-            let ty = expr.get_type()?.to_expr();
+            // let expr =
+            //     parse_file_str(&expr_file_path)?.resolve()?.typecheck()?;
+            // let ty = expr.get_type()?.to_expr();
             // let expr = parse_file_str(&expr_file_path)?.resolve()?.to_expr();
             // let tyexpr = crate::semantics::nze::nzexpr::typecheck(expr)?;
             // let ty = tyexpr.get_type()?.to_expr();
+            let expr = parse_file_str(&expr_file_path)?.resolve()?.to_expr();
+            let ty = crate::semantics::tck::typecheck::typecheck(&expr)?
+                .get_type()?
+                .to_expr(crate::semantics::phase::ToExprOptions {
+                    alpha: false,
+                    normalize: true,
+                });
             let expected = parse_file_str(&expected_file_path)?.to_expr();
             assert_eq_display!(ty, expected);
         }
@@ -226,7 +233,7 @@ pub fn run_test(test: Test<'_>) -> Result<()> {
             //         normalize: true,
             //     });
             let expr = parse_file_str(&expr_file_path)?.resolve()?.to_expr();
-            let expr = crate::semantics::tck::typecheck::typecheck(expr)?
+            let expr = crate::semantics::tck::typecheck::typecheck(&expr)?
                 .normalize_whnf_noenv()
                 .to_expr(crate::semantics::phase::ToExprOptions {
                     alpha: false,