Don't discard normalization work done by typechecking

2 files changed, 54 insertions, 31 deletions

diff --git a/dhall/src/core/context.rs b/dhall/src/core/context.rs
index aeec7fb..56dec03 100644
--- a/dhall/src/core/context.rs
+++ b/dhall/src/core/context.rs
@@ -1,4 +1,3 @@
-use std::borrow::Cow;
 use std::rc::Rc;
 
 use dhall_syntax::context::Context as SimpleContext;
@@ -25,28 +24,10 @@ pub(crate) struct NormalizationContext(Context<()>);
 #[derive(Debug, Clone)]
 pub(crate) struct TypecheckContext(Context<Type>);
 
-impl<T> CtxItem<T> {
-    fn forget(&self) -> CtxItem<()> {
-        match self {
-            CtxItem::Kept(var, _) => CtxItem::Kept(var.clone(), ()),
-            CtxItem::Replaced(e, _) => CtxItem::Replaced(e.clone(), ()),
-        }
-    }
-}
-
 impl<T> Context<T> {
     pub(crate) fn new() -> Self {
         Context(Rc::new(SimpleContext::new()))
     }
-    pub(crate) fn forget(&self) -> Context<()> {
-        let mut ctx = SimpleContext::new();
-        for (k, vs) in self.0.iter_keys() {
-            for v in vs.iter() {
-                ctx = ctx.insert(k.clone(), v.forget());
-            }
-        }
-        Context(Rc::new(ctx))
-    }
     pub(crate) fn insert_kept(&self, x: &Label, t: T) -> Self
     where
         T: Shift + Clone,
@@ -103,16 +84,18 @@ impl TypecheckContext {
             e.get_type()?.into_owned(),
         )))
     }
-    pub(crate) fn lookup(&self, var: &V<Label>) -> Option<Cow<'_, Type>> {
+    pub(crate) fn lookup(&self, var: &V<Label>) -> Option<Typed> {
         match self.0.lookup(var) {
-            Some(CtxItem::Kept(_, t)) => Some(Cow::Borrowed(&t)),
-            Some(CtxItem::Replaced(_, t)) => Some(Cow::Borrowed(&t)),
+            Some(CtxItem::Kept(newvar, t)) => Some(Typed::from_thunk_and_type(
+                Thunk::from_value(Value::Var(newvar.clone())),
+                t.clone(),
+            )),
+            Some(CtxItem::Replaced(th, t)) => {
+                Some(Typed::from_thunk_and_type(th.clone(), t.clone()))
+            }
             None => None,
         }
     }
-    pub(crate) fn to_normalization_ctx(&self) -> NormalizationContext {
-        NormalizationContext(self.0.forget())
-    }
 }
 
 impl<T: Shift> Shift for CtxItem<T> {
diff --git a/dhall/src/core/thunk.rs b/dhall/src/core/thunk.rs
index b3817dc..530762b 100644
--- a/dhall/src/core/thunk.rs
+++ b/dhall/src/core/thunk.rs
@@ -1,13 +1,15 @@
 use std::cell::{Ref, RefCell};
 use std::rc::Rc;
 
+use dhall_syntax::{ExprF, Label, X};
+
 use crate::core::context::NormalizationContext;
 use crate::core::context::TypecheckContext;
 use crate::core::value::Value;
 use crate::core::var::{AlphaVar, Shift, Subst};
 use crate::error::TypeError;
 use crate::phase::normalize::{
-    apply_any, normalize_whnf, InputSubExpr, OutputSubExpr,
+    apply_any, normalize_one_layer, normalize_whnf, InputSubExpr, OutputSubExpr,
 };
 use crate::phase::typecheck::mktype;
 use crate::phase::{Type, Typed};
@@ -25,6 +27,10 @@ use Marker::{NF, WHNF};
 enum ThunkInternal {
     /// Non-normalized value alongside a normalization context
     Unnormalized(NormalizationContext, InputSubExpr),
+    /// Partially normalized value whose subexpressions have been thunked (this is returned from
+    /// typechecking). Note that this is different from `Value::PartialExpr` because there is no
+    /// requirement of WHNF here.
+    PartialExpr(ExprF<Thunk, Label, X>),
     /// Partially normalized value.
     /// Invariant: if the marker is `NF`, the value must be fully normalized
     Value(Marker, Value),
@@ -57,6 +63,10 @@ impl ThunkInternal {
                     normalize_whnf(ctx.clone(), e.clone()),
                 )
             }
+            ThunkInternal::PartialExpr(e) => {
+                *self =
+                    ThunkInternal::Value(WHNF, normalize_one_layer(e.clone()))
+            }
             // Already at least in WHNF
             ThunkInternal::Value(_, _) => {}
         }
@@ -64,7 +74,8 @@ impl ThunkInternal {
 
     fn normalize_nf(&mut self) {
         match self {
-            ThunkInternal::Unnormalized(_, _) => {
+            ThunkInternal::Unnormalized(_, _)
+            | ThunkInternal::PartialExpr(_) => {
                 self.normalize_whnf();
                 self.normalize_nf();
             }
@@ -80,7 +91,8 @@ impl ThunkInternal {
     // Always use normalize_whnf before
     fn as_whnf(&self) -> &Value {
         match self {
-            ThunkInternal::Unnormalized(_, _) => unreachable!(),
+            ThunkInternal::Unnormalized(_, _)
+            | ThunkInternal::PartialExpr(_) => unreachable!(),
             ThunkInternal::Value(_, v) => v,
         }
     }
@@ -88,8 +100,9 @@ impl ThunkInternal {
     // Always use normalize_nf before
     fn as_nf(&self) -> &Value {
         match self {
-            ThunkInternal::Unnormalized(_, _) => unreachable!(),
-            ThunkInternal::Value(WHNF, _) => unreachable!(),
+            ThunkInternal::Unnormalized(_, _)
+            | ThunkInternal::PartialExpr(_)
+            | ThunkInternal::Value(WHNF, _) => unreachable!(),
             ThunkInternal::Value(NF, v) => v,
         }
     }
@@ -108,6 +121,10 @@ impl Thunk {
         Thunk::new(NormalizationContext::new(), e.absurd())
     }
 
+    pub(crate) fn from_partial_expr(e: ExprF<Thunk, Label, X>) -> Thunk {
+        ThunkInternal::PartialExpr(e).into_thunk()
+    }
+
     // Normalizes contents to normal form; faster than `normalize_nf` if
     // no one else shares this thunk
     pub(crate) fn normalize_mut(&mut self) {
@@ -122,7 +139,8 @@ impl Thunk {
     fn do_normalize_whnf(&self) {
         let borrow = self.0.borrow();
         match &*borrow {
-            ThunkInternal::Unnormalized(_, _) => {
+            ThunkInternal::Unnormalized(_, _)
+            | ThunkInternal::PartialExpr(_) => {
                 drop(borrow);
                 self.0.borrow_mut().normalize_whnf();
             }
@@ -135,6 +153,7 @@ impl Thunk {
         let borrow = self.0.borrow();
         match &*borrow {
             ThunkInternal::Unnormalized(_, _)
+            | ThunkInternal::PartialExpr(_)
             | ThunkInternal::Value(WHNF, _) => {
                 drop(borrow);
                 self.0.borrow_mut().normalize_nf();
@@ -249,6 +268,14 @@ impl Shift for ThunkInternal {
             ThunkInternal::Unnormalized(ctx, e) => {
                 ThunkInternal::Unnormalized(ctx.shift(delta, var), e.clone())
             }
+            ThunkInternal::PartialExpr(e) => ThunkInternal::PartialExpr(
+                e.map_ref_with_special_handling_of_binders(
+                    |v| v.shift(delta, var),
+                    |x, v| v.shift(delta, &var.shift(1, &x.into())),
+                    X::clone,
+                    Label::clone,
+                ),
+            ),
             ThunkInternal::Value(m, v) => {
                 ThunkInternal::Value(*m, v.shift(delta, var))
             }
@@ -278,6 +305,19 @@ impl Subst<Typed> for ThunkInternal {
                 ctx.subst_shift(var, val),
                 e.clone(),
             ),
+            ThunkInternal::PartialExpr(e) => ThunkInternal::PartialExpr(
+                e.map_ref_with_special_handling_of_binders(
+                    |v| v.subst_shift(var, val),
+                    |x, v| {
+                        v.subst_shift(
+                            &var.shift(1, &x.into()),
+                            &val.shift(1, &x.into()),
+                        )
+                    },
+                    X::clone,
+                    Label::clone,
+                ),
+            ),
             ThunkInternal::Value(_, v) => {
                 // The resulting value may not stay in normal form after substitution
                 ThunkInternal::Value(WHNF, v.subst_shift(var, val))