Rewrite map methods with Visitor traits

Closes #69

3 files changed, 261 insertions, 110 deletions

diff --git a/dhall_core/src/core.rs b/dhall_core/src/core.rs
index 1de363c..fe53bb2 100644
--- a/dhall_core/src/core.rs
+++ b/dhall_core/src/core.rs
@@ -1,8 +1,10 @@
 #![allow(non_snake_case)]
-use crate::*;
 use std::collections::BTreeMap;
 use std::rc::Rc;
 
+use crate::visitor::{TraverseRefSimpleVisitor, TraverseRefVisitor};
+use crate::*;
+
 pub type Integer = isize;
 pub type Natural = usize;
 pub type Double = NaiveDouble;
@@ -209,8 +211,8 @@ impl<S, A> Expr<S, A> {
         T: Clone,
         S: Clone,
         F1: Fn(&Self) -> Expr<T, B>,
-        F2: FnOnce(&S) -> T,
-        F3: FnOnce(&A) -> B,
+        F2: Fn(&S) -> T,
+        F3: Fn(&A) -> B,
         F4: Fn(&Label) -> Label,
     {
         self.map_ref(
@@ -340,13 +342,13 @@ impl<SE, L, N, E> ExprF<SE, L, N, E> {
         }
     }
 
-    fn traverse_ref<'a, SE2, L2, N2, E2, Err, F1, F2, F3, F4, F5>(
+    pub fn traverse_ref<'a, SE2, L2, N2, E2, Err, F1, F2, F3, F4, F5>(
         &'a self,
-        map_subexpr: F1,
-        map_under_binder: F2,
-        map_note: F3,
-        map_embed: F4,
-        mut map_label: F5,
+        visit_subexpr: F1,
+        visit_under_binder: F2,
+        visit_note: F3,
+        visit_embed: F4,
+        visit_label: F5,
     ) -> Result<ExprF<SE2, L2, N2, E2>, Err>
     where
         L: Ord,
@@ -357,95 +359,30 @@ impl<SE, L, N, E> ExprF<SE, L, N, E> {
         F4: FnOnce(&'a E) -> Result<E2, Err>,
         F5: FnMut(&'a L) -> Result<L2, Err>,
     {
-        let mut map = map_subexpr;
-        fn vec<'a, T, U, Err, F: FnMut(&'a T) -> Result<U, Err>>(
-            x: &'a Vec<T>,
-            f: F,
-        ) -> Result<Vec<U>, Err> {
-            x.iter().map(f).collect()
-        }
-        fn opt<'a, T, U, Err, F: FnOnce(&'a T) -> Result<U, Err>>(
-            x: &'a Option<T>,
-            f: F,
-        ) -> Result<Option<U>, Err> {
-            Ok(match x {
-                Some(x) => Some(f(x)?),
-                None => None,
-            })
-        }
-        fn btmap<'a, K, L, T, U, Err, FK, FV>(
-            x: &'a BTreeMap<K, T>,
-            mut fk: FK,
-            mut fv: FV,
-        ) -> Result<BTreeMap<L, U>, Err>
-        where
-            K: Ord,
-            L: Ord,
-            FK: FnMut(&'a K) -> Result<L, Err>,
-            FV: FnMut(&'a T) -> Result<U, Err>,
-        {
-            x.into_iter().map(|(k, v)| Ok((fk(k)?, fv(v)?))).collect()
-        }
-
-        use crate::ExprF::*;
-        Ok(match self {
-            Var(V(l, n)) => Var(V(map_label(l)?, *n)),
-            Lam(l, t, b) => {
-                let b = map_under_binder(&l, b)?;
-                Lam(map_label(l)?, map(t)?, b)
-            }
-            Pi(l, t, b) => {
-                let b = map_under_binder(&l, b)?;
-                Pi(map_label(l)?, map(t)?, b)
-            }
-            Let(l, t, a, b) => {
-                let b = map_under_binder(&l, b)?;
-                Let(map_label(l)?, opt(t, &mut map)?, map(a)?, b)
-            }
-            App(f, args) => App(map(f)?, vec(args, map)?),
-            Annot(x, t) => Annot(map(x)?, map(t)?),
-            Const(k) => Const(*k),
-            Builtin(v) => Builtin(*v),
-            BoolLit(b) => BoolLit(*b),
-            NaturalLit(n) => NaturalLit(*n),
-            IntegerLit(n) => IntegerLit(*n),
-            DoubleLit(n) => DoubleLit(*n),
-            TextLit(t) => TextLit(t.traverse_ref(map)?),
-            BinOp(o, x, y) => BinOp(*o, map(x)?, map(y)?),
-            BoolIf(b, t, f) => BoolIf(map(b)?, map(t)?, map(f)?),
-            EmptyListLit(t) => EmptyListLit(map(t)?),
-            NEListLit(es) => NEListLit(vec(es, map)?),
-            EmptyOptionalLit(t) => EmptyOptionalLit(map(t)?),
-            NEOptionalLit(e) => NEOptionalLit(map(e)?),
-            RecordType(kts) => RecordType(btmap(kts, map_label, map)?),
-            RecordLit(kvs) => RecordLit(btmap(kvs, map_label, map)?),
-            UnionType(kts) => UnionType(btmap(kts, map_label, map)?),
-            UnionLit(k, v, kvs) => {
-                UnionLit(map_label(k)?, map(v)?, btmap(kvs, map_label, map)?)
-            }
-            Merge(x, y, t) => Merge(map(x)?, map(y)?, opt(t, map)?),
-            Field(e, l) => Field(map(e)?, map_label(l)?),
-            Projection(e, ls) => Projection(map(e)?, vec(ls, map_label)?),
-            Note(n, e) => Note(map_note(n)?, map(e)?),
-            Embed(a) => Embed(map_embed(a)?),
+        self.visit(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>(
         &'a self,
         mut map_subexpr: F1,
-        map_under_binder: F2,
-        map_note: F3,
-        map_embed: F4,
+        mut map_under_binder: F2,
+        mut map_note: F3,
+        mut map_embed: F4,
         mut map_label: F5,
     ) -> ExprF<SE2, L2, N2, E2>
     where
         L: Ord,
         L2: Ord,
         F1: FnMut(&'a SE) -> SE2,
-        F2: FnOnce(&'a L, &'a SE) -> SE2,
-        F3: FnOnce(&'a N) -> N2,
-        F4: FnOnce(&'a E) -> E2,
+        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(
@@ -457,44 +394,31 @@ impl<SE, L, N, E> ExprF<SE, L, N, E> {
         ))
     }
 
-    pub fn map_ref_simple<'a, SE2, F1>(
+    pub fn traverse_ref_simple<'a, SE2, Err, F1>(
         &'a self,
-        map_subexpr: F1,
-    ) -> ExprF<SE2, L, N, E>
+        visit_subexpr: F1,
+    ) -> Result<ExprF<SE2, L, N, E>, Err>
     where
-        L: Ord,
-        L: Clone,
+        L: Ord + Clone,
         N: Clone,
         E: Clone,
-        F1: Fn(&'a SE) -> SE2,
+        F1: FnMut(&'a SE) -> Result<SE2, Err>,
     {
-        self.map_ref(
-            &map_subexpr,
-            |_, e| map_subexpr(e),
-            N::clone,
-            E::clone,
-            L::clone,
-        )
+        self.visit(TraverseRefSimpleVisitor { visit_subexpr })
     }
 
-    pub fn traverse_ref_simple<'a, SE2, Err, F1>(
+    pub fn map_ref_simple<'a, SE2, F1>(
         &'a self,
         map_subexpr: F1,
-    ) -> Result<ExprF<SE2, L, N, E>, Err>
+    ) -> ExprF<SE2, L, N, E>
     where
         L: Ord,
         L: Clone,
         N: Clone,
         E: Clone,
-        F1: Fn(&'a SE) -> Result<SE2, Err>,
+        F1: Fn(&'a SE) -> SE2,
     {
-        self.traverse_ref(
-            &map_subexpr,
-            |_, e| map_subexpr(e),
-            |x| Ok(N::clone(x)),
-            |x| Ok(E::clone(x)),
-            |x| Ok(L::clone(x)),
-        )
+        trivial_result(self.traverse_ref_simple(|x| Ok(map_subexpr(x))))
     }
 
     // pub fn zip<SE2, L2, N2, E2>(
@@ -552,7 +476,7 @@ impl<N, E> SubExpr<N, E> {
     ) -> Self
     where
         F1: FnMut(&'a Self) -> Self,
-        F2: FnOnce(&'a Label, &'a Self) -> Self,
+        F2: FnMut(&'a Label, &'a Self) -> Self,
     {
         match self.as_ref() {
             ExprF::Embed(_) => SubExpr::clone(self),
diff --git a/dhall_core/src/lib.rs b/dhall_core/src/lib.rs
index 887f29e..62c0ecd 100644
--- a/dhall_core/src/lib.rs
+++ b/dhall_core/src/lib.rs
@@ -20,3 +20,4 @@ pub use crate::printer::*;
 mod parser;
 pub use crate::parser::*;
 pub mod context;
+pub mod visitor;
diff --git a/dhall_core/src/visitor.rs b/dhall_core/src/visitor.rs
new file mode 100644
index 0000000..1b50e46
--- /dev/null
+++ b/dhall_core/src/visitor.rs
@@ -0,0 +1,226 @@
+use std::collections::BTreeMap;
+
+use crate::*;
+
+pub trait ExprVisitor<'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<SE, L, N, E> ExprF<SE, L, N, E> {
+    pub fn visit<'a, V, SE2, L2, N2, E2>(
+        &'a self,
+        mut v: V,
+    ) -> Result<ExprF<SE2, L2, N2, E2>, V::Error>
+    where
+        L: Ord,
+        L2: Ord,
+        V: ExprVisitor<'a, SE, SE2, L, L2, N, N2, E, E2>,
+    {
+        fn vec<'a, T, U, Err, F: FnMut(&'a T) -> Result<U, Err>>(
+            x: &'a Vec<T>,
+            f: F,
+        ) -> Result<Vec<U>, Err> {
+            x.iter().map(f).collect()
+        }
+        fn opt<'a, T, U, Err, F: FnOnce(&'a T) -> Result<U, Err>>(
+            x: &'a Option<T>,
+            f: F,
+        ) -> Result<Option<U>, Err> {
+            Ok(match x {
+                Some(x) => Some(f(x)?),
+                None => None,
+            })
+        }
+        fn btmap<'a, V, SE, L, N, E, SE2, L2, N2, E2>(
+            x: &'a BTreeMap<L, SE>,
+            mut v: V,
+        ) -> Result<BTreeMap<L2, SE2>, V::Error>
+        where
+            L: Ord,
+            L2: Ord,
+            V: ExprVisitor<'a, SE, SE2, L, L2, N, N2, E, E2>,
+        {
+            x.into_iter()
+                .map(|(k, x)| Ok((v.visit_label(k)?, v.visit_subexpr(x)?)))
+                .collect()
+        }
+
+        use crate::ExprF::*;
+        Ok(match self {
+            Var(V(l, n)) => Var(V(v.visit_label(l)?, *n)),
+            Lam(l, t, e) => {
+                let t = v.visit_subexpr(t)?;
+                let (l, e) = v.visit_binder(l, e)?;
+                Lam(l, t, e)
+            }
+            Pi(l, t, e) => {
+                let t = v.visit_subexpr(t)?;
+                let (l, e) = v.visit_binder(l, e)?;
+                Pi(l, t, e)
+            }
+            Let(l, t, a, e) => {
+                let t = opt(t, &mut |e| v.visit_subexpr(e))?;
+                let a = v.visit_subexpr(a)?;
+                let (l, e) = v.visit_binder(l, e)?;
+                Let(l, t, a, e)
+            }
+            App(f, args) => {
+                App(v.visit_subexpr(f)?, vec(args, |e| v.visit_subexpr(e))?)
+            }
+            Annot(x, t) => Annot(v.visit_subexpr(x)?, v.visit_subexpr(t)?),
+            Const(k) => Const(*k),
+            Builtin(v) => Builtin(*v),
+            BoolLit(b) => BoolLit(*b),
+            NaturalLit(n) => NaturalLit(*n),
+            IntegerLit(n) => IntegerLit(*n),
+            DoubleLit(n) => DoubleLit(*n),
+            TextLit(t) => TextLit(t.traverse_ref(|e| v.visit_subexpr(e))?),
+            BinOp(o, x, y) => {
+                BinOp(*o, v.visit_subexpr(x)?, v.visit_subexpr(y)?)
+            }
+            BoolIf(b, t, f) => BoolIf(
+                v.visit_subexpr(b)?,
+                v.visit_subexpr(t)?,
+                v.visit_subexpr(f)?,
+            ),
+            EmptyListLit(t) => EmptyListLit(v.visit_subexpr(t)?),
+            NEListLit(es) => NEListLit(vec(es, |e| v.visit_subexpr(e))?),
+            EmptyOptionalLit(t) => EmptyOptionalLit(v.visit_subexpr(t)?),
+            NEOptionalLit(e) => NEOptionalLit(v.visit_subexpr(e)?),
+            RecordType(kts) => RecordType(btmap(kts, v)?),
+            RecordLit(kvs) => RecordLit(btmap(kvs, v)?),
+            UnionType(kts) => UnionType(btmap(kts, v)?),
+            UnionLit(k, x, kvs) => {
+                UnionLit(v.visit_label(k)?, v.visit_subexpr(x)?, btmap(kvs, v)?)
+            }
+            Merge(x, y, t) => Merge(
+                v.visit_subexpr(x)?,
+                v.visit_subexpr(y)?,
+                opt(t, |e| v.visit_subexpr(e))?,
+            ),
+            Field(e, l) => Field(v.visit_subexpr(e)?, v.visit_label(l)?),
+            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)?,
+        })
+    }
+}
+
+pub struct TraverseRefVisitor<F1, F2, F3, F4, F5> {
+    pub visit_subexpr: F1,
+    pub visit_under_binder: F2,
+    pub visit_note: F3,
+    pub visit_embed: F4,
+    pub visit_label: F5,
+}
+
+impl<'a, SE, L, N, E, SE2, L2, N2, E2, Err, F1, F2, F3, F4, F5>
+    ExprVisitor<'a, SE, SE2, L, L2, N, N2, E, E2>
+    for TraverseRefVisitor<F1, F2, F3, F4, F5>
+where
+    SE: 'a,
+    L: 'a,
+    N: 'a,
+    E: 'a,
+    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>,
+{
+    type Error = Err;
+
+    fn visit_subexpr(&mut self, subexpr: &'a SE) -> Result<SE2, Self::Error> {
+        (self.visit_subexpr)(subexpr)
+    }
+    fn visit_subexpr_under_binder(
+        self,
+        label: &'a L,
+        subexpr: &'a SE,
+    ) -> Result<SE2, Self::Error> {
+        (self.visit_under_binder)(label, subexpr)
+    }
+    fn visit_note(self, note: &'a N) -> Result<N2, Self::Error> {
+        (self.visit_note)(note)
+    }
+    fn visit_embed(self, embed: &'a E) -> Result<E2, Self::Error> {
+        (self.visit_embed)(embed)
+    }
+    fn visit_label(&mut self, label: &'a L) -> Result<L2, Self::Error> {
+        (self.visit_label)(label)
+    }
+}
+
+pub struct TraverseRefSimpleVisitor<F1> {
+    pub visit_subexpr: F1,
+}
+
+impl<'a, SE, L, N, E, SE2, Err, F1> ExprVisitor<'a, SE, SE2, L, L, N, N, E, E>
+    for TraverseRefSimpleVisitor<F1>
+where
+    SE: 'a,
+    L: Ord + Clone + 'a,
+    N: Clone + 'a,
+    E: Clone + 'a,
+    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)
+    }
+    fn visit_note(self, note: &'a N) -> Result<N, Self::Error> {
+        Ok(N::clone(note))
+    }
+    fn visit_embed(self, embed: &'a E) -> Result<E, Self::Error> {
+        Ok(E::clone(embed))
+    }
+    fn visit_label(&mut self, label: &'a L) -> Result<L, Self::Error> {
+        Ok(L::clone(label))
+    }
+}