Simplify ExprKind visitor

2 files changed, 84 insertions, 140 deletions

diff --git a/dhall/src/syntax/ast/expr.rs b/dhall/src/syntax/ast/expr.rs
index 722630f..512010a 100644
--- a/dhall/src/syntax/ast/expr.rs
+++ b/dhall/src/syntax/ast/expr.rs
@@ -1,5 +1,5 @@
 use crate::syntax::map::{DupTreeMap, DupTreeSet};
-use crate::syntax::visitor::{self, ExprKindVisitor};
+use crate::syntax::visitor;
 use crate::syntax::*;
 
 pub type Integer = isize;
@@ -176,7 +176,7 @@ impl<SE> ExprKind<SE> {
         &'a self,
         visit: impl FnMut(Option<&'a Label>, &'a SE) -> Result<SE2, Err>,
     ) -> Result<ExprKind<SE2>, Err> {
-        visitor::TraverseRefMaybeBinderVisitor(visit).visit(self)
+        visitor::visit_ref(self, visit)
     }
 
     pub fn traverse_ref_with_special_handling_of_binders<'a, SE2, Err>(
diff --git a/dhall/src/syntax/ast/visitor.rs b/dhall/src/syntax/ast/visitor.rs
index fc90efd..39959ac 100644
--- a/dhall/src/syntax/ast/visitor.rs
+++ b/dhall/src/syntax/ast/visitor.rs
@@ -2,171 +2,115 @@ use std::iter::FromIterator;
 
 use crate::syntax::*;
 
-/// A visitor trait that can be used to traverse `ExprKind`s. We need this pattern so that Rust lets
-/// us have as much mutability as we can.
-/// For example, `traverse_ref_with_special_handling_of_binders` cannot be made using only
-/// `traverse_ref`, because `traverse_ref` takes a `FnMut` so we would need to pass multiple
-/// mutable reverences to this argument to `traverse_ref`. But Rust's ownership system is all about
-/// preventing exactly this ! So we have to be more clever. The visitor pattern allows us to have
-/// only one mutable thing the whole 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.
-pub trait ExprKindVisitor<'a, SE1, SE2>: Sized {
-    type Error;
+fn vec<'a, T, U, Err>(
+    x: &'a [T],
+    f: impl FnMut(&'a T) -> Result<U, Err>,
+) -> Result<Vec<U>, Err> {
+    x.iter().map(f).collect()
+}
 
-    fn visit_subexpr(&mut self, subexpr: &'a SE1) -> Result<SE2, Self::Error>;
+fn opt<'a, T, U, Err>(
+    x: &'a Option<T>,
+    f: impl FnOnce(&'a T) -> Result<U, Err>,
+) -> Result<Option<U>, Err> {
+    Ok(match x {
+        Some(x) => Some(f(x)?),
+        None => None,
+    })
+}
 
-    fn visit_subexpr_under_binder(
-        mut self,
-        _label: &'a Label,
-        subexpr: &'a SE1,
-    ) -> Result<SE2, Self::Error> {
-        self.visit_subexpr(subexpr)
-    }
+fn dupmap<'a, SE1, SE2, T, Err>(
+    x: impl IntoIterator<Item = (&'a Label, &'a SE1)>,
+    mut f: impl FnMut(&'a SE1) -> Result<SE2, Err>,
+) -> Result<T, Err>
+where
+    SE1: 'a,
+    T: FromIterator<(Label, SE2)>,
+{
+    x.into_iter().map(|(k, x)| Ok((k.clone(), f(x)?))).collect()
+}
 
-    fn visit(
-        self,
-        input: &'a ExprKind<SE1>,
-    ) -> Result<ExprKind<SE2>, Self::Error> {
-        visit_ref(self, input)
-    }
+fn optdupmap<'a, SE1, SE2, T, Err>(
+    x: impl IntoIterator<Item = (&'a Label, &'a Option<SE1>)>,
+    mut f: impl FnMut(&'a SE1) -> Result<SE2, Err>,
+) -> Result<T, Err>
+where
+    SE1: 'a,
+    T: FromIterator<(Label, Option<SE2>)>,
+{
+    x.into_iter()
+        .map(|(k, x)| {
+            Ok((
+                k.clone(),
+                match x {
+                    Some(x) => Some(f(x)?),
+                    None => None,
+                },
+            ))
+        })
+        .collect()
 }
 
-fn visit_ref<'a, V, SE1, SE2>(
-    mut v: V,
+pub(crate) fn visit_ref<'a, F, SE1, SE2, Err>(
     input: &'a ExprKind<SE1>,
-) -> Result<ExprKind<SE2>, V::Error>
+    mut f: F,
+) -> Result<ExprKind<SE2>, Err>
 where
-    V: ExprKindVisitor<'a, SE1, SE2>,
+    F: FnMut(Option<&'a Label>, &'a SE1) -> Result<SE2, Err>,
 {
-    fn vec<'a, T, U, Err, F: FnMut(&'a T) -> Result<U, Err>>(
-        x: &'a [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 dupmap<'a, V, SE1, SE2, T>(
-        x: impl IntoIterator<Item = (&'a Label, &'a SE1)>,
-        mut v: V,
-    ) -> Result<T, V::Error>
-    where
-        SE1: 'a,
-        T: FromIterator<(Label, SE2)>,
-        V: ExprKindVisitor<'a, SE1, SE2>,
-    {
-        x.into_iter()
-            .map(|(k, x)| Ok((k.clone(), v.visit_subexpr(x)?)))
-            .collect()
-    }
-    fn optdupmap<'a, V, SE1, SE2, T>(
-        x: impl IntoIterator<Item = (&'a Label, &'a Option<SE1>)>,
-        mut v: V,
-    ) -> Result<T, V::Error>
-    where
-        SE1: 'a,
-        T: FromIterator<(Label, Option<SE2>)>,
-        V: ExprKindVisitor<'a, SE1, SE2>,
-    {
-        x.into_iter()
-            .map(|(k, x)| {
-                Ok((
-                    k.clone(),
-                    match x {
-                        Some(x) => Some(v.visit_subexpr(x)?),
-                        None => None,
-                    },
-                ))
-            })
-            .collect()
+    // Can't use closures because of borrowing rules
+    macro_rules! expr {
+        ($e:expr) => {
+            f(None, $e)
+        };
+        ($l:expr, $e:expr) => {
+            f(Some($l), $e)
+        };
     }
 
     use crate::syntax::ExprKind::*;
     Ok(match input {
         Var(v) => Var(v.clone()),
         Lam(l, t, e) => {
-            let t = v.visit_subexpr(t)?;
-            let e = v.visit_subexpr_under_binder(l, e)?;
+            let t = expr!(t)?;
+            let e = expr!(l, e)?;
             Lam(l.clone(), t, e)
         }
         Pi(l, t, e) => {
-            let t = v.visit_subexpr(t)?;
-            let e = v.visit_subexpr_under_binder(l, e)?;
+            let t = expr!(t)?;
+            let e = expr!(l, e)?;
             Pi(l.clone(), t, e)
         }
         Let(l, t, a, e) => {
-            let t = opt(t, &mut |e| v.visit_subexpr(e))?;
-            let a = v.visit_subexpr(a)?;
-            let e = v.visit_subexpr_under_binder(l, e)?;
+            let t = opt(t, &mut |e| expr!(e))?;
+            let a = expr!(a)?;
+            let e = expr!(l, e)?;
             Let(l.clone(), t, a, e)
         }
-        App(f, a) => App(v.visit_subexpr(f)?, v.visit_subexpr(a)?),
-        Annot(x, t) => Annot(v.visit_subexpr(x)?, v.visit_subexpr(t)?),
+        App(f, a) => App(expr!(f)?, expr!(a)?),
+        Annot(x, t) => Annot(expr!(x)?, expr!(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))?),
-        SomeLit(e) => SomeLit(v.visit_subexpr(e)?),
-        RecordType(kts) => RecordType(dupmap(kts, v)?),
-        RecordLit(kvs) => RecordLit(dupmap(kvs, v)?),
-        UnionType(kts) => UnionType(optdupmap(kts, v)?),
-        Merge(x, y, t) => Merge(
-            v.visit_subexpr(x)?,
-            v.visit_subexpr(y)?,
-            opt(t, |e| v.visit_subexpr(e))?,
-        ),
-        ToMap(x, t) => {
-            ToMap(v.visit_subexpr(x)?, opt(t, |e| v.visit_subexpr(e))?)
-        }
-        Field(e, l) => Field(v.visit_subexpr(e)?, l.clone()),
-        Projection(e, ls) => Projection(v.visit_subexpr(e)?, ls.clone()),
-        ProjectionByExpr(e, x) => {
-            ProjectionByExpr(v.visit_subexpr(e)?, v.visit_subexpr(x)?)
-        }
-        Completion(e, x) => {
-            Completion(v.visit_subexpr(e)?, v.visit_subexpr(x)?)
-        }
-        Assert(e) => Assert(v.visit_subexpr(e)?),
-        Import(i) => Import(i.traverse_ref(|e| v.visit_subexpr(e))?),
+        TextLit(t) => TextLit(t.traverse_ref(|e| expr!(e))?),
+        BinOp(o, x, y) => BinOp(*o, expr!(x)?, expr!(y)?),
+        BoolIf(b, t, f) => BoolIf(expr!(b)?, expr!(t)?, expr!(f)?),
+        EmptyListLit(t) => EmptyListLit(expr!(t)?),
+        NEListLit(es) => NEListLit(vec(es, |e| expr!(e))?),
+        SomeLit(e) => SomeLit(expr!(e)?),
+        RecordType(kts) => RecordType(dupmap(kts, |e| expr!(e))?),
+        RecordLit(kvs) => RecordLit(dupmap(kvs, |e| expr!(e))?),
+        UnionType(kts) => UnionType(optdupmap(kts, |e| expr!(e))?),
+        Merge(x, y, t) => Merge(expr!(x)?, expr!(y)?, opt(t, |e| expr!(e))?),
+        ToMap(x, t) => ToMap(expr!(x)?, opt(t, |e| expr!(e))?),
+        Field(e, l) => Field(expr!(e)?, l.clone()),
+        Projection(e, ls) => Projection(expr!(e)?, ls.clone()),
+        ProjectionByExpr(e, x) => ProjectionByExpr(expr!(e)?, expr!(x)?),
+        Completion(e, x) => Completion(expr!(e)?, expr!(x)?),
+        Assert(e) => Assert(expr!(e)?),
+        Import(i) => Import(i.traverse_ref(|e| expr!(e))?),
     })
 }
-
-pub struct TraverseRefMaybeBinderVisitor<F>(pub F);
-
-impl<'a, SE, SE2, Err, F> ExprKindVisitor<'a, SE, SE2>
-    for TraverseRefMaybeBinderVisitor<F>
-where
-    SE: 'a,
-    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.0)(None, subexpr)
-    }
-    fn visit_subexpr_under_binder(
-        mut self,
-        label: &'a Label,
-        subexpr: &'a SE,
-    ) -> Result<SE2, Self::Error> {
-        (self.0)(Some(label), subexpr)
-    }
-}