4 files changed, 22 insertions, 40 deletions
diff --git a/dhall/src/expr.rs b/dhall/src/expr.rs
index b2f2bec..c794324 100644
--- a/dhall/src/expr.rs
+++ b/dhall/src/expr.rs
@@ -38,12 +38,10 @@ derive_other_traits!(Resolved);
 
 #[derive(Debug, Clone)]
 pub(crate) struct Typed<'a>(
-    pub(crate) SubExpr<X, Normalized<'static>>,
+    pub(crate) crate::normalize::Thunk,
     pub(crate) Option<Type<'static>>,
-    pub(crate) crate::typecheck::TypecheckContext,
     pub(crate) PhantomData<&'a ()>,
 );
-derive_other_traits!(Typed);
 
 #[derive(Debug, Clone)]
 pub(crate) struct PartiallyNormalized<'a>(
@@ -112,23 +110,17 @@ impl<'a> From<SubExpr<X, X>> for SimpleType<'a> {
 impl<'a> From<Normalized<'a>> for Typed<'a> {
     fn from(x: Normalized<'a>) -> Typed<'a> {
         Typed(
-            x.0.embed_absurd(),
+            crate::normalize::Thunk::new(
+                crate::normalize::NormalizationContext::new(),
+                x.0.embed_absurd(),
+            ),
             x.1,
-            crate::typecheck::TypecheckContext::new(),
             x.2,
         )
     }
 }
 
 #[doc(hidden)]
-#[allow(dead_code)]
-impl<'a> Typed<'a> {
-    pub(crate) fn as_expr(&self) -> &SubExpr<X, Normalized<'a>> {
-        &self.0
-    }
-}
-
-#[doc(hidden)]
 impl<'a> Normalized<'a> {
     pub(crate) fn as_expr(&self) -> &SubExpr<X, X> {
         &self.0
diff --git a/dhall/src/normalize.rs b/dhall/src/normalize.rs
index f2ad8ea..5151790 100644
--- a/dhall/src/normalize.rs
+++ b/dhall/src/normalize.rs
@@ -28,23 +28,7 @@ impl<'a> Typed<'a> {
         self.partially_normalize().normalize()
     }
     pub(crate) fn partially_normalize(self) -> PartiallyNormalized<'a> {
-        PartiallyNormalized(
-            normalize_whnf(
-                NormalizationContext::from_typecheck_ctx(&self.2),
-                self.0,
-            ),
-            self.1,
-            self.3,
-        )
-    }
-    /// Pretends this expression is normalized. Use with care.
-    #[allow(dead_code)]
-    pub fn skip_normalize(self) -> Normalized<'a> {
-        Normalized(
-            self.0.unroll().squash_embed(|e| e.0.clone()),
-            self.1,
-            self.3,
-        )
+        PartiallyNormalized(self.0.as_whnf(), self.1, self.2)
     }
 }
 
@@ -300,7 +284,7 @@ impl EnvItem {
 pub(crate) struct NormalizationContext(Rc<Context<Label, EnvItem>>);
 
 impl NormalizationContext {
-    fn new() -> Self {
+    pub(crate) fn new() -> Self {
         NormalizationContext(Rc::new(Context::new()))
     }
     fn insert(&self, x: &Label, e: Thunk) -> Self {
@@ -323,7 +307,9 @@ impl NormalizationContext {
             None => Value::Var(var.clone()),
         }
     }
-    fn from_typecheck_ctx(tc_ctx: &crate::typecheck::TypecheckContext) -> Self {
+    pub(crate) fn from_typecheck_ctx(
+        tc_ctx: &crate::typecheck::TypecheckContext,
+    ) -> Self {
         use crate::typecheck::EnvItem::*;
         let mut ctx = Context::new();
         for (k, vs) in tc_ctx.0.iter_keys() {
diff --git a/dhall/src/tests.rs b/dhall/src/tests.rs
index 4d8fabd..ebca256 100644
--- a/dhall/src/tests.rs
+++ b/dhall/src/tests.rs
@@ -130,8 +130,7 @@ pub fn run_test(
             match feature {
                 Parser => unreachable!(),
                 Import => {
-                    // Not sure whether to normalize or not
-                    let expr = expr.skip_typecheck().skip_normalize();
+                    let expr = expr.skip_typecheck().normalize();
                     assert_eq_display!(expr, expected);
                 }
                 Typecheck => {
diff --git a/dhall/src/typecheck.rs b/dhall/src/typecheck.rs
index d924f41..85de42a 100644
--- a/dhall/src/typecheck.rs
+++ b/dhall/src/typecheck.rs
@@ -6,7 +6,7 @@ use std::fmt;
 use std::marker::PhantomData;
 
 use crate::expr::*;
-use crate::normalize::{TypeThunk, Value};
+use crate::normalize::{NormalizationContext, Thunk, TypeThunk, Value};
 use crate::traits::DynamicType;
 use dhall_core;
 use dhall_core::context::Context;
@@ -30,7 +30,11 @@ impl<'a> Resolved<'a> {
     /// Pretends this expression has been typechecked. Use with care.
     #[allow(dead_code)]
     pub fn skip_typecheck(self) -> Typed<'a> {
-        Typed(self.0.unnote(), None, TypecheckContext::new(), PhantomData)
+        Typed(
+            Thunk::new(NormalizationContext::new(), self.0.unnote()),
+            None,
+            PhantomData,
+        )
     }
 }
 impl<'a> Typed<'a> {
@@ -325,6 +329,9 @@ impl TypecheckContext {
             None => None,
         }
     }
+    fn to_normalization_ctx(&self) -> NormalizationContext {
+        NormalizationContext::from_typecheck_ctx(self)
+    }
 }
 
 impl PartialEq for TypecheckContext {
@@ -813,15 +820,13 @@ fn type_with(
     }?;
     match ret {
         RetExpr(ret) => Ok(TypedOrType::Typed(Typed(
-            e,
+            Thunk::new(ctx.to_normalization_ctx(), e),
             Some(mktype(ctx, rc(ret))?),
-            ctx.clone(),
             PhantomData,
         ))),
         RetType(typ) => Ok(TypedOrType::Typed(Typed(
-            e,
+            Thunk::new(ctx.to_normalization_ctx(), e),
             Some(typ),
-            ctx.clone(),
             PhantomData,
         ))),
         RetTypedOrType(tt) => Ok(tt),