Store context in Typed

3 files changed, 39 insertions, 33 deletions

diff --git a/dhall/src/expr.rs b/dhall/src/expr.rs
index 44d1612..1840eba 100644
--- a/dhall/src/expr.rs
+++ b/dhall/src/expr.rs
@@ -10,6 +10,8 @@ macro_rules! derive_other_traits {
             }
         }
 
+        impl<'a> std::cmp::Eq for $ty<'a> {}
+
         impl<'a> std::fmt::Display for $ty<'a> {
             fn fmt(
                 &self,
@@ -21,28 +23,29 @@ macro_rules! derive_other_traits {
     };
 }
 
-#[derive(Debug, Clone, Eq)]
+#[derive(Debug, Clone)]
 pub(crate) struct Parsed<'a>(
     pub(crate) SubExpr<Span<'a>, Import>,
     pub(crate) ImportRoot,
 );
 derive_other_traits!(Parsed);
 
-#[derive(Debug, Clone, Eq)]
+#[derive(Debug, Clone)]
 pub(crate) struct Resolved<'a>(
     pub(crate) SubExpr<Span<'a>, Normalized<'static>>,
 );
 derive_other_traits!(Resolved);
 
-#[derive(Debug, Clone, Eq)]
+#[derive(Debug, Clone)]
 pub(crate) struct Typed<'a>(
     pub(crate) SubExpr<X, Normalized<'static>>,
     pub(crate) Option<Type<'static>>,
+    pub(crate) crate::typecheck::TypecheckContext,
     pub(crate) PhantomData<&'a ()>,
 );
 derive_other_traits!(Typed);
 
-#[derive(Debug, Clone, Eq)]
+#[derive(Debug, Clone)]
 pub(crate) struct Normalized<'a>(
     pub(crate) SubExpr<X, X>,
     pub(crate) Option<Type<'static>>,
@@ -56,7 +59,7 @@ derive_other_traits!(Normalized);
 /// `Bool`, `{ x: Integer }` or `Natural -> Text`.
 ///
 /// For a more general notion of "type", see [Type].
-#[derive(Debug, Clone, Eq)]
+#[derive(Debug, Clone)]
 pub struct SimpleType<'a>(
     pub(crate) SubExpr<X, X>,
     pub(crate) PhantomData<&'a ()>,
@@ -101,7 +104,12 @@ impl<'a> From<SubExpr<X, X>> for SimpleType<'a> {
 #[doc(hidden)]
 impl<'a> From<Normalized<'a>> for Typed<'a> {
     fn from(x: Normalized<'a>) -> Typed<'a> {
-        Typed(x.0.embed_absurd(), x.1, x.2)
+        Typed(
+            x.0.embed_absurd(),
+            x.1,
+            crate::typecheck::TypecheckContext::new(),
+            x.2,
+        )
     }
 }
 
diff --git a/dhall/src/normalize.rs b/dhall/src/normalize.rs
index 0301e35..35a6b23 100644
--- a/dhall/src/normalize.rs
+++ b/dhall/src/normalize.rs
@@ -16,20 +16,13 @@ type OutputSubExpr = SubExpr<X, X>;
 
 impl<'a> Typed<'a> {
     pub fn normalize(self) -> Normalized<'a> {
-        Normalized(normalize(self.0), self.1, self.2)
-    }
-    pub fn normalize_ctx(
-        self,
-        ctx: &crate::typecheck::TypecheckContext,
-    ) -> Normalized<'a> {
         Normalized(
-            normalize_whnf(
-                NormalizationContext::from_typecheck_ctx(ctx),
+            normalize(
+                NormalizationContext::from_typecheck_ctx(&self.2),
                 self.0,
-            )
-            .normalize_to_expr(),
+            ),
             self.1,
-            self.2,
+            self.3,
         )
     }
     /// Pretends this expression is normalized. Use with care.
@@ -38,7 +31,7 @@ impl<'a> Typed<'a> {
         Normalized(
             self.0.unroll().squash_embed(|e| e.0.clone()),
             self.1,
-            self.2,
+            self.3,
         )
     }
 }
@@ -832,8 +825,8 @@ fn normalize_last_layer(expr: ExprF<WHNF, Label, X, X>) -> WHNF {
 /// However, `normalize` will not fail if the expression is ill-typed and will
 /// leave ill-typed sub-expressions unevaluated.
 ///
-fn normalize(e: InputSubExpr) -> OutputSubExpr {
-    normalize_whnf(NormalizationContext::new(), e).normalize_to_expr()
+fn normalize(ctx: NormalizationContext, e: InputSubExpr) -> OutputSubExpr {
+    normalize_whnf(ctx, e).normalize_to_expr()
 }
 
 #[cfg(test)]
diff --git a/dhall/src/typecheck.rs b/dhall/src/typecheck.rs
index e397316..a53e88a 100644
--- a/dhall/src/typecheck.rs
+++ b/dhall/src/typecheck.rs
@@ -28,7 +28,7 @@ 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, PhantomData)
+        Typed(self.0.unnote(), None, TypecheckContext::new(), PhantomData)
     }
 }
 impl<'a> Typed<'a> {
@@ -420,7 +420,7 @@ fn mktype(
     ctx: &TypecheckContext,
     e: SubExpr<X, Normalized<'static>>,
 ) -> Result<Type<'static>, TypeError> {
-    Ok(type_with(ctx, e)?.normalize_ctx(ctx).into_type())
+    Ok(type_with(ctx, e)?.normalize().into_type())
 }
 
 fn into_simple_type<'a>(e: SubExpr<X, X>) -> Type<'a> {
@@ -498,7 +498,7 @@ fn type_with(
                 v.clone()
             };
 
-            let v = type_with(ctx, v)?.normalize_ctx(ctx);
+            let v = type_with(ctx, v)?.normalize();
             let e = type_with(&ctx.insert_value(x, v.clone()), e.clone())?;
 
             Ok(RetType(e.get_type_move()?))
@@ -513,8 +513,13 @@ fn type_with(
         ),
     }?;
     match ret {
-        RetExpr(ret) => Ok(Typed(e, Some(mktype(ctx, rc(ret))?), PhantomData)),
-        RetType(typ) => Ok(Typed(e, Some(typ), PhantomData)),
+        RetExpr(ret) => Ok(Typed(
+            e,
+            Some(mktype(ctx, rc(ret))?),
+            ctx.clone(),
+            PhantomData,
+        )),
+        RetType(typ) => Ok(Typed(e, Some(typ), ctx.clone(), PhantomData)),
     }
 }
 
@@ -556,12 +561,12 @@ fn type_last_layer(
             Ok(RetExpr(Let(
                 x.clone(),
                 None,
-                a.normalize_ctx(ctx).embed(),
+                a.normalize().embed(),
                 tb.embed_absurd(),
             )))
         }
         Annot(x, t) => {
-            let t = t.normalize_ctx(ctx).into_type();
+            let t = t.normalize().into_type();
             ensure_equal!(
                 &t,
                 x.get_type()?,
@@ -595,7 +600,7 @@ fn type_last_layer(
             Ok(RetType(y.get_type_move()?))
         }
         EmptyListLit(t) => {
-            let t = t.normalize_ctx(ctx).into_type();
+            let t = t.normalize().into_type();
             ensure_simple_type!(
                 t,
                 mkerr(InvalidListType(t.into_normalized()?)),
@@ -626,13 +631,13 @@ fn type_last_layer(
             Ok(RetExpr(dhall::expr!(List t)))
         }
         OldOptionalLit(None, t) => {
-            let t = t.normalize_ctx(ctx).embed();
+            let t = t.normalize().embed();
             let e = dhall::subexpr!(None t);
             Ok(RetType(type_with(ctx, e)?.get_type()?.into_owned()))
         }
         OldOptionalLit(Some(x), t) => {
-            let t = t.normalize_ctx(ctx).embed();
-            let x = x.normalize_ctx(ctx).embed();
+            let t = t.normalize().embed();
+            let x = x.normalize().embed();
             let e = dhall::subexpr!(Some x : Optional t);
             Ok(RetType(type_with(ctx, e)?.get_type()?.into_owned()))
         }
@@ -702,7 +707,7 @@ fn type_last_layer(
             let mut kts: std::collections::BTreeMap<_, _> = kvs
                 .into_iter()
                 .map(|(x, v)| {
-                    let t = v.map(|x| x.normalize_ctx(ctx).embed());
+                    let t = v.map(|x| x.normalize().embed());
                     Ok((x, t))
                 })
                 .collect::<Result<_, _>>()?;
@@ -716,7 +721,7 @@ fn type_last_layer(
                 None => Err(mkerr(MissingRecordField(x, r))),
             },
             _ => {
-                let r = r.normalize_ctx(ctx);
+                let r = r.normalize();
                 match r.as_expr().as_ref() {
                     UnionType(kts) => match kts.get(&x) {
                         // Constructor has type T -> < x: T, ... >