2 files changed, 42 insertions, 0 deletions
diff --git a/dhall_core/src/core.rs b/dhall_core/src/core.rs
index a56c5a3..89506ec 100644
--- a/dhall_core/src/core.rs
+++ b/dhall_core/src/core.rs
@@ -68,6 +68,7 @@ impl From<NaiveDouble> for f64 {
 pub enum Const {
     Type,
     Kind,
+    Sort,
 }
 
 /// Bound variable
@@ -285,6 +286,28 @@ impl<S, A> Expr<S, A> {
         self.map_shallow(recurse, |x| x.clone(), map_embed, |x| x.clone())
     }
 
+    #[inline(always)]
+    pub fn traverse_embed<B, Err, F>(
+        &self,
+        map_embed: &F,
+    ) -> Result<Expr<S, B>, Err>
+    where
+        S: Clone,
+        B: Clone,
+        F: Fn(&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.as_ref().traverse(
+            |e| recurse(e),
+            |_, e| recurse(e),
+            |x| Ok(S::clone(x)),
+            map_embed,
+            |x| Ok(Label::clone(x)),
+        )
+    }
+
     pub fn map_label<F>(&self, map_label: &F) -> Self
     where
         A: Clone,
@@ -319,6 +342,23 @@ impl<S: Clone, A: Clone> Expr<S, Expr<S, A>> {
     }
 }
 
+impl<S: Clone, A: Clone> Expr<S, SubExpr<S, A>> {
+    pub fn squash_embed(&self) -> SubExpr<S, A> {
+        match self.as_ref() {
+            ExprF::Embed(e) => e.clone(),
+            e => e
+                .map(
+                    |e| e.as_ref().squash_embed(),
+                    |_, e| e.as_ref().squash_embed(),
+                    S::clone,
+                    |_| unreachable!(),
+                    Label::clone,
+                )
+                .roll(),
+        }
+    }
+}
+
 impl<SE, L, N, E> ExprF<SE, L, N, E> {
     #[inline(always)]
     pub fn as_ref(&self) -> ExprF<&SE, &L, &N, &E>
diff --git a/dhall_core/src/parser.rs b/dhall_core/src/parser.rs
index 67bcf77..41a2ce7 100644
--- a/dhall_core/src/parser.rs
+++ b/dhall_core/src/parser.rs
@@ -764,6 +764,7 @@ make_parser! {
                     "False" => BoolLit(false),
                     "Type" => Const(crate::Const::Type),
                     "Kind" => Const(crate::Const::Kind),
+                    "Sort" => Const(crate::Const::Sort),
                     _ => Var(V(l, idx)),
                 }
             }
@@ -777,6 +778,7 @@ make_parser! {
                     "False" => BoolLit(false),
                     "Type" => Const(crate::Const::Type),
                     "Kind" => Const(crate::Const::Kind),
+                    "Sort" => Const(crate::Const::Sort),
                     _ => Var(V(l, 0)),
                 }
             }