Transform more mapping functions to use the Visitor trait

2 files changed, 173 insertions, 73 deletions

diff --git a/dhall_core/src/core.rs b/dhall_core/src/core.rs
index fe53bb2..2c23ed1 100644
--- a/dhall_core/src/core.rs
+++ b/dhall_core/src/core.rs
@@ -2,7 +2,7 @@
 use std::collections::BTreeMap;
 use std::rc::Rc;
 
-use crate::visitor::{TraverseRefSimpleVisitor, TraverseRefVisitor};
+use crate::visitor;
 use crate::*;
 
 pub type Integer = isize;
@@ -236,23 +236,14 @@ impl<S, A> Expr<S, A> {
 
     pub fn traverse_embed<B, Err, F>(
         &self,
-        map_embed: &F,
+        map_embed: F,
     ) -> Result<Expr<S, B>, Err>
     where
         S: Clone,
         B: Clone,
-        F: Fn(&A) -> Result<B, Err>,
+        F: FnMut(&A) -> Result<B, Err>,
     {
-        let recurse = |e: &SubExpr<S, A>| -> Result<SubExpr<S, B>, Err> {
-            Ok(e.as_ref().traverse_embed(map_embed)?.roll())
-        };
-        self.traverse_ref(
-            |e| recurse(e),
-            |_, e| recurse(e),
-            |x| Ok(S::clone(x)),
-            map_embed,
-            |x| Ok(Label::clone(x)),
-        )
+        self.visit(&mut visitor::TraverseEmbedVisitor(map_embed))
     }
 
     pub fn map_label<F>(&self, map_label: &F) -> Self
@@ -277,20 +268,11 @@ impl<S, A> Expr<S, A> {
 impl<N: Clone, E> Expr<N, E> {
     pub fn squash_embed<E2: Clone>(
         &self,
-        f: &impl Fn(&E) -> SubExpr<N, E2>,
+        f: impl FnMut(&E) -> SubExpr<N, E2>,
     ) -> SubExpr<N, E2> {
-        match self.as_ref() {
-            ExprF::Embed(e) => f(e),
-            _ => self
-                .map_ref(
-                    |e| e.as_ref().squash_embed(f),
-                    |_, e| e.as_ref().squash_embed(f),
-                    N::clone,
-                    |_| unreachable!(),
-                    Label::clone,
-                )
-                .roll(),
-        }
+        rc(trivial_result(
+            self.visit(&mut visitor::SquashEmbedVisitor(f)),
+        ))
     }
 }
 
@@ -359,7 +341,7 @@ impl<SE, L, N, E> ExprF<SE, L, N, E> {
         F4: FnOnce(&'a E) -> Result<E2, Err>,
         F5: FnMut(&'a L) -> Result<L2, Err>,
     {
-        self.visit(TraverseRefVisitor {
+        self.visit(visitor::TraverseRefVisitor {
             visit_subexpr,
             visit_under_binder,
             visit_note,
@@ -404,7 +386,7 @@ impl<SE, L, N, E> ExprF<SE, L, N, E> {
         E: Clone,
         F1: FnMut(&'a SE) -> Result<SE2, Err>,
     {
-        self.visit(TraverseRefSimpleVisitor { visit_subexpr })
+        self.visit(visitor::TraverseRefSimpleVisitor { visit_subexpr })
     }
 
     pub fn map_ref_simple<'a, SE2, F1>(
@@ -517,65 +499,31 @@ impl<N: Clone> SubExpr<N, X> {
 
 impl<E: Clone> SubExpr<X, E> {
     pub fn note_absurd<N>(&self) -> SubExpr<N, E> {
-        rc(self.as_ref().map_ref(
-            |e| e.note_absurd(),
-            |_, e| e.note_absurd(),
-            |_| unreachable!(),
-            E::clone,
-            Label::clone,
-        ))
+        rc(self.as_ref().note_absurd())
     }
 }
 
 impl<E: Clone> Expr<X, E> {
-    pub fn note_absurd<N: Clone>(&self) -> Expr<N, E> {
-        self.roll().note_absurd().unroll()
+    pub fn note_absurd<N>(&self) -> Expr<N, E> {
+        trivial_result(self.visit(&mut visitor::NoteAbsurdVisitor))
     }
 }
 
-impl<N: Clone, E: Clone> SubExpr<N, E> {
+impl<N, E: Clone> SubExpr<N, E> {
     pub fn unnote(&self) -> SubExpr<X, E> {
-        match self.as_ref() {
-            ExprF::Note(_, e) => e.unnote(),
-            e => rc(e.map_ref(
-                |e| e.unnote(),
-                |_, e| e.unnote(),
-                |_| unreachable!(),
-                E::clone,
-                Label::clone,
-            )),
-        }
+        rc(trivial_result(
+            self.as_ref().visit(&mut visitor::UnNoteVisitor),
+        ))
     }
 }
 
 impl<N: Clone> Expr<N, X> {
-    // This is all very sad and I hope this can be avoided sometime
+    // Deprecated, use embed_absurd instead
     pub fn absurd_rec<T>(&self) -> Expr<N, T> {
-        self.map_ref(
-            |e| e.absurd(),
-            |_, e| e.absurd(),
-            |_| unreachable!(),
-            |_| unreachable!(),
-            Label::clone,
-        )
+        self.embed_absurd()
     }
-}
-
-impl<SE, L, N, E> ExprF<SE, L, N, E>
-where
-    SE: Clone,
-    L: Clone + Ord,
-    N: Clone,
-{
-    // When we know there is no Embed
-    pub fn absurd<T>(&self) -> ExprF<SE, L, N, T> {
-        self.map_ref(
-            |e| e.clone(),
-            |_, e| e.clone(),
-            N::clone,
-            |_| unreachable!(),
-            L::clone,
-        )
+    pub fn embed_absurd<T>(&self) -> Expr<N, T> {
+        trivial_result(self.visit(&mut visitor::EmbedAbsurdVisitor))
     }
 }
 
diff --git a/dhall_core/src/visitor.rs b/dhall_core/src/visitor.rs
index 1b50e46..3a94c7b 100644
--- a/dhall_core/src/visitor.rs
+++ b/dhall_core/src/visitor.rs
@@ -224,3 +224,155 @@ where
         Ok(L::clone(label))
     }
 }
+
+pub struct TraverseEmbedVisitor<F1>(pub F1);
+
+impl<'a, 'b, N, E, E2, Err, F1>
+    ExprVisitor<'a, SubExpr<N, E>, SubExpr<N, E2>, Label, Label, N, N, E, E2>
+    for &'b mut TraverseEmbedVisitor<F1>
+where
+    N: Clone + 'a,
+    E2: Clone,
+    F1: FnMut(&E) -> Result<E2, Err>,
+{
+    type Error = Err;
+
+    fn visit_subexpr(
+        &mut self,
+        subexpr: &'a SubExpr<N, E>,
+    ) -> Result<SubExpr<N, E2>, Self::Error> {
+        Ok(rc(subexpr.as_ref().visit(&mut **self)?))
+    }
+    fn visit_note(self, note: &'a N) -> Result<N, Self::Error> {
+        Ok(N::clone(note))
+    }
+    fn visit_embed(self, embed: &'a E) -> Result<E2, Self::Error> {
+        (self.0)(embed)
+    }
+    fn visit_label(&mut self, label: &'a Label) -> Result<Label, Self::Error> {
+        Ok(Label::clone(label))
+    }
+}
+
+pub struct SquashEmbedVisitor<F1>(pub F1);
+
+impl<'a, 'b, N, E, E2, F1>
+    ExprVisitor<'a, SubExpr<N, E>, SubExpr<N, E2>, Label, Label, N, N, E, E2>
+    for &'b mut SquashEmbedVisitor<F1>
+where
+    N: Clone + 'a,
+    E2: Clone,
+    F1: FnMut(&E) -> SubExpr<N, E2>,
+{
+    type Error = X;
+
+    fn visit_subexpr(
+        &mut self,
+        subexpr: &'a SubExpr<N, E>,
+    ) -> Result<SubExpr<N, E2>, Self::Error> {
+        Ok(rc(subexpr.as_ref().visit(&mut **self)?))
+    }
+    fn visit_note(self, note: &'a N) -> Result<N, Self::Error> {
+        Ok(N::clone(note))
+    }
+    fn visit_embed(self, _: &'a E) -> Result<E2, Self::Error> {
+        unreachable!()
+    }
+    fn visit_embed_squash(
+        self,
+        embed: &'a E,
+    ) -> Result<Expr<N, E2>, Self::Error> {
+        Ok((self.0)(embed).unroll())
+    }
+    fn visit_label(&mut self, label: &'a Label) -> Result<Label, Self::Error> {
+        Ok(Label::clone(label))
+    }
+}
+
+pub struct UnNoteVisitor;
+
+impl<'a, 'b, N, E>
+    ExprVisitor<'a, SubExpr<N, E>, SubExpr<X, E>, Label, Label, N, X, E, E>
+    for &'b mut UnNoteVisitor
+where
+    E: Clone + 'a,
+{
+    type Error = X;
+
+    fn visit_subexpr(
+        &mut self,
+        subexpr: &'a SubExpr<N, E>,
+    ) -> Result<SubExpr<X, E>, Self::Error> {
+        Ok(rc(subexpr.as_ref().visit(&mut **self)?))
+    }
+    fn visit_note(self, _: &'a N) -> Result<X, Self::Error> {
+        unreachable!()
+    }
+    fn visit_note_squash(
+        self,
+        _: &'a N,
+        subexpr: &'a SubExpr<N, E>,
+    ) -> Result<Expr<X, E>, Self::Error> {
+        subexpr.as_ref().visit(self)
+    }
+    fn visit_embed(self, embed: &'a E) -> Result<E, Self::Error> {
+        Ok(E::clone(embed))
+    }
+    fn visit_label(&mut self, label: &'a Label) -> Result<Label, Self::Error> {
+        Ok(Label::clone(label))
+    }
+}
+
+pub struct NoteAbsurdVisitor;
+
+impl<'a, 'b, N, E>
+    ExprVisitor<'a, SubExpr<X, E>, SubExpr<N, E>, Label, Label, X, N, E, E>
+    for &'b mut NoteAbsurdVisitor
+where
+    E: Clone + 'a,
+{
+    type Error = X;
+
+    fn visit_subexpr(
+        &mut self,
+        subexpr: &'a SubExpr<X, E>,
+    ) -> Result<SubExpr<N, E>, Self::Error> {
+        Ok(rc(subexpr.as_ref().visit(&mut **self)?))
+    }
+    fn visit_note(self, note: &'a X) -> Result<N, Self::Error> {
+        match *note {}
+    }
+    fn visit_embed(self, embed: &'a E) -> Result<E, Self::Error> {
+        Ok(E::clone(embed))
+    }
+    fn visit_label(&mut self, label: &'a Label) -> Result<Label, Self::Error> {
+        Ok(Label::clone(label))
+    }
+}
+
+pub struct EmbedAbsurdVisitor;
+
+impl<'a, 'b, N, E>
+    ExprVisitor<'a, SubExpr<N, X>, SubExpr<N, E>, Label, Label, N, N, X, E>
+    for &'b mut EmbedAbsurdVisitor
+where
+    N: Clone + 'a,
+{
+    type Error = X;
+
+    fn visit_subexpr(
+        &mut self,
+        subexpr: &'a SubExpr<N, X>,
+    ) -> Result<SubExpr<N, E>, Self::Error> {
+        Ok(rc(subexpr.as_ref().visit(&mut **self)?))
+    }
+    fn visit_note(self, note: &'a N) -> Result<N, Self::Error> {
+        Ok(N::clone(note))
+    }
+    fn visit_embed(self, embed: &'a X) -> Result<E, Self::Error> {
+        match *embed {}
+    }
+    fn visit_label(&mut self, label: &'a Label) -> Result<Label, Self::Error> {
+        Ok(Label::clone(label))
+    }
+}