subst_shift in a single pass

2 files changed, 71 insertions, 5 deletions

diff --git a/dhall_core/src/context.rs b/dhall_core/src/context.rs
index 877843d..b500226 100644
--- a/dhall_core/src/context.rs
+++ b/dhall_core/src/context.rs
@@ -43,6 +43,20 @@ impl<K: Hash + Eq + Clone, T> Context<K, T> {
                 .collect(),
         )
     }
+
+    pub fn lookup_all<'a>(&'a self, k: &K) -> impl Iterator<Item = &T> {
+        self.0.get(k).into_iter().flat_map(|v| v.iter())
+    }
+
+    pub fn iter<'a>(&'a self) -> impl Iterator<Item = (&K, &T)> {
+        self.0
+            .iter()
+            .flat_map(|(k, vs)| vs.iter().map(move |v| (k, v)))
+    }
+
+    pub fn iter_keys<'a>(&'a self) -> impl Iterator<Item = (&K, &Vec<T>)> {
+        self.0.iter()
+    }
 }
 
 impl<K: Hash + Eq + Clone, T: Clone> Context<K, T> {
diff --git a/dhall_core/src/core.rs b/dhall_core/src/core.rs
index aeb6f23..56ab66c 100644
--- a/dhall_core/src/core.rs
+++ b/dhall_core/src/core.rs
@@ -1,7 +1,9 @@
 #![allow(non_snake_case)]
 use std::collections::BTreeMap;
+use std::collections::HashMap;
 use std::rc::Rc;
 
+use crate::context::Context;
 use crate::visitor;
 use crate::*;
 
@@ -486,6 +488,7 @@ fn shift_var(delta: isize, var: &V<Label>, in_expr: &V<Label>) -> V<Label> {
 /// capture by shifting variable indices
 /// See https://github.com/dhall-lang/dhall-lang/blob/master/standard/semantics.md#shift
 /// for details
+///
 pub fn shift<N, E>(
     delta: isize,
     var: &V<Label>,
@@ -505,6 +508,41 @@ pub fn shift<N, E>(
     }
 }
 
+pub fn shift0<N, E>(
+    delta: isize,
+    label: &Label,
+    in_expr: &SubExpr<N, E>,
+) -> SubExpr<N, E> {
+    shift(delta, &V(label.clone(), 0), in_expr)
+}
+
+fn shift0_multiple<N, E>(
+    deltas: &HashMap<Label, isize>,
+    in_expr: &SubExpr<N, E>,
+) -> SubExpr<N, E> {
+    shift0_multiple_inner(&Context::new(), deltas, in_expr)
+}
+
+fn shift0_multiple_inner<N, E>(
+    ctx: &Context<Label, ()>,
+    deltas: &HashMap<Label, isize>,
+    in_expr: &SubExpr<N, E>,
+) -> SubExpr<N, E> {
+    use crate::ExprF::*;
+    match in_expr.as_ref() {
+        Var(V(y, m)) if ctx.lookup(y, *m).is_none() => {
+            let delta = deltas.get(y).unwrap_or(&0);
+            rc(Var(V(y.clone(), add_ui(*m, *delta))))
+        }
+        _ => in_expr.map_subexprs_with_special_handling_of_binders(
+            |e| shift0_multiple_inner(ctx, deltas, e),
+            |l: &Label, e| {
+                shift0_multiple_inner(&ctx.insert(l.clone(), ()), deltas, e)
+            },
+        ),
+    }
+}
+
 /// Substitute all occurrences of a variable with an expression, and
 /// removes the variable from the environment by shifting the indices
 /// of other variables appropriately.
@@ -518,15 +556,29 @@ pub fn subst_shift<N, E>(
     value: &SubExpr<N, E>,
     in_expr: &SubExpr<N, E>,
 ) -> SubExpr<N, E> {
+    subst_shift_inner(&HashMap::new(), var, value, in_expr)
+}
+
+fn subst_shift_inner<N, E>(
+    ctx: &HashMap<Label, isize>,
+    var: &V<Label>,
+    value: &SubExpr<N, E>,
+    in_expr: &SubExpr<N, E>,
+) -> SubExpr<N, E> {
     use crate::ExprF::*;
+    let V(x, n) = var;
+    let dn = ctx.get(x).unwrap_or(&0);
+    let actual_var = V(x.clone(), add_ui(*n, *dn));
     match in_expr.as_ref() {
-        Var(v) if v == var => SubExpr::clone(value),
-        Var(v) => rc(Var(shift_var(-1, var, v))),
+        Var(v) if v == &actual_var => shift0_multiple(ctx, value),
+        Var(v) => rc(Var(shift_var(-1, &actual_var, v))),
         _ => in_expr.map_subexprs_with_special_handling_of_binders(
-            |e| subst_shift(var, &value, e),
+            |e| subst_shift_inner(ctx, var, value, e),
             |l: &Label, e| {
-                let vl = V(l.clone(), 0);
-                subst_shift(&shift_var(1, &vl, var), &shift(1, &vl, value), e)
+                let mut ctx = ctx.clone();
+                let count = ctx.entry(l.clone()).or_insert(0);
+                *count += 1;
+                subst_shift_inner(&ctx, var, value, e)
             },
         ),
     }