Introduce Thunks and normalize lazily

7 files changed, 95 insertions, 45 deletions

diff --git a/dhall/src/semantics/builtins.rs b/dhall/src/semantics/builtins.rs
index a536261..9d10d89 100644
--- a/dhall/src/semantics/builtins.rs
+++ b/dhall/src/semantics/builtins.rs
@@ -286,7 +286,7 @@ fn apply_builtin(
         Value(Value),
         DoneAsIs,
     }
-    let make_closure = |e| typecheck(&e).unwrap().normalize_whnf(&env);
+    let make_closure = |e| typecheck(&e).unwrap().eval(&env);
 
     let ret = match (b, args.as_slice()) {
         (OptionalNone, [t]) => Ret::ValueKind(EmptyOptionalLit(t.clone())),
diff --git a/dhall/src/semantics/core/value.rs b/dhall/src/semantics/core/value.rs
index 4c844ee..c881a9a 100644
--- a/dhall/src/semantics/core/value.rs
+++ b/dhall/src/semantics/core/value.rs
@@ -4,9 +4,11 @@ use std::rc::Rc;
 
 use crate::error::{TypeError, TypeMessage};
 use crate::semantics::core::var::Binder;
-use crate::semantics::phase::normalize::{apply_any, normalize_whnf};
+use crate::semantics::phase::normalize::{
+    apply_any, normalize_tyexpr_whnf, normalize_whnf,
+};
 use crate::semantics::phase::{Normalized, NormalizedExpr, ToExprOptions};
-use crate::semantics::{type_of_builtin, TyExpr, TyExprKind};
+use crate::semantics::{type_of_builtin, typecheck, TyExpr, TyExprKind};
 use crate::semantics::{BuiltinClosure, NzEnv, NzVar, VarEnv};
 use crate::syntax::{
     BinOp, Builtin, Const, ExprKind, Integer, InterpolatedTextContents, Label,
@@ -47,6 +49,14 @@ pub(crate) enum Form {
     WHNF,
 }
 
+/// An unevaluated subexpression
+#[derive(Debug, Clone)]
+pub(crate) struct Thunk {
+    env: NzEnv,
+    body: TyExpr,
+}
+
+/// An unevaluated subexpression that takes an argument.
 #[derive(Debug, Clone)]
 pub(crate) enum Closure {
     /// Normal closure
@@ -99,6 +109,8 @@ pub(crate) enum ValueKind<Value> {
     Equivalence(Value, Value),
     // Invariant: in whnf, this must not contain a value captured by one of the variants above.
     PartialExpr(ExprKind<Value, Normalized>),
+    // Invariant: absent in whnf
+    Thunk(Thunk),
 }
 
 impl Value {
@@ -120,6 +132,15 @@ impl Value {
         }
         .into_value()
     }
+    pub(crate) fn new_thunk(env: &NzEnv, tye: TyExpr) -> Value {
+        ValueInternal {
+            form: Unevaled,
+            kind: ValueKind::Thunk(Thunk::new(env, tye.clone())),
+            ty: tye.get_type().ok(),
+            span: tye.span().clone(),
+        }
+        .into_value()
+    }
     pub(crate) fn from_kind_and_type(v: ValueKind<Value>, t: Value) -> Value {
         Value::new(v, Unevaled, t, Span::Artificial)
     }
@@ -147,9 +168,7 @@ impl Value {
     pub(crate) fn from_builtin_env(b: Builtin, env: &NzEnv) -> Self {
         Value::from_kind_and_type(
             ValueKind::from_builtin_env(b, env.clone()),
-            crate::semantics::tck::typecheck::typecheck(&type_of_builtin(b))
-                .unwrap()
-                .normalize_whnf_noenv(),
+            typecheck(&type_of_builtin(b)).unwrap().eval_closed_expr(),
         )
     }
 
@@ -325,6 +344,7 @@ impl Value {
                     v.to_tyexpr(venv),
                 ))
             }
+            ValueKind::Thunk(th) => return th.to_tyexpr(venv),
             self_kind => {
                 let self_kind = self_kind.map_ref(|v| v.to_tyexpr(venv));
                 let expr = match self_kind {
@@ -333,7 +353,8 @@ impl Value {
                     | ValueKind::PiClosure { .. }
                     | ValueKind::AppliedBuiltin(..)
                     | ValueKind::UnionConstructor(..)
-                    | ValueKind::UnionLit(..) => unreachable!(),
+                    | ValueKind::UnionLit(..)
+                    | ValueKind::Thunk(..) => unreachable!(),
                     ValueKind::Const(c) => ExprKind::Const(c),
                     ValueKind::BoolLit(b) => ExprKind::BoolLit(b),
                     ValueKind::NaturalLit(n) => ExprKind::NaturalLit(n),
@@ -501,6 +522,9 @@ impl ValueKind<Value> {
             ValueKind::PartialExpr(e) => {
                 e.map_mut(Value::normalize_mut);
             }
+            ValueKind::Thunk(..) => {
+                unreachable!("Should only normalize_mut values in WHNF")
+            }
         }
     }
 
@@ -553,6 +577,22 @@ impl<V> ValueKind<V> {
     }
 }
 
+impl Thunk {
+    pub fn new(env: &NzEnv, body: TyExpr) -> Self {
+        Thunk {
+            env: env.clone(),
+            body,
+        }
+    }
+
+    pub fn eval(&self) -> Value {
+        normalize_tyexpr_whnf(&self.body, &self.env)
+    }
+    pub fn to_tyexpr(&self, venv: VarEnv) -> TyExpr {
+        self.eval().to_tyexpr(venv)
+    }
+}
+
 impl Closure {
     pub fn new(arg_ty: Value, env: &NzEnv, body: TyExpr) -> Self {
         Closure::Closure {
@@ -572,11 +612,9 @@ impl Closure {
     pub fn apply(&self, val: Value) -> Value {
         match self {
             Closure::Closure { env, body, .. } => {
-                body.normalize_whnf(&env.insert_value(val))
-            }
-            Closure::ConstantClosure { env, body, .. } => {
-                body.normalize_whnf(env)
+                body.eval(&env.insert_value(val))
             }
+            Closure::ConstantClosure { env, body, .. } => body.eval(env),
         }
     }
     fn apply_var(&self, var: NzVar) -> Value {
@@ -588,9 +626,7 @@ impl Closure {
                 );
                 self.apply(val)
             }
-            Closure::ConstantClosure { env, body, .. } => {
-                body.normalize_whnf(env)
-            }
+            Closure::ConstantClosure { env, body, .. } => body.eval(env),
         }
     }
 
@@ -626,6 +662,15 @@ impl std::cmp::PartialEq for Value {
 }
 impl std::cmp::Eq for Value {}
 
+impl std::cmp::PartialEq for Thunk {
+    fn eq(&self, _other: &Self) -> bool {
+        unreachable!(
+            "Trying to compare thunks but we should only compare WHNFs"
+        )
+    }
+}
+impl std::cmp::Eq for Thunk {}
+
 impl std::cmp::PartialEq for Closure {
     fn eq(&self, other: &Self) -> bool {
         let v = NzVar::fresh();
diff --git a/dhall/src/semantics/core/visitor.rs b/dhall/src/semantics/core/visitor.rs
index a78c219..d1a85d8 100644
--- a/dhall/src/semantics/core/visitor.rs
+++ b/dhall/src/semantics/core/visitor.rs
@@ -133,6 +133,7 @@ where
         ),
         Equivalence(x, y) => Equivalence(v.visit_val(x)?, v.visit_val(y)?),
         PartialExpr(e) => PartialExpr(e.traverse_ref(|e| v.visit_val(e))?),
+        Thunk(th) => Thunk(th.clone()),
     })
 }
 
diff --git a/dhall/src/semantics/phase/mod.rs b/dhall/src/semantics/phase/mod.rs
index f24a668..a25f096 100644
--- a/dhall/src/semantics/phase/mod.rs
+++ b/dhall/src/semantics/phase/mod.rs
@@ -91,7 +91,7 @@ impl Resolved {
 impl Typed {
     /// Reduce an expression to its normal form, performing beta reduction
     pub fn normalize(&self) -> Normalized {
-        Normalized(self.0.normalize_nf_noenv())
+        Normalized(self.0.rec_eval_closed_expr())
     }
 
     /// Converts a value back to the corresponding AST expression.
diff --git a/dhall/src/semantics/phase/normalize.rs b/dhall/src/semantics/phase/normalize.rs
index f36ec4a..ea1a25b 100644
--- a/dhall/src/semantics/phase/normalize.rs
+++ b/dhall/src/semantics/phase/normalize.rs
@@ -457,6 +457,7 @@ pub(crate) fn normalize_whnf(
     match v {
         ValueKind::AppliedBuiltin(closure) => closure.ensure_whnf(ty),
         ValueKind::PartialExpr(e) => normalize_one_layer(e, ty, &NzEnv::new()),
+        ValueKind::Thunk(th) => th.eval().kind().clone(),
         ValueKind::TextLit(elts) => {
             ValueKind::TextLit(squash_textlit(elts.into_iter()))
         }
@@ -475,7 +476,7 @@ pub(crate) fn normalize_tyexpr_whnf(tye: &TyExpr, env: &NzEnv) -> Value {
     let kind = match tye.kind() {
         TyExprKind::Var(var) => return env.lookup_val(var),
         TyExprKind::Expr(ExprKind::Lam(binder, annot, body)) => {
-            let annot = annot.normalize_whnf(env);
+            let annot = annot.eval(env);
             ValueKind::LamClosure {
                 binder: Binder::new(binder.clone()),
                 annot: annot.clone(),
@@ -483,7 +484,7 @@ pub(crate) fn normalize_tyexpr_whnf(tye: &TyExpr, env: &NzEnv) -> Value {
             }
         }
         TyExprKind::Expr(ExprKind::Pi(binder, annot, body)) => {
-            let annot = annot.normalize_whnf(env);
+            let annot = annot.eval(env);
             let closure = Closure::new(annot.clone(), env, body.clone());
             ValueKind::PiClosure {
                 binder: Binder::new(binder.clone()),
@@ -492,11 +493,11 @@ pub(crate) fn normalize_tyexpr_whnf(tye: &TyExpr, env: &NzEnv) -> Value {
             }
         }
         TyExprKind::Expr(ExprKind::Let(_, None, val, body)) => {
-            let val = val.normalize_whnf(env);
-            return body.normalize_whnf(&env.insert_value(val));
+            let val = val.eval(env);
+            return body.eval(&env.insert_value(val));
         }
         TyExprKind::Expr(e) => {
-            let e = e.map_ref(|tye| tye.normalize_whnf(env));
+            let e = e.map_ref(|tye| tye.eval(env));
             normalize_one_layer(e, &ty, env)
         }
     };
diff --git a/dhall/src/semantics/tck/tyexpr.rs b/dhall/src/semantics/tck/tyexpr.rs
index 5492b8b..0a27f32 100644
--- a/dhall/src/semantics/tck/tyexpr.rs
+++ b/dhall/src/semantics/tck/tyexpr.rs
@@ -49,6 +49,9 @@ impl TyExpr {
     pub fn kind(&self) -> &TyExprKind {
         &*self.kind
     }
+    pub fn span(&self) -> &Span {
+        &self.span
+    }
     pub fn get_type(&self) -> Result<Type, TypeError> {
         match &self.ty {
             Some(t) => Ok(t.clone()),
@@ -69,17 +72,18 @@ impl TyExpr {
         tyexpr_to_expr(self, opts, &mut env)
     }
 
-    pub fn normalize_whnf(&self, env: &NzEnv) -> Value {
-        normalize_tyexpr_whnf(self, env)
-    }
-    pub fn normalize_whnf_noenv(&self) -> Value {
-        self.normalize_whnf(&NzEnv::new())
+    /// Eval the TyExpr. It will actually get evaluated only as needed on demand.
+    pub fn eval(&self, env: &NzEnv) -> Value {
+        Value::new_thunk(env, self.clone())
     }
-    pub fn normalize_nf(&self, env: &NzEnv) -> Value {
-        self.normalize_whnf(env)
+    /// Eval a closed TyExpr (i.e. without free variables). It will actually get evaluated only as
+    /// needed on demand.
+    pub fn eval_closed_expr(&self) -> Value {
+        self.eval(&NzEnv::new())
     }
-    pub fn normalize_nf_noenv(&self) -> Value {
-        self.normalize_nf(&NzEnv::new())
+    /// Eval a closed TyExpr fully and recursively;
+    pub fn rec_eval_closed_expr(&self) -> Value {
+        normalize_tyexpr_whnf(self, &NzEnv::new())
     }
 }
 
diff --git a/dhall/src/semantics/tck/typecheck.rs b/dhall/src/semantics/tck/typecheck.rs
index bc1c87f..d741b40 100644
--- a/dhall/src/semantics/tck/typecheck.rs
+++ b/dhall/src/semantics/tck/typecheck.rs
@@ -74,7 +74,7 @@ fn type_one_layer(
         ExprKind::Builtin(b) => {
             let t_expr = type_of_builtin(*b);
             let t_tyexpr = type_with(env, &t_expr)?;
-            t_tyexpr.normalize_whnf(env.as_nzenv())
+            t_tyexpr.eval(env.as_nzenv())
         }
         ExprKind::BoolLit(_) => Value::from_builtin(Builtin::Bool),
         ExprKind::NaturalLit(_) => Value::from_builtin(Builtin::Natural),
@@ -93,7 +93,7 @@ fn type_one_layer(
             text_type
         }
         ExprKind::EmptyListLit(t) => {
-            let t = t.normalize_nf(env.as_nzenv());
+            let t = t.eval(env.as_nzenv());
             match &*t.kind() {
                 ValueKind::AppliedBuiltin(BuiltinClosure {
                     b: Builtin::List,
@@ -196,7 +196,7 @@ fn type_one_layer(
                 },
                 // TODO: branch here only when scrut.get_type() is a Const
                 _ => {
-                    let scrut_nf = scrut.normalize_nf(env.as_nzenv());
+                    let scrut_nf = scrut.eval(env.as_nzenv());
                     let scrut_nf_borrow = scrut_nf.kind();
                     match &*scrut_nf_borrow {
                         ValueKind::UnionType(kts) => match kts.get(x) {
@@ -224,7 +224,7 @@ fn type_one_layer(
             }
         }
         ExprKind::Annot(x, t) => {
-            let t = t.normalize_whnf(env.as_nzenv());
+            let t = t.eval(env.as_nzenv());
             let x_ty = x.get_type()?;
             if x_ty != t {
                 return mkerr(format!(
@@ -243,7 +243,7 @@ fn type_one_layer(
             x_ty
         }
         ExprKind::Assert(t) => {
-            let t = t.normalize_whnf(env.as_nzenv());
+            let t = t.eval(env.as_nzenv());
             match &*t.kind() {
                 ValueKind::Equivalence(x, y) if x == y => {}
                 ValueKind::Equivalence(..) => return mkerr("AssertMismatch"),
@@ -268,7 +268,7 @@ fn type_one_layer(
                         ));
                     }
 
-                    let arg_nf = arg.normalize_nf(env.as_nzenv());
+                    let arg_nf = arg.eval(env.as_nzenv());
                     closure.apply(arg_nf)
                 }
                 _ => return mkerr(format!("apply to not Pi")),
@@ -329,11 +329,11 @@ fn type_one_layer(
             );
             let ty = type_one_layer(env, &ekind)?;
             TyExpr::new(TyExprKind::Expr(ekind), Some(ty), Span::Artificial)
-                .normalize_nf(env.as_nzenv())
+                .eval(env.as_nzenv())
         }
         ExprKind::BinOp(BinOp::RecursiveRecordTypeMerge, x, y) => {
-            let x_val = x.normalize_whnf(env.as_nzenv());
-            let y_val = y.normalize_whnf(env.as_nzenv());
+            let x_val = x.eval(env.as_nzenv());
+            let y_val = y.eval(env.as_nzenv());
             let x_val_borrow = x_val.kind();
             let y_val_borrow = y_val.kind();
             let kts_x = match &*x_val_borrow {
@@ -478,9 +478,8 @@ fn type_one_layer(
                 }
             }
 
-            let type_annot = type_annot
-                .as_ref()
-                .map(|t| t.normalize_whnf(env.as_nzenv()));
+            let type_annot =
+                type_annot.as_ref().map(|t| t.eval(env.as_nzenv()));
             match (inferred_type, type_annot) {
                 (Some(t1), Some(t2)) => {
                     if t1 != t2 {
@@ -542,7 +541,7 @@ pub(crate) fn type_with(
         },
         ExprKind::Lam(binder, annot, body) => {
             let annot = type_with(env, annot)?;
-            let annot_nf = annot.normalize_nf(env.as_nzenv());
+            let annot_nf = annot.eval(env.as_nzenv());
             let body_env = env.insert_type(&binder, annot_nf.clone());
             let body = type_with(&body_env, body)?;
             let body_ty = body.get_type()?;
@@ -555,7 +554,7 @@ pub(crate) fn type_with(
                 Some(type_of_function(annot.get_type()?, body_ty.get_type()?)?),
                 Span::Artificial,
             );
-            let ty = ty.normalize_whnf(env.as_nzenv());
+            let ty = ty.eval(env.as_nzenv());
             (
                 TyExprKind::Expr(ExprKind::Lam(binder.clone(), annot, body)),
                 Some(ty),
@@ -563,7 +562,7 @@ pub(crate) fn type_with(
         }
         ExprKind::Pi(binder, annot, body) => {
             let annot = type_with(env, annot)?;
-            let annot_nf = annot.normalize_whnf(env.as_nzenv());
+            let annot_nf = annot.eval(env.as_nzenv());
             let body =
                 type_with(&env.insert_type(binder, annot_nf.clone()), body)?;
             let ty = type_of_function(annot.get_type()?, body.get_type()?)?;
@@ -580,7 +579,7 @@ pub(crate) fn type_with(
             };
 
             let val = type_with(env, &val)?;
-            let val_nf = val.normalize_nf(&env.as_nzenv());
+            let val_nf = val.eval(&env.as_nzenv());
             let body = type_with(&env.insert_value(&binder, val_nf), body)?;
             let body_ty = body.get_type().ok();
             (