Preserve notes when normalizing

1 files changed, 13 insertions, 6 deletions

diff --git a/dhall/src/normalize.rs b/dhall/src/normalize.rs
index fa669d8..076f889 100644
--- a/dhall/src/normalize.rs
+++ b/dhall/src/normalize.rs
@@ -221,16 +221,20 @@ enum WhatNext<'a, N, E> {
     // The following expression is the normal form
     Done(Expr<N, E>),
     DoneRef(&'a Expr<N, E>),
+    DoneSub(SubExpr<N, E>),
     DoneRefSub(&'a SubExpr<N, E>),
     // The current expression is already in normal form
     DoneAsIs,
 }
 
-fn normalize_ref(expr: &Expr<X, Normalized<'static>>) -> Expr<X, X> {
+fn normalize_ref<N>(expr: &Expr<N, Normalized<'static>>) -> Expr<N, X>
+where
+    N: fmt::Debug + Clone,
+{
     use dhall_core::BinOp::*;
     use dhall_core::ExprF::*;
     // Recursively normalize all subexpressions
-    let expr: ExprF<Expr<X, X>, Label, X, Normalized<'static>> =
+    let expr: ExprF<Expr<N, X>, Label, N, Normalized<'static>> =
         expr.map_ref_simple(|e| normalize_ref(e.as_ref()));
 
     use WhatNext::*;
@@ -241,7 +245,6 @@ fn normalize_ref(expr: &Expr<X, Normalized<'static>>) -> Expr<X, X> {
             ContinueSub(subst_shift(vf0, &r.roll(), &b.roll()))
         }
         Annot(x, _) => DoneRef(x),
-        Note(_, e) => DoneRef(e),
         App(f, args) if args.is_empty() => DoneRef(f),
         App(App(f, args1), args2) => Continue(App(
             f.clone(),
@@ -303,7 +306,7 @@ fn normalize_ref(expr: &Expr<X, Normalized<'static>>) -> Expr<X, X> {
                 .filter_map(|l| kvs.get(l).map(|x| (l.clone(), x.clone())))
                 .collect(),
         )),
-        Embed(e) => DoneRefSub(&e.0),
+        Embed(e) => DoneSub(e.0.note_absurd()),
         _ => DoneAsIs,
     };
 
@@ -312,12 +315,13 @@ fn normalize_ref(expr: &Expr<X, Normalized<'static>>) -> Expr<X, X> {
         ContinueSub(e) => normalize_ref(e.absurd().as_ref()),
         Done(e) => e,
         DoneRef(e) => e.clone(),
+        DoneSub(e) => e.unroll(),
         DoneRefSub(e) => e.unroll(),
         DoneAsIs => match expr.map_ref_simple(ExprF::roll) {
             e => e.map_ref(
                 SubExpr::clone,
                 |_, e| e.clone(),
-                X::clone,
+                N::clone,
                 |_| unreachable!(),
                 Label::clone,
             ),
@@ -334,7 +338,10 @@ fn normalize_ref(expr: &Expr<X, Normalized<'static>>) -> Expr<X, X> {
 /// However, `normalize` will not fail if the expression is ill-typed and will
 /// leave ill-typed sub-expressions unevaluated.
 ///
-fn normalize(e: SubExpr<X, Normalized<'static>>) -> SubExpr<X, X> {
+fn normalize<N>(e: SubExpr<N, Normalized<'static>>) -> SubExpr<N, X>
+where
+    N: fmt::Debug + Clone,
+{
     normalize_ref(e.as_ref()).roll()
 }