3 files changed, 374 insertions, 274 deletions

diff --git a/dhall_core/src/core.rs b/dhall_core/src/core.rs
index bfe769f..aeb6f23 100644
--- a/dhall_core/src/core.rs
+++ b/dhall_core/src/core.rs
@@ -205,135 +205,66 @@ pub enum ExprF<SubExpr, Label, Note, Embed> {
 }
 
 impl<SE, L, N, E> ExprF<SE, L, N, E> {
-    pub fn visit<'a, V, Return>(&'a self, v: V) -> Return
+    pub(crate) fn visit<'a, V, Return>(&'a self, v: V) -> Return
     where
         V: visitor::GenericVisitor<&'a ExprF<SE, L, N, E>, Return>,
     {
         v.visit(self)
     }
-}
-
-impl<S, A> Expr<S, A> {
-    pub fn map_shallow<T, B, F1, F2, F3, F4>(
-        &self,
-        map_expr: F1,
-        map_note: F2,
-        map_embed: F3,
-        map_label: F4,
-    ) -> Expr<T, B>
-    where
-        A: Clone,
-        T: Clone,
-        S: Clone,
-        F1: Fn(&Self) -> Expr<T, B>,
-        F2: Fn(&S) -> T,
-        F3: Fn(&A) -> B,
-        F4: Fn(&Label) -> Label,
-    {
-        self.map_ref(
-            |x| rc(map_expr(x.as_ref())),
-            |_, x| rc(map_expr(x.as_ref())),
-            map_note,
-            map_embed,
-            map_label,
-        )
-    }
-
-    pub fn map_embed<B, F>(&self, map_embed: &F) -> Expr<S, B>
-    where
-        A: Clone,
-        S: Clone,
-        F: Fn(&A) -> B,
-    {
-        let recurse = |e: &Expr<S, A>| -> Expr<S, B> { e.map_embed(map_embed) };
-        self.map_shallow(recurse, S::clone, map_embed, Label::clone)
-    }
-
-    pub fn traverse_embed<B, Err, F>(
-        &self,
-        map_embed: F,
-    ) -> Result<Expr<S, B>, Err>
-    where
-        S: Clone,
-        B: Clone,
-        F: FnMut(&A) -> Result<B, Err>,
-    {
-        self.visit(&mut visitor::TraverseEmbedVisitor(map_embed))
-    }
-
-    pub fn map_label<F>(&self, map_label: &F) -> Self
-    where
-        A: Clone,
-        S: Clone,
-        F: Fn(&Label) -> Label,
-    {
-        let recurse = |e: &Self| -> Self { e.map_label(map_label) };
-        self.map_shallow(recurse, S::clone, A::clone, map_label)
-    }
 
-    pub fn roll(&self) -> SubExpr<S, A>
+    fn traverse_ref_with_special_handling_of_binders<'a, SE2, L2, N2, E2, Err>(
+        &'a self,
+        visit_subexpr: impl FnMut(&'a SE) -> Result<SE2, Err>,
+        visit_under_binder: impl FnOnce(&'a L, &'a SE) -> Result<SE2, Err>,
+        visit_note: impl FnOnce(&'a N) -> Result<N2, Err>,
+        visit_embed: impl FnOnce(&'a E) -> Result<E2, Err>,
+        visit_label: impl FnMut(&'a L) -> Result<L2, Err>,
+    ) -> Result<ExprF<SE2, L2, N2, E2>, Err>
     where
-        S: Clone,
-        A: Clone,
+        L: Ord,
+        L2: Ord,
     {
-        rc(ExprF::clone(self))
-    }
-}
-
-impl<N: Clone, E> Expr<N, E> {
-    pub fn squash_embed<E2: Clone>(
-        &self,
-        f: impl FnMut(&E) -> SubExpr<N, E2>,
-    ) -> SubExpr<N, E2> {
-        rc(self.visit(&mut visitor::SquashEmbedVisitor(f)))
+        self.visit(visitor::TraverseRefWithBindersVisitor {
+            visit_subexpr,
+            visit_under_binder,
+            visit_note,
+            visit_embed,
+            visit_label,
+        })
     }
-}
 
-impl<SE, L, N, E> ExprF<SE, L, N, E> {
-    pub fn traverse_ref<'a, SE2, L2, N2, E2, Err, F1, F2, F3, F4, F5>(
+    fn traverse_ref<'a, SE2, L2, N2, E2, Err>(
         &'a self,
-        visit_subexpr: F1,
-        visit_under_binder: F2,
-        visit_note: F3,
-        visit_embed: F4,
-        visit_label: F5,
+        visit_subexpr: impl FnMut(&'a SE) -> Result<SE2, Err>,
+        visit_note: impl FnOnce(&'a N) -> Result<N2, Err>,
+        visit_embed: impl FnOnce(&'a E) -> Result<E2, Err>,
+        visit_label: impl FnMut(&'a L) -> Result<L2, Err>,
     ) -> Result<ExprF<SE2, L2, N2, E2>, Err>
     where
         L: Ord,
         L2: Ord,
-        F1: FnMut(&'a SE) -> Result<SE2, Err>,
-        F2: FnOnce(&'a L, &'a SE) -> Result<SE2, Err>,
-        F3: FnOnce(&'a N) -> Result<N2, Err>,
-        F4: FnOnce(&'a E) -> Result<E2, Err>,
-        F5: FnMut(&'a L) -> Result<L2, Err>,
     {
         self.visit(visitor::TraverseRefVisitor {
             visit_subexpr,
-            visit_under_binder,
             visit_note,
             visit_embed,
             visit_label,
         })
     }
 
-    pub fn map_ref<'a, SE2, L2, N2, E2, F1, F2, F3, F4, F5>(
+    pub fn map_ref_with_special_handling_of_binders<'a, SE2, L2, N2, E2>(
         &'a self,
-        mut map_subexpr: F1,
-        mut map_under_binder: F2,
-        mut map_note: F3,
-        mut map_embed: F4,
-        mut map_label: F5,
+        mut map_subexpr: impl FnMut(&'a SE) -> SE2,
+        mut map_under_binder: impl FnMut(&'a L, &'a SE) -> SE2,
+        map_note: impl FnOnce(&'a N) -> N2,
+        map_embed: impl FnOnce(&'a E) -> E2,
+        mut map_label: impl FnMut(&'a L) -> L2,
     ) -> ExprF<SE2, L2, N2, E2>
     where
         L: Ord,
         L2: Ord,
-        F1: FnMut(&'a SE) -> SE2,
-        F2: FnMut(&'a L, &'a SE) -> SE2,
-        F3: FnMut(&'a N) -> N2,
-        F4: FnMut(&'a E) -> E2,
-        F5: FnMut(&'a L) -> L2,
     {
-        trivial_result(self.traverse_ref(
+        trivial_result(self.traverse_ref_with_special_handling_of_binders(
             |x| Ok(map_subexpr(x)),
             |l, x| Ok(map_under_binder(l, x)),
             |x| Ok(map_note(x)),
@@ -342,31 +273,101 @@ impl<SE, L, N, E> ExprF<SE, L, N, E> {
         ))
     }
 
-    pub fn traverse_ref_simple<'a, SE2, Err, F1>(
+    pub fn map_ref<'a, SE2, L2, N2, E2>(
         &'a self,
-        visit_subexpr: F1,
+        mut map_subexpr: impl FnMut(&'a SE) -> SE2,
+        map_note: impl FnOnce(&'a N) -> N2,
+        map_embed: impl FnOnce(&'a E) -> E2,
+        mut map_label: impl FnMut(&'a L) -> L2,
+    ) -> ExprF<SE2, L2, N2, E2>
+    where
+        L: Ord,
+        L2: Ord,
+    {
+        trivial_result(self.traverse_ref(
+            |x| Ok(map_subexpr(x)),
+            |x| Ok(map_note(x)),
+            |x| Ok(map_embed(x)),
+            |x| Ok(map_label(x)),
+        ))
+    }
+
+    pub fn traverse_ref_simple<'a, SE2, Err>(
+        &'a self,
+        visit_subexpr: impl FnMut(&'a SE) -> Result<SE2, Err>,
     ) -> Result<ExprF<SE2, L, N, E>, Err>
     where
         L: Ord + Clone,
         N: Clone,
         E: Clone,
-        F1: FnMut(&'a SE) -> Result<SE2, Err>,
     {
-        self.visit(visitor::TraverseRefSimpleVisitor { visit_subexpr })
+        self.traverse_ref(
+            visit_subexpr,
+            |x| Ok(N::clone(x)),
+            |x| Ok(E::clone(x)),
+            |x| Ok(L::clone(x)),
+        )
     }
 
-    pub fn map_ref_simple<'a, SE2, F1>(
+    pub fn map_ref_simple<'a, SE2>(
         &'a self,
-        map_subexpr: F1,
+        map_subexpr: impl Fn(&'a SE) -> SE2,
     ) -> ExprF<SE2, L, N, E>
     where
-        L: Ord,
-        L: Clone,
+        L: Ord + Clone,
         N: Clone,
         E: Clone,
-        F1: Fn(&'a SE) -> SE2,
     {
-        trivial_result(self.traverse_ref_simple(|x| Ok(map_subexpr(x))))
+        self.map_ref(map_subexpr, N::clone, E::clone, L::clone)
+    }
+}
+
+impl<N, E> Expr<N, E> {
+    fn traverse_embed<E2, Err>(
+        &self,
+        visit_embed: impl FnMut(&E) -> Result<E2, Err>,
+    ) -> Result<Expr<N, E2>, Err>
+    where
+        N: Clone,
+    {
+        self.visit(&mut visitor::TraverseEmbedVisitor(visit_embed))
+    }
+
+    fn map_embed<E2>(&self, mut map_embed: impl FnMut(&E) -> E2) -> Expr<N, E2>
+    where
+        N: Clone,
+    {
+        trivial_result(self.traverse_embed(|x| Ok(map_embed(x))))
+    }
+
+    pub fn roll(&self) -> SubExpr<N, E>
+    where
+        N: Clone,
+        E: Clone,
+    {
+        rc(ExprF::clone(self))
+    }
+
+    pub fn squash_embed<E2>(
+        &self,
+        f: impl FnMut(&E) -> SubExpr<N, E2>,
+    ) -> SubExpr<N, E2>
+    where
+        N: Clone,
+    {
+        trivial_result(self.visit(&mut visitor::SquashEmbedVisitor(f)))
+    }
+}
+
+impl<E: Clone> Expr<X, E> {
+    pub fn note_absurd<N>(&self) -> Expr<N, E> {
+        self.visit(&mut visitor::NoteAbsurdVisitor)
+    }
+}
+
+impl<N: Clone> Expr<N, X> {
+    pub fn embed_absurd<E>(&self) -> Expr<N, E> {
+        self.visit(&mut visitor::EmbedAbsurdVisitor)
     }
 }
 
@@ -375,17 +376,42 @@ impl<N, E> SubExpr<N, E> {
         self.0.as_ref()
     }
 
-    fn map_ref<'a, F1, F2>(&'a self, map_expr: F1, map_under_binder: F2) -> Self
+    pub fn traverse_embed<E2, Err>(
+        &self,
+        visit_embed: impl FnMut(&E) -> Result<E2, Err>,
+    ) -> Result<SubExpr<N, E2>, Err>
     where
-        F1: FnMut(&'a Self) -> Self,
-        F2: FnMut(&'a Label, &'a Self) -> Self,
+        N: Clone,
+    {
+        Ok(rc(self.as_ref().traverse_embed(visit_embed)?))
+    }
+
+    pub fn map_embed<E2>(
+        &self,
+        map_embed: impl FnMut(&E) -> E2,
+    ) -> SubExpr<N, E2>
+    where
+        N: Clone,
     {
+        rc(self.as_ref().map_embed(map_embed))
+    }
+
+    pub fn map_subexprs_with_special_handling_of_binders<'a>(
+        &'a self,
+        map_expr: impl FnMut(&'a Self) -> Self,
+        map_under_binder: impl FnMut(&'a Label, &'a Self) -> Self,
+    ) -> Self {
         match self.as_ref() {
             ExprF::Embed(_) => SubExpr::clone(self),
             // Recursive call
-            ExprF::Note(_, e) => e.map_ref(map_expr, map_under_binder),
+            // TODO: don't discard the note !
+            ExprF::Note(_, e) => e
+                .map_subexprs_with_special_handling_of_binders(
+                    map_expr,
+                    map_under_binder,
+                ),
             // Call ExprF::map_ref
-            e => rc(e.map_ref(
+            e => rc(e.map_ref_with_special_handling_of_binders(
                 map_expr,
                 map_under_binder,
                 |_| unreachable!(),
@@ -395,25 +421,25 @@ impl<N, E> SubExpr<N, E> {
         }
     }
 
-    pub fn map_ref_simple<F1>(&self, map_expr: F1) -> Self
+    pub fn unroll(&self) -> Expr<N, E>
     where
-        F1: Fn(&Self) -> Self,
+        N: Clone,
+        E: Clone,
     {
-        self.map_ref(&map_expr, |_, e| map_expr(e))
+        ExprF::clone(self.as_ref())
     }
 
-    pub fn unroll(&self) -> Expr<N, E>
+    pub fn unnote(&self) -> SubExpr<X, E>
     where
-        N: Clone,
         E: Clone,
     {
-        ExprF::clone(self.as_ref())
+        rc(self.as_ref().visit(&mut visitor::UnNoteVisitor))
     }
 }
 
 impl<N: Clone> SubExpr<N, X> {
-    pub fn absurd<T>(&self) -> SubExpr<N, T> {
-        rc(self.as_ref().absurd_rec())
+    pub fn embed_absurd<T>(&self) -> SubExpr<N, T> {
+        rc(self.as_ref().embed_absurd())
     }
 }
 
@@ -423,34 +449,6 @@ impl<E: Clone> SubExpr<X, E> {
     }
 }
 
-impl<E: Clone> Expr<X, E> {
-    pub fn note_absurd<N>(&self) -> Expr<N, E> {
-        self.visit(&mut visitor::NoteAbsurdVisitor)
-    }
-}
-
-impl<N, E: Clone> SubExpr<N, E> {
-    pub fn unnote(&self) -> SubExpr<X, E> {
-        rc(self.as_ref().visit(&mut visitor::UnNoteVisitor))
-    }
-}
-
-impl<N: Clone> Expr<N, X> {
-    // Deprecated, use embed_absurd instead
-    pub fn absurd_rec<T>(&self) -> Expr<N, T> {
-        self.embed_absurd()
-    }
-    pub fn embed_absurd<T>(&self) -> Expr<N, T> {
-        self.visit(&mut visitor::EmbedAbsurdVisitor)
-    }
-}
-
-impl<N: Clone> SubExpr<N, X> {
-    pub fn embed_absurd<T>(&self) -> SubExpr<N, T> {
-        rc(self.as_ref().embed_absurd())
-    }
-}
-
 impl<N, E> Clone for SubExpr<N, E> {
     fn clone(&self) -> Self {
         SubExpr(Rc::clone(&self.0))
@@ -462,25 +460,8 @@ pub fn rc<N, E>(x: Expr<N, E>) -> SubExpr<N, E> {
     SubExpr(Rc::new(x))
 }
 
-pub fn app<N, E>(f: Expr<N, E>, args: Vec<SubExpr<N, E>>) -> Expr<N, E> {
-    if args.is_empty() {
-        f
-    } else {
-        ExprF::App(rc(f), args)
-    }
-}
-
-pub fn app_rc<N, E>(
-    f: SubExpr<N, E>,
-    args: Vec<SubExpr<N, E>>,
-) -> SubExpr<N, E> {
-    if args.is_empty() {
-        f
-    } else {
-        rc(ExprF::App(f, args))
-    }
-}
-
+/// Add an isize to an usize
+/// Panics on over/underflow
 fn add_ui(u: usize, i: isize) -> usize {
     if i < 0 {
         u.checked_sub(i.checked_neg().unwrap() as usize).unwrap()
@@ -505,15 +486,15 @@ fn shift_var(delta: isize, var: &V<Label>, in_expr: &V<Label>) -> V<Label> {
 /// capture by shifting variable indices
 /// See https://github.com/dhall-lang/dhall-lang/blob/master/standard/semantics.md#shift
 /// for details
-pub fn shift<S, A>(
+pub fn shift<N, E>(
     delta: isize,
     var: &V<Label>,
-    in_expr: &SubExpr<S, A>,
-) -> SubExpr<S, A> {
+    in_expr: &SubExpr<N, E>,
+) -> SubExpr<N, E> {
     use crate::ExprF::*;
     match in_expr.as_ref() {
         Var(v) => rc(Var(shift_var(delta, var, v))),
-        _ => in_expr.map_ref(
+        _ => in_expr.map_subexprs_with_special_handling_of_binders(
             |e| shift(delta, var, e),
             |l: &Label, e| {
                 let vl = V(l.clone(), 0);
@@ -532,16 +513,16 @@ pub fn shift<S, A>(
 /// subst_shift(x, v, e) = ↑(-1, x, e[x := ↑(1, x, v)])
 /// ```
 ///
-pub fn subst_shift<S, A>(
+pub fn subst_shift<N, E>(
     var: &V<Label>,
-    value: &SubExpr<S, A>,
-    in_expr: &SubExpr<S, A>,
-) -> SubExpr<S, A> {
+    value: &SubExpr<N, E>,
+    in_expr: &SubExpr<N, E>,
+) -> SubExpr<N, E> {
     use crate::ExprF::*;
     match in_expr.as_ref() {
         Var(v) if v == var => SubExpr::clone(value),
         Var(v) => rc(Var(shift_var(-1, var, v))),
-        _ => in_expr.map_ref(
+        _ => in_expr.map_subexprs_with_special_handling_of_binders(
             |e| subst_shift(var, &value, e),
             |l: &Label, e| {
                 let vl = V(l.clone(), 0);
diff --git a/dhall_core/src/parser.rs b/dhall_core/src/parser.rs
index 1e47c16..2a30b2b 100644
--- a/dhall_core/src/parser.rs
+++ b/dhall_core/src/parser.rs
@@ -746,7 +746,7 @@ make_parser! {
     rule!(application_expression<ParsedExpr<'a>> as expression; span; children!(
         [expression(e)] => e,
         [expression(first), expression(rest)..] => {
-            spanned(span, app(first, rest.map(rc).collect()))
+            spanned(span, App(rc(first), rest.map(rc).collect()))
         },
     ));
 
diff --git a/dhall_core/src/visitor.rs b/dhall_core/src/visitor.rs
index c123275..16ad418 100644
--- a/dhall_core/src/visitor.rs
+++ b/dhall_core/src/visitor.rs
@@ -17,67 +17,49 @@ pub trait GenericVisitor<Input, Return>: Sized {
 /// time: the visitor itself. The visitor can then carry around multiple closures
 /// or just one, and Rust is ok with either. See for example TraverseRefVisitor
 /// and TraverseEmbedVisitor.
-pub trait ExprFFallibleVisitor<'a, SE1, SE2, L1, L2, N1, N2, E1, E2>:
-    Sized
-{
+/// This is very generic. For a more legible trait, see ExprFInFallibleVisitor
+pub trait ExprFVeryGenericVisitor<'a, Ret, SE1, L1, N1, E1>: Sized {
     type Error;
+    type SE2;
+    type L2;
+    type N2;
+    type E2;
 
-    fn visit_subexpr(&mut self, subexpr: &'a SE1) -> Result<SE2, Self::Error>;
-    fn visit_label(&mut self, label: &'a L1) -> Result<L2, Self::Error>;
-    fn visit_note(self, note: &'a N1) -> Result<N2, Self::Error>;
-    fn visit_embed(self, embed: &'a E1) -> Result<E2, Self::Error>;
-
-    fn visit_subexpr_under_binder(
-        mut self,
-        _label: &'a L1,
+    fn visit_subexpr(
+        &mut self,
         subexpr: &'a SE1,
-    ) -> Result<SE2, Self::Error> {
-        self.visit_subexpr(subexpr)
-    }
+    ) -> Result<Self::SE2, Self::Error>;
+    fn visit_label(&mut self, label: &'a L1) -> Result<Self::L2, Self::Error>;
 
     fn visit_binder(
-        mut self,
+        self,
         label: &'a L1,
         subexpr: &'a SE1,
-    ) -> Result<(L2, SE2), Self::Error> {
-        Ok((
-            self.visit_label(label)?,
-            self.visit_subexpr_under_binder(label, subexpr)?,
-        ))
-    }
+    ) -> Result<(Self::L2, Self::SE2), Self::Error>;
 
-    fn visit_embed_squash(
-        self,
-        embed: &'a E1,
-    ) -> Result<ExprF<SE2, L2, N2, E2>, Self::Error> {
-        Ok(ExprF::Embed(self.visit_embed(embed)?))
-    }
+    fn visit_embed_squash(self, embed: &'a E1) -> Result<Ret, Self::Error>;
 
     fn visit_note_squash(
-        mut self,
+        self,
         note: &'a N1,
         subexpr: &'a SE1,
-    ) -> Result<ExprF<SE2, L2, N2, E2>, Self::Error> {
-        let subexpr = self.visit_subexpr(subexpr)?;
-        let note = self.visit_note(note)?;
-        Ok(ExprF::Note(note, subexpr))
-    }
+    ) -> Result<Ret, Self::Error>;
+
+    // Called with the result of the map, in the non-embed/note case.
+    // Useful to change the result type, and/or avoid some loss of info
+    fn visit_resulting_exprf(
+        result: ExprF<Self::SE2, Self::L2, Self::N2, Self::E2>,
+    ) -> Result<Ret, Self::Error>;
 }
 
-impl<'a, T, SE1, SE2, L1, L2, N1, N2, E1, E2>
-    GenericVisitor<
-        &'a ExprF<SE1, L1, N1, E1>,
-        Result<ExprF<SE2, L2, N2, E2>, T::Error>,
-    > for T
+impl<'a, T, Ret, SE1, L1, N1, E1>
+    GenericVisitor<&'a ExprF<SE1, L1, N1, E1>, Result<Ret, T::Error>> for T
 where
     L1: Ord,
-    L2: Ord,
-    T: ExprFFallibleVisitor<'a, SE1, SE2, L1, L2, N1, N2, E1, E2>,
+    T::L2: Ord,
+    T: ExprFVeryGenericVisitor<'a, Ret, SE1, L1, N1, E1>,
 {
-    fn visit(
-        self,
-        input: &'a ExprF<SE1, L1, N1, E1>,
-    ) -> Result<ExprF<SE2, L2, N2, E2>, T::Error> {
+    fn visit(self, input: &'a ExprF<SE1, L1, N1, E1>) -> Result<Ret, T::Error> {
         fn vec<'a, T, U, Err, F: FnMut(&'a T) -> Result<U, Err>>(
             x: &'a [T],
             f: F,
@@ -93,27 +75,27 @@ where
                 None => None,
             })
         }
-        fn btmap<'a, V, SE, L, N, E, SE2, L2, N2, E2>(
+        fn btmap<'a, V, Ret, SE, L, N, E>(
             x: &'a BTreeMap<L, SE>,
             mut v: V,
-        ) -> Result<BTreeMap<L2, SE2>, V::Error>
+        ) -> Result<BTreeMap<V::L2, V::SE2>, V::Error>
         where
             L: Ord,
-            L2: Ord,
-            V: ExprFFallibleVisitor<'a, SE, SE2, L, L2, N, N2, E, E2>,
+            V::L2: Ord,
+            V: ExprFVeryGenericVisitor<'a, Ret, SE, L, N, E>,
         {
             x.iter()
                 .map(|(k, x)| Ok((v.visit_label(k)?, v.visit_subexpr(x)?)))
                 .collect()
         }
-        fn btoptmap<'a, V, SE, L, N, E, SE2, L2, N2, E2>(
+        fn btoptmap<'a, V, Ret, SE, L, N, E>(
             x: &'a BTreeMap<L, Option<SE>>,
             mut v: V,
-        ) -> Result<BTreeMap<L2, Option<SE2>>, V::Error>
+        ) -> Result<BTreeMap<V::L2, Option<V::SE2>>, V::Error>
         where
             L: Ord,
-            L2: Ord,
-            V: ExprFFallibleVisitor<'a, SE, SE2, L, L2, N, N2, E, E2>,
+            V::L2: Ord,
+            V: ExprFVeryGenericVisitor<'a, Ret, SE, L, N, E>,
         {
             x.iter()
                 .map(|(k, x)| {
@@ -130,7 +112,7 @@ where
 
         let mut v = self;
         use crate::ExprF::*;
-        Ok(match input {
+        T::visit_resulting_exprf(match input {
             Var(V(l, n)) => Var(V(v.visit_label(l)?, *n)),
             Lam(l, t, e) => {
                 let t = v.visit_subexpr(t)?;
@@ -195,12 +177,115 @@ where
             Projection(e, ls) => {
                 Projection(v.visit_subexpr(e)?, vec(ls, |l| v.visit_label(l))?)
             }
-            Note(n, e) => v.visit_note_squash(n, e)?,
-            Embed(a) => v.visit_embed_squash(a)?,
+            Note(n, e) => return v.visit_note_squash(n, e),
+            Embed(a) => return v.visit_embed_squash(a),
         })
     }
 }
 
+/// Like ExprFVeryGenericVisitor, but sets the return
+/// type to ExprF<_>
+pub trait ExprFFallibleVisitor<'a, SE1, SE2, L1, L2, N1, N2, E1, E2>:
+    Sized
+{
+    type Error;
+
+    fn visit_subexpr(&mut self, subexpr: &'a SE1) -> Result<SE2, Self::Error>;
+    fn visit_label(&mut self, label: &'a L1) -> Result<L2, Self::Error>;
+    fn visit_note(self, note: &'a N1) -> Result<N2, Self::Error>;
+    fn visit_embed(self, embed: &'a E1) -> Result<E2, Self::Error>;
+
+    fn visit_subexpr_under_binder(
+        mut self,
+        _label: &'a L1,
+        subexpr: &'a SE1,
+    ) -> Result<SE2, Self::Error> {
+        self.visit_subexpr(subexpr)
+    }
+
+    fn visit_binder(
+        mut self,
+        label: &'a L1,
+        subexpr: &'a SE1,
+    ) -> Result<(L2, SE2), Self::Error> {
+        Ok((
+            self.visit_label(label)?,
+            self.visit_subexpr_under_binder(label, subexpr)?,
+        ))
+    }
+
+    fn visit_embed_squash(
+        self,
+        embed: &'a E1,
+    ) -> Result<ExprF<SE2, L2, N2, E2>, Self::Error> {
+        Ok(ExprF::Embed(self.visit_embed(embed)?))
+    }
+
+    fn visit_note_squash(
+        mut self,
+        note: &'a N1,
+        subexpr: &'a SE1,
+    ) -> Result<ExprF<SE2, L2, N2, E2>, Self::Error> {
+        let subexpr = self.visit_subexpr(subexpr)?;
+        let note = self.visit_note(note)?;
+        Ok(ExprF::Note(note, subexpr))
+    }
+}
+
+impl<'a, T, SE1, SE2, L1, L2, N1, N2, E1, E2>
+    ExprFVeryGenericVisitor<'a, ExprF<SE2, L2, N2, E2>, SE1, L1, N1, E1> for T
+where
+    T: ExprFFallibleVisitor<'a, SE1, SE2, L1, L2, N1, N2, E1, E2>,
+{
+    type Error = T::Error;
+    type SE2 = SE2;
+    type L2 = L2;
+    type N2 = N2;
+    type E2 = E2;
+
+    fn visit_subexpr(
+        &mut self,
+        subexpr: &'a SE1,
+    ) -> Result<Self::SE2, Self::Error> {
+        self.visit_subexpr(subexpr)
+    }
+
+    fn visit_label(&mut self, label: &'a L1) -> Result<Self::L2, Self::Error> {
+        self.visit_label(label)
+    }
+
+    fn visit_binder(
+        self,
+        label: &'a L1,
+        subexpr: &'a SE1,
+    ) -> Result<(Self::L2, Self::SE2), Self::Error> {
+        self.visit_binder(label, subexpr)
+    }
+
+    fn visit_embed_squash(
+        self,
+        embed: &'a E1,
+    ) -> Result<ExprF<SE2, L2, N2, E2>, Self::Error> {
+        self.visit_embed_squash(embed)
+    }
+
+    fn visit_note_squash(
+        self,
+        note: &'a N1,
+        subexpr: &'a SE1,
+    ) -> Result<ExprF<SE2, L2, N2, E2>, Self::Error> {
+        self.visit_note_squash(note, subexpr)
+    }
+
+    // Called with the result of the map, in the non-embed/note case.
+    // Useful to change the result type, and/or avoid some loss of info
+    fn visit_resulting_exprf(
+        result: ExprF<Self::SE2, Self::L2, Self::N2, Self::E2>,
+    ) -> Result<ExprF<SE2, L2, N2, E2>, Self::Error> {
+        Ok(result)
+    }
+}
+
 /// Like ExprFFallibleVisitor, but without the error handling.
 pub trait ExprFInFallibleVisitor<'a, SE1, SE2, L1, L2, N1, N2, E1, E2>:
     Sized
@@ -302,7 +387,7 @@ where
     }
 }
 
-pub struct TraverseRefVisitor<F1, F2, F3, F4, F5> {
+pub struct TraverseRefWithBindersVisitor<F1, F2, F3, F4, F5> {
     pub visit_subexpr: F1,
     pub visit_under_binder: F2,
     pub visit_note: F3,
@@ -312,7 +397,7 @@ pub struct TraverseRefVisitor<F1, F2, F3, F4, F5> {
 
 impl<'a, SE, L, N, E, SE2, L2, N2, E2, Err, F1, F2, F3, F4, F5>
     ExprFFallibleVisitor<'a, SE, SE2, L, L2, N, N2, E, E2>
-    for TraverseRefVisitor<F1, F2, F3, F4, F5>
+    for TraverseRefWithBindersVisitor<F1, F2, F3, F4, F5>
 where
     SE: 'a,
     L: 'a,
@@ -349,33 +434,41 @@ where
     }
 }
 
-pub struct TraverseRefSimpleVisitor<F1> {
+pub struct TraverseRefVisitor<F1, F2, F3, F4> {
     pub visit_subexpr: F1,
+    pub visit_note: F2,
+    pub visit_embed: F3,
+    pub visit_label: F4,
 }
 
-impl<'a, SE, L, N, E, SE2, Err, F1>
-    ExprFFallibleVisitor<'a, SE, SE2, L, L, N, N, E, E>
-    for TraverseRefSimpleVisitor<F1>
+impl<'a, SE, L, N, E, SE2, L2, N2, E2, Err, F1, F2, F3, F4>
+    ExprFFallibleVisitor<'a, SE, SE2, L, L2, N, N2, E, E2>
+    for TraverseRefVisitor<F1, F2, F3, F4>
 where
     SE: 'a,
-    L: Ord + Clone + 'a,
-    N: Clone + 'a,
-    E: Clone + 'a,
+    L: 'a,
+    N: 'a,
+    E: 'a,
+    L: Ord,
+    L2: Ord,
     F1: FnMut(&'a SE) -> Result<SE2, Err>,
+    F2: FnOnce(&'a N) -> Result<N2, Err>,
+    F3: FnOnce(&'a E) -> Result<E2, Err>,
+    F4: FnMut(&'a L) -> Result<L2, Err>,
 {
     type Error = Err;
 
     fn visit_subexpr(&mut self, subexpr: &'a SE) -> Result<SE2, Self::Error> {
         (self.visit_subexpr)(subexpr)
     }
-    fn visit_note(self, note: &'a N) -> Result<N, Self::Error> {
-        Ok(N::clone(note))
+    fn visit_note(self, note: &'a N) -> Result<N2, Self::Error> {
+        (self.visit_note)(note)
     }
-    fn visit_embed(self, embed: &'a E) -> Result<E, Self::Error> {
-        Ok(E::clone(embed))
+    fn visit_embed(self, embed: &'a E) -> Result<E2, Self::Error> {
+        (self.visit_embed)(embed)
     }
-    fn visit_label(&mut self, label: &'a L) -> Result<L, Self::Error> {
-        Ok(L::clone(label))
+    fn visit_label(&mut self, label: &'a L) -> Result<L2, Self::Error> {
+        (self.visit_label)(label)
     }
 }
 
@@ -395,7 +488,6 @@ impl<'a, 'b, N, E, E2, Err, F1>
     > for &'b mut TraverseEmbedVisitor<F1>
 where
     N: Clone + 'a,
-    E2: Clone,
     F1: FnMut(&E) -> Result<E2, Err>,
 {
     type Error = Err;
@@ -419,37 +511,64 @@ where
 
 pub struct SquashEmbedVisitor<F1>(pub F1);
 
-impl<'a, 'b, N, E, E2, F1>
-    ExprFInFallibleVisitor<
-        'a,
-        SubExpr<N, E>,
-        SubExpr<N, E2>,
-        Label,
-        Label,
-        N,
-        N,
-        E,
-        E2,
-    > for &'b mut SquashEmbedVisitor<F1>
+impl<'a, 'b, N, E1, E2, F1>
+    ExprFVeryGenericVisitor<'a, SubExpr<N, E2>, SubExpr<N, E1>, Label, N, E1>
+    for &'b mut SquashEmbedVisitor<F1>
 where
     N: Clone + 'a,
-    E2: Clone,
-    F1: FnMut(&E) -> SubExpr<N, E2>,
+    F1: FnMut(&E1) -> SubExpr<N, E2>,
 {
-    fn visit_subexpr(&mut self, subexpr: &'a SubExpr<N, E>) -> SubExpr<N, E2> {
-        rc(subexpr.as_ref().visit(&mut **self))
+    type Error = X;
+    type SE2 = SubExpr<N, E2>;
+    type L2 = Label;
+    type N2 = N;
+    type E2 = E2;
+
+    fn visit_subexpr(
+        &mut self,
+        subexpr: &'a SubExpr<N, E1>,
+    ) -> Result<Self::SE2, Self::Error> {
+        Ok(subexpr.as_ref().visit(&mut **self)?)
     }
-    fn visit_note(self, note: &'a N) -> N {
-        N::clone(note)
+
+    fn visit_label(
+        &mut self,
+        label: &'a Label,
+    ) -> Result<Self::L2, Self::Error> {
+        Ok(Label::clone(label))
     }
-    fn visit_embed(self, _: &'a E) -> E2 {
-        unreachable!()
+
+    fn visit_binder(
+        mut self,
+        label: &'a Label,
+        subexpr: &'a SubExpr<N, E1>,
+    ) -> Result<(Self::L2, Self::SE2), Self::Error> {
+        Ok((self.visit_label(label)?, self.visit_subexpr(subexpr)?))
+    }
+
+    fn visit_embed_squash(
+        self,
+        embed: &'a E1,
+    ) -> Result<SubExpr<N, E2>, Self::Error> {
+        Ok((self.0)(embed))
     }
-    fn visit_embed_squash(self, embed: &'a E) -> Expr<N, E2> {
-        (self.0)(embed).unroll()
+
+    fn visit_note_squash(
+        mut self,
+        note: &'a N,
+        subexpr: &'a SubExpr<N, E1>,
+    ) -> Result<SubExpr<N, E2>, Self::Error> {
+        let subexpr = self.visit_subexpr(subexpr)?;
+        let note = N::clone(note);
+        Ok(rc(ExprF::Note(note, subexpr)))
     }
-    fn visit_label(&mut self, label: &'a Label) -> Label {
-        Label::clone(label)
+
+    // Called with the result of the map, in the non-embed/note case.
+    // Useful to change the result type, and/or avoid some loss of info
+    fn visit_resulting_exprf(
+        result: ExprF<Self::SE2, Self::L2, Self::N2, Self::E2>,
+    ) -> Result<SubExpr<N, E2>, Self::Error> {
+        Ok(rc(result))
     }
 }