Add Expr visitor and improve tyexpr_to_expr

3 files changed, 52 insertions, 56 deletions

diff --git a/dhall/src/semantics/tck/tyexpr.rs b/dhall/src/semantics/tck/tyexpr.rs
index a8b8e58..f0fcdd1 100644
--- a/dhall/src/semantics/tck/tyexpr.rs
+++ b/dhall/src/semantics/tck/tyexpr.rs
@@ -36,15 +36,14 @@ impl TyExpr {
 
     /// Converts a value back to the corresponding AST expression.
     pub fn to_expr<'a>(&'a self, opts: ToExprOptions) -> NormalizedExpr {
-        tyexpr_to_expr(self, opts, &Vec::new())
+        tyexpr_to_expr(self, opts, &mut Vec::new())
     }
 }
 
-// TODO: mutate context once map_ref gets simplified
 fn tyexpr_to_expr<'a>(
     tyexpr: &'a TyExpr,
     opts: ToExprOptions,
-    ctx: &Vec<&'a Label>,
+    ctx: &mut Vec<&'a Label>,
 ) -> NormalizedExpr {
     rc(match tyexpr.kind() {
         TyExprKind::Var(v) if opts.alpha => {
@@ -61,13 +60,16 @@ fn tyexpr_to_expr<'a>(
             ExprKind::Var(V(name.clone(), idx))
         }
         TyExprKind::Expr(e) => {
-            let e = e.map_ref_with_special_handling_of_binders(
-                |tye| tyexpr_to_expr(tye, opts, ctx),
-                |l, tye| {
-                    let ctx = ctx.iter().copied().chain(Some(l)).collect();
-                    tyexpr_to_expr(tye, opts, &ctx)
-                },
-            );
+            let e = e.map_ref_maybe_binder(|l, tye| {
+                if let Some(l) = l {
+                    ctx.push(l);
+                }
+                let e = tyexpr_to_expr(tye, opts, ctx);
+                if let Some(_) = l {
+                    ctx.pop();
+                }
+                e
+            });
 
             match e {
                 ExprKind::Lam(_, t, e) if opts.alpha => {
diff --git a/dhall/src/syntax/ast/expr.rs b/dhall/src/syntax/ast/expr.rs
index f22a279..dbe233b 100644
--- a/dhall/src/syntax/ast/expr.rs
+++ b/dhall/src/syntax/ast/expr.rs
@@ -180,29 +180,38 @@ impl<Label> V<Label> {
 }
 
 impl<SE, E> ExprKind<SE, E> {
+    pub fn traverse_ref_maybe_binder<'a, SE2, Err>(
+        &'a self,
+        visit: impl FnMut(Option<&'a Label>, &'a SE) -> Result<SE2, Err>,
+    ) -> Result<ExprKind<SE2, E>, Err>
+    where
+        E: Clone,
+    {
+        visitor::TraverseRefMaybeBinderVisitor(visit).visit(self)
+    }
+
     pub fn traverse_ref_with_special_handling_of_binders<'a, SE2, Err>(
         &'a self,
-        visit_subexpr: impl FnMut(&'a SE) -> Result<SE2, Err>,
-        visit_under_binder: impl FnOnce(&'a Label, &'a SE) -> Result<SE2, Err>,
+        mut visit_subexpr: impl FnMut(&'a SE) -> Result<SE2, Err>,
+        mut visit_under_binder: impl FnMut(&'a Label, &'a SE) -> Result<SE2, Err>,
     ) -> Result<ExprKind<SE2, E>, Err>
     where
         E: Clone,
     {
-        visitor::TraverseRefWithBindersVisitor {
-            visit_subexpr,
-            visit_under_binder,
-        }
-        .visit(self)
+        self.traverse_ref_maybe_binder(|l, x| match l {
+            None => visit_subexpr(x),
+            Some(l) => visit_under_binder(l, x),
+        })
     }
 
     pub(crate) fn traverse_ref<'a, SE2, Err>(
         &'a self,
-        visit_subexpr: impl FnMut(&'a SE) -> Result<SE2, Err>,
+        mut visit_subexpr: impl FnMut(&'a SE) -> Result<SE2, Err>,
     ) -> Result<ExprKind<SE2, E>, Err>
     where
         E: Clone,
     {
-        visitor::TraverseRefVisitor { visit_subexpr }.visit(self)
+        self.traverse_ref_maybe_binder(|_, e| visit_subexpr(e))
     }
 
     fn traverse_mut<'a, Err>(
@@ -212,6 +221,16 @@ impl<SE, E> ExprKind<SE, E> {
         visitor::TraverseMutVisitor { visit_subexpr }.visit(self)
     }
 
+    pub fn map_ref_maybe_binder<'a, SE2>(
+        &'a self,
+        mut map: impl FnMut(Option<&'a Label>, &'a SE) -> SE2,
+    ) -> ExprKind<SE2, E>
+    where
+        E: Clone,
+    {
+        trivial_result(self.traverse_ref_maybe_binder(|l, x| Ok(map(l, x))))
+    }
+
     pub fn map_ref_with_special_handling_of_binders<'a, SE2>(
         &'a self,
         mut map_subexpr: impl FnMut(&'a SE) -> SE2,
@@ -220,10 +239,10 @@ impl<SE, E> ExprKind<SE, E> {
     where
         E: Clone,
     {
-        trivial_result(self.traverse_ref_with_special_handling_of_binders(
-            |x| Ok(map_subexpr(x)),
-            |l, x| Ok(map_under_binder(l, x)),
-        ))
+        self.map_ref_maybe_binder(|l, x| match l {
+            None => map_subexpr(x),
+            Some(l) => map_under_binder(l, x),
+        })
     }
 
     pub fn map_ref<'a, SE2>(
@@ -233,7 +252,7 @@ impl<SE, E> ExprKind<SE, E> {
     where
         E: Clone,
     {
-        trivial_result(self.traverse_ref(|x| Ok(map_subexpr(x))))
+        self.map_ref_maybe_binder(|_, e| map_subexpr(e))
     }
 
     pub fn map_mut<'a>(&'a mut self, mut map_subexpr: impl FnMut(&'a mut SE)) {
diff --git a/dhall/src/syntax/ast/visitor.rs b/dhall/src/syntax/ast/visitor.rs
index f9fd379..6a1ce7d 100644
--- a/dhall/src/syntax/ast/visitor.rs
+++ b/dhall/src/syntax/ast/visitor.rs
@@ -296,51 +296,26 @@ where
     Ok(())
 }
 
-pub struct TraverseRefWithBindersVisitor<F1, F2> {
-    pub visit_subexpr: F1,
-    pub visit_under_binder: F2,
-}
+pub struct TraverseRefMaybeBinderVisitor<F>(pub F);
 
-impl<'a, SE, E, SE2, Err, F1, F2> ExprKindVisitor<'a, SE, SE2, E, E>
-    for TraverseRefWithBindersVisitor<F1, F2>
+impl<'a, SE, E, SE2, Err, F> ExprKindVisitor<'a, SE, SE2, E, E>
+    for TraverseRefMaybeBinderVisitor<F>
 where
     SE: 'a,
     E: 'a + Clone,
-    F1: FnMut(&'a SE) -> Result<SE2, Err>,
-    F2: FnOnce(&'a Label, &'a SE) -> Result<SE2, Err>,
+    F: FnMut(Option<&'a Label>, &'a SE) -> Result<SE2, Err>,
 {
     type Error = Err;
 
     fn visit_subexpr(&mut self, subexpr: &'a SE) -> Result<SE2, Self::Error> {
-        (self.visit_subexpr)(subexpr)
+        (self.0)(None, subexpr)
     }
     fn visit_subexpr_under_binder(
-        self,
+        mut self,
         label: &'a Label,
         subexpr: &'a SE,
     ) -> Result<SE2, Self::Error> {
-        (self.visit_under_binder)(label, subexpr)
-    }
-    fn visit_embed(self, embed: &'a E) -> Result<E, Self::Error> {
-        Ok(embed.clone())
-    }
-}
-
-pub struct TraverseRefVisitor<F1> {
-    pub visit_subexpr: F1,
-}
-
-impl<'a, SE, E, SE2, Err, F1> ExprKindVisitor<'a, SE, SE2, E, E>
-    for TraverseRefVisitor<F1>
-where
-    SE: 'a,
-    E: 'a + Clone,
-    F1: FnMut(&'a SE) -> Result<SE2, Err>,
-{
-    type Error = Err;
-
-    fn visit_subexpr(&mut self, subexpr: &'a SE) -> Result<SE2, Self::Error> {
-        (self.visit_subexpr)(subexpr)
+        (self.0)(Some(label), subexpr)
     }
     fn visit_embed(self, embed: &'a E) -> Result<E, Self::Error> {
         Ok(embed.clone())