Store Thunk in Normalized

5 files changed, 44 insertions, 41 deletions

diff --git a/dhall/src/expr.rs b/dhall/src/expr.rs
index 1ddc4ba..a753ffd 100644
--- a/dhall/src/expr.rs
+++ b/dhall/src/expr.rs
@@ -45,11 +45,24 @@ pub(crate) struct Typed<'a>(
 
 #[derive(Debug, Clone)]
 pub(crate) struct Normalized<'a>(
-    pub(crate) SubExpr<X, X>,
+    pub(crate) crate::normalize::Thunk,
     pub(crate) Option<Type<'static>>,
     pub(crate) PhantomData<&'a ()>,
 );
-derive_other_traits!(Normalized);
+
+impl<'a> std::cmp::PartialEq for Normalized<'a> {
+    fn eq(&self, other: &Self) -> bool {
+        self.0.normalize_to_expr() == other.0.normalize_to_expr()
+    }
+}
+
+impl<'a> std::cmp::Eq for Normalized<'a> {}
+
+impl<'a> std::fmt::Display for Normalized<'a> {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
+        self.0.normalize_to_expr().fmt(f)
+    }
+}
 
 /// A Dhall expression representing a simple type.
 ///
@@ -102,21 +115,20 @@ 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(
-            crate::normalize::Thunk::new(
-                crate::normalize::NormalizationContext::new(),
-                x.0.embed_absurd(),
-            ),
-            x.1,
-            x.2,
-        )
+        Typed(x.0, x.1, x.2)
     }
 }
 
-#[doc(hidden)]
 impl<'a> Normalized<'a> {
-    pub(crate) fn as_expr(&self) -> &SubExpr<X, X> {
-        &self.0
+    // Deprecated
+    pub(crate) fn as_expr(&self) -> SubExpr<X, X> {
+        self.0.normalize_to_expr()
+    }
+    pub(crate) fn to_expr(&self) -> SubExpr<X, X> {
+        self.0.normalize_to_expr()
+    }
+    pub(crate) fn to_value(&self) -> crate::normalize::Value {
+        self.0.normalize_nf().clone()
     }
     #[allow(dead_code)]
     pub(crate) fn unnote<'b>(self) -> Normalized<'b> {
diff --git a/dhall/src/normalize.rs b/dhall/src/normalize.rs
index 7a69bea..390911a 100644
--- a/dhall/src/normalize.rs
+++ b/dhall/src/normalize.rs
@@ -25,7 +25,11 @@ impl<'a> Typed<'a> {
     /// leave ill-typed sub-expressions unevaluated.
     ///
     pub fn normalize(self) -> Normalized<'a> {
-        Normalized(self.0.normalize_whnf().normalize_to_expr(), self.1, self.2)
+        // TODO: stupid but needed for now
+        let thunk = Thunk::from_normalized_expr(self.0.normalize_to_expr());
+        // let thunk = self.0;
+        thunk.normalize_nf();
+        Normalized(thunk, self.1, self.2)
     }
 
     pub(crate) fn shift0(&self, delta: isize, x: &Label) -> Self {
@@ -1062,6 +1066,10 @@ mod thunk {
             ThunkInternal::Value(WHNF, v).into_thunk()
         }
 
+        pub(crate) fn from_normalized_expr(e: OutputSubExpr) -> Thunk {
+            Thunk::new(NormalizationContext::new(), e.embed_absurd())
+        }
+
         // Deprecated
         pub(crate) fn normalize(&self) -> Thunk {
             self.normalize_nf();
@@ -1189,10 +1197,7 @@ fn normalize_whnf(ctx: NormalizationContext, expr: InputSubExpr) -> Value {
         ExprF::Var(v) => ctx.lookup(&v),
         ExprF::Annot(x, _) => normalize_whnf(ctx, x.clone()),
         ExprF::Note(_, e) => normalize_whnf(ctx, e.clone()),
-        // TODO: wasteful to retraverse everything
-        ExprF::Embed(e) => {
-            normalize_whnf(NormalizationContext::new(), e.0.embed_absurd())
-        }
+        ExprF::Embed(e) => e.to_value(),
         ExprF::Let(x, _, r, b) => {
             let t = Thunk::new(ctx.clone(), r.clone());
             normalize_whnf(ctx.insert(x, t), b.clone())
diff --git a/dhall/src/serde.rs b/dhall/src/serde.rs
index 196bda1..6f143cb 100644
--- a/dhall/src/serde.rs
+++ b/dhall/src/serde.rs
@@ -7,7 +7,7 @@ use std::borrow::Cow;
 impl<'a, T: serde::Deserialize<'a>> Deserialize<'a> for T {
     fn from_str(s: &'a str, ty: Option<&Type>) -> Result<Self> {
         let expr = Normalized::from_str(s, ty)?;
-        T::deserialize(Deserializer(Cow::Owned(expr.0)))
+        T::deserialize(Deserializer(Cow::Owned(expr.to_expr())))
     }
 }
 
diff --git a/dhall/src/traits/static_type.rs b/dhall/src/traits/static_type.rs
index 8b5acbf..4dd9961 100644
--- a/dhall/src/traits/static_type.rs
+++ b/dhall/src/traits/static_type.rs
@@ -39,7 +39,9 @@ fn mktype<'a>(x: SubExpr<X, X>) -> SimpleType<'a> {
 impl<T: SimpleStaticType> StaticType for T {
     fn get_static_type() -> Type<'static> {
         crate::expr::Normalized(
-            T::get_simple_static_type().into(),
+            crate::normalize::Thunk::from_normalized_expr(
+                T::get_simple_static_type().into(),
+            ),
             Some(Type::const_type()),
             std::marker::PhantomData,
         )
diff --git a/dhall/src/typecheck.rs b/dhall/src/typecheck.rs
index e8bc26d..8afbb2b 100644
--- a/dhall/src/typecheck.rs
+++ b/dhall/src/typecheck.rs
@@ -50,30 +50,14 @@ impl<'a> Normalized<'a> {
     }
     fn shift(&self, delta: isize, var: &V<Label>) -> Self {
         Normalized(
-            shift(delta, var, &self.0),
+            self.0.shift(delta, var),
             self.1.as_ref().map(|t| t.shift(delta, var)),
             self.2,
         )
     }
     pub(crate) fn into_type(self) -> Result<Type<'a>, TypeError> {
-        self.into_type_ctx(&TypecheckContext::new())
-    }
-    fn into_type_ctx(
-        self,
-        ctx: &TypecheckContext,
-    ) -> Result<Type<'a>, TypeError> {
-        Ok(match self.0.as_ref() {
-            ExprF::Const(c) => Type(TypeInternal::Const(*c)),
-            ExprF::Pi(_, _, _) | ExprF::RecordType(_) | ExprF::UnionType(_) => {
-                // TODO: wasteful
-                type_with(ctx, self.0.embed_absurd())?.to_type()
-            }
-            _ => Type(TypeInternal::Typed(Box::new(Typed(
-                Value::Expr(self.0).into_thunk(),
-                self.1,
-                self.2,
-            )))),
-        })
+        let typed: Typed = self.into();
+        Ok(typed.to_type())
     }
     fn get_type_move(self) -> Result<Type<'static>, TypeError> {
         self.1.ok_or_else(|| {
@@ -394,8 +378,8 @@ where
             let mut ctx = vec![];
             go(
                 &mut ctx,
-                eL0.borrow().as_normalized().unwrap().as_expr(),
-                eR0.borrow().as_normalized().unwrap().as_expr(),
+                &eL0.borrow().as_normalized().unwrap().to_expr(),
+                &eR0.borrow().as_normalized().unwrap().to_expr(),
             )
         }
         // _ => false,