1 files changed, 85 insertions, 36 deletions

diff --git a/dhall/src/syntax/ast/expr.rs b/dhall/src/syntax/ast/expr.rs
index 9e935dc..51426ce 100644
--- a/dhall/src/syntax/ast/expr.rs
+++ b/dhall/src/syntax/ast/expr.rs
@@ -126,63 +126,73 @@ pub enum NumKind {
     Double(Double),
 }
 
+/// Operations
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub enum OpKind<SubExpr> {
+    ///  `f a`
+    App(SubExpr, SubExpr),
+    /// Binary operations
+    BinOp(BinOp, SubExpr, SubExpr),
+    ///  `if x then y else z`
+    BoolIf(SubExpr, SubExpr, SubExpr),
+    ///  `merge x y : t`
+    Merge(SubExpr, SubExpr, Option<SubExpr>),
+    ///  `toMap x : t`
+    ToMap(SubExpr, Option<SubExpr>),
+    ///  `e.x`
+    Field(SubExpr, Label),
+    ///  `e.{ x, y, z }`
+    Projection(SubExpr, DupTreeSet<Label>),
+    ///  `e.(t)`
+    ProjectionByExpr(SubExpr, SubExpr),
+    ///  `x::y`
+    Completion(SubExpr, SubExpr),
+}
+
 /// Syntax tree for expressions
 // Having the recursion out of the enum definition enables writing
 // much more generic code and improves pattern-matching behind
 // smart pointers.
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum ExprKind<SubExpr> {
+    /// `Type`, `Kind` and `Sort`
     Const(Const),
+    /// Numbers and booleans
     Num(NumKind),
-    ///  `x`
-    ///  `x@n`
-    Var(V),
-    ///  `λ(x : A) -> b`
-    Lam(Label, SubExpr, SubExpr),
-    ///  `A -> B`
-    ///  `∀(x : A) -> B`
-    Pi(Label, SubExpr, SubExpr),
-    ///  `f a`
-    App(SubExpr, SubExpr),
-    ///  `let x     = r in e`
-    ///  `let x : t = r in e`
-    Let(Label, Option<SubExpr>, SubExpr, SubExpr),
-    ///  `x : t`
-    Annot(SubExpr, SubExpr),
-    ///  `assert : t`
-    Assert(SubExpr),
-    /// Built-in values
+    /// Built-in functions and types
     Builtin(Builtin),
-    // Binary operations
-    BinOp(BinOp, SubExpr, SubExpr),
-    ///  `if x then y else z`
-    BoolIf(SubExpr, SubExpr, SubExpr),
+    ///  `Some e`
+    SomeLit(SubExpr),
     ///  `"Some ${interpolated} text"`
     TextLit(InterpolatedText<SubExpr>),
     ///  `[] : t`
     EmptyListLit(SubExpr),
     ///  `[x, y, z]`
     NEListLit(Vec<SubExpr>),
-    ///  `Some e`
-    SomeLit(SubExpr),
     ///  `{ k1 : t1, k2 : t1 }`
     RecordType(DupTreeMap<Label, SubExpr>),
     ///  `{ k1 = v1, k2 = v2 }`
     RecordLit(BTreeMap<Label, SubExpr>),
     ///  `< k1 : t1, k2 >`
     UnionType(DupTreeMap<Label, Option<SubExpr>>),
-    ///  `merge x y : t`
-    Merge(SubExpr, SubExpr, Option<SubExpr>),
-    ///  `toMap x : t`
-    ToMap(SubExpr, Option<SubExpr>),
-    ///  `e.x`
-    Field(SubExpr, Label),
-    ///  `e.{ x, y, z }`
-    Projection(SubExpr, DupTreeSet<Label>),
-    ///  `e.(t)`
-    ProjectionByExpr(SubExpr, SubExpr),
-    ///  `x::y`
-    Completion(SubExpr, SubExpr),
+
+    ///  `x`, `x@n`
+    Var(V),
+    ///  `λ(x : A) -> b`
+    Lam(Label, SubExpr, SubExpr),
+    ///  `A -> B`, `∀(x : A) -> B`
+    Pi(Label, SubExpr, SubExpr),
+    ///  `let x : t = r in e`
+    Let(Label, Option<SubExpr>, SubExpr, SubExpr),
+
+    /// Operations
+    Op(OpKind<SubExpr>),
+
+    ///  `x : t`
+    Annot(SubExpr, SubExpr),
+    ///  `assert : t`
+    Assert(SubExpr),
+
     /// `./some/path`
     Import(Import<SubExpr>),
 }
@@ -239,6 +249,45 @@ impl<SE> ExprKind<SE> {
     }
 }
 
+impl<SE> OpKind<SE> {
+    pub fn traverse_ref<'a, SE2, Err>(
+        &'a self,
+        mut f: impl FnMut(&'a SE) -> Result<SE2, Err>,
+    ) -> Result<OpKind<SE2>, Err> {
+        // Can't use closures because of borrowing rules
+        macro_rules! expr {
+            ($e:expr) => {
+                f($e)?
+            };
+        }
+        macro_rules! opt {
+            ($e:expr) => {
+                $e.as_ref().map(|e| Ok(expr!(e))).transpose()?
+            };
+        }
+
+        use OpKind::*;
+        Ok(match self {
+            App(f, a) => App(expr!(f), expr!(a)),
+            BinOp(o, x, y) => BinOp(*o, expr!(x), expr!(y)),
+            BoolIf(b, t, f) => BoolIf(expr!(b), expr!(t), expr!(f)),
+            Merge(x, y, t) => Merge(expr!(x), expr!(y), opt!(t)),
+            ToMap(x, t) => ToMap(expr!(x), opt!(t)),
+            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)),
+        })
+    }
+
+    pub fn map_ref<'a, SE2>(
+        &'a self,
+        mut f: impl FnMut(&'a SE) -> SE2,
+    ) -> OpKind<SE2> {
+        trivial_result(self.traverse_ref(|x| Ok(f(x))))
+    }
+}
+
 impl Expr {
     pub fn as_ref(&self) -> &UnspannedExpr {
         &self.kind