Implement bulk shifting

4 files changed, 134 insertions, 63 deletions

diff --git a/dhall/src/semantics/core/context.rs b/dhall/src/semantics/core/context.rs
index a6ab79e..f8d6ff0 100644
--- a/dhall/src/semantics/core/context.rs
+++ b/dhall/src/semantics/core/context.rs
@@ -40,32 +40,43 @@ impl TyCtx {
     }
     pub fn lookup(&self, var: &V<Label>) -> Option<Value> {
         let mut var = var.clone();
-        let mut shift_map: HashMap<Label, _> = HashMap::new();
-        for (b, i) in self.ctx.iter().rev() {
-            let l = b.to_label();
-            match var.over_binder(&l) {
-                None => {
-                    return Some(match i {
-                        CtxItem::Kept(newvar, t) => Value::from_kind_and_type(
-                            ValueKind::Var(
-                                newvar.under_multiple_binders(&shift_map),
-                            ),
-                            t.under_multiple_binders(&shift_map),
-                        ),
-                        CtxItem::Replaced(v) => {
-                            v.under_multiple_binders(&shift_map)
-                        }
-                    });
+        let mut shift_map: HashMap<Label, usize> = HashMap::new();
+        let mut rev_ctx = self.ctx.iter().rev().map(|(b, i)| (b.to_label(), i));
+
+        let found_item = loop {
+            if let Some((label, item)) = rev_ctx.next() {
+                var = match var.over_binder(&label) {
+                    Some(newvar) => newvar,
+                    None => break item,
+                };
+                if let CtxItem::Kept(_, _) = item {
+                    *shift_map.entry(label).or_insert(0) += 1;
                 }
-                Some(newvar) => var = newvar,
-            };
-            if let CtxItem::Kept(_, _) = i {
-                *shift_map.entry(l).or_insert(0) += 1;
+            } else {
+                // Unbound variable
+                return None;
             }
-        }
-        // Unbound variable
-        None
+        };
+
+        let v = match found_item {
+            CtxItem::Kept(newvar, t) => Value::from_kind_and_type(
+                ValueKind::Var(newvar.clone()),
+                t.clone(),
+            ),
+            CtxItem::Replaced(v) => v.clone(),
+        };
+
+        let v = v.under_multiple_binders(&shift_map);
+        // for (x, n) in shift_map {
+        //     let x: AlphaVar = (&x).into();
+        //     // Can't fail since delta is positive
+        //     v = v.shift(n as isize, &x).unwrap();
+        // }
+        return Some(v);
     }
+    // pub fn lookup_alpha(&self, var: &AlphaVar) -> Option<Value> {
+    //     self.lookup(&var.normal)
+    // }
     pub fn new_binder(&self, l: &Label) -> Binder {
         Binder::new(l.clone())
     }
diff --git a/dhall/src/semantics/core/value.rs b/dhall/src/semantics/core/value.rs
index 049c2a9..5aa337d 100644
--- a/dhall/src/semantics/core/value.rs
+++ b/dhall/src/semantics/core/value.rs
@@ -240,7 +240,7 @@ impl Value {
     }
 
     pub(crate) fn shift(&self, delta: isize, var: &AlphaVar) -> Option<Self> {
-        Some(self.0.borrow().shift(delta, var)?.into_value())
+        Some(self.as_internal().shift(delta, var)?.into_value())
     }
     pub(crate) fn over_binder<T>(&self, x: T) -> Option<Self>
     where
@@ -258,16 +258,10 @@ impl Value {
     pub(crate) fn under_multiple_binders(
         &self,
         shift_map: &HashMap<Label, usize>,
-    ) -> Self
-    where
-        Self: Clone,
-    {
-        let mut v = self.clone();
-        for (x, n) in shift_map {
-            // Can't fail since delta is positive
-            v = v.shift(*n as isize, &x.into()).unwrap();
-        }
-        v
+    ) -> Self {
+        self.as_internal()
+            .under_multiple_binders(shift_map)
+            .into_value()
     }
     pub(crate) fn subst_shift(&self, var: &AlphaVar, val: &Value) -> Self {
         match &*self.as_kind() {
@@ -279,9 +273,30 @@ impl Value {
                 }
                 val.clone()
             }
-            _ => self.0.borrow().subst_shift(var, val).into_value(),
+            _ => self.as_internal().subst_shift(var, val).into_value(),
         }
     }
+    // pub(crate) fn subst_ctx(&self, ctx: &TyCtx) -> Self {
+    //     match &*self.as_kind() {
+    //         ValueKind::Var(var) => match ctx.lookup_alpha(var) {
+    //             Some(val) => val,
+    //             None => unreachable!("Internal type error"),
+    //         },
+    //         kind => {
+    //             let kind = kind.map_ref_with_special_handling_of_binders(
+    //                 |x| x.subst_ctx(ctx),
+    //                 |b, t, x| x.subst_ctx(&ctx.insert_type(b, t.clone())),
+    //             );
+    //             ValueInternal {
+    //                 form: Unevaled,
+    //                 kind,
+    //                 ty: self.as_internal().ty.clone(),
+    //                 span: self.as_internal().span.clone(),
+    //             }
+    //             .into_value()
+    //         }
+    //     }
+    // }
 }
 
 impl ValueInternal {
@@ -343,7 +358,7 @@ impl ValueInternal {
         }
     }
 
-    pub(crate) fn shift(&self, delta: isize, var: &AlphaVar) -> Option<Self> {
+    fn shift(&self, delta: isize, var: &AlphaVar) -> Option<Self> {
         Some(ValueInternal {
             form: self.form,
             kind: self.kind.shift(delta, var)?,
@@ -354,7 +369,21 @@ impl ValueInternal {
             span: self.span.clone(),
         })
     }
-    pub(crate) fn subst_shift(&self, var: &AlphaVar, val: &Value) -> Self {
+    fn under_multiple_binders(
+        &self,
+        shift_map: &HashMap<Label, usize>,
+    ) -> Self {
+        ValueInternal {
+            form: self.form,
+            kind: self.kind.under_multiple_binders(shift_map),
+            ty: match &self.ty {
+                None => None,
+                Some(ty) => Some(ty.under_multiple_binders(shift_map)),
+            },
+            span: self.span.clone(),
+        }
+    }
+    fn subst_shift(&self, var: &AlphaVar, val: &Value) -> Self {
         ValueInternal {
             // The resulting value may not stay in wnhf after substitution
             form: Unevaled,
@@ -449,17 +478,40 @@ impl ValueKind<Value> {
             ValueKind::Var(v) => ValueKind::Var(v.shift(delta, var)?),
             _ => self.traverse_ref_with_special_handling_of_binders(
                 |x| Ok(x.shift(delta, var)?),
-                |b, x| Ok(x.shift(delta, &var.under_binder(b))?),
+                |b, _, x| Ok(x.shift(delta, &var.under_binder(b))?),
             )?,
         })
     }
+    #[allow(dead_code)]
+    fn under_multiple_binders(
+        &self,
+        shift_map: &HashMap<Label, usize>,
+    ) -> Self {
+        match self {
+            ValueKind::Var(v) => {
+                ValueKind::Var(v.under_multiple_binders(shift_map))
+            }
+            _ => self.map_ref_with_special_handling_of_binders(
+                |v| v.under_multiple_binders(shift_map),
+                |b, _, v| {
+                    let mut v = v.clone();
+                    for (x, n) in shift_map {
+                        let x: AlphaVar = x.into();
+                        // Can't fail since delta is positive
+                        v = v.shift(*n as isize, &x.under_binder(b)).unwrap();
+                    }
+                    v
+                },
+            ),
+        }
+    }
     fn subst_shift(&self, var: &AlphaVar, val: &Value) -> Self {
         match self {
             ValueKind::Var(v) if v == var => val.to_whnf_ignore_type(),
-            ValueKind::Var(v) => ValueKind::Var(v.shift(-1, var).unwrap()),
+            ValueKind::Var(v) => ValueKind::Var(v.over_binder(var).unwrap()),
             _ => self.map_ref_with_special_handling_of_binders(
                 |x| x.subst_shift(var, val),
-                |b, x| {
+                |b, _, x| {
                     x.subst_shift(&var.under_binder(b), &val.under_binder(b))
                 },
             ),
@@ -471,7 +523,11 @@ impl<V> ValueKind<V> {
     pub(crate) fn traverse_ref_with_special_handling_of_binders<'a, V2, E>(
         &'a self,
         visit_val: impl FnMut(&'a V) -> Result<V2, E>,
-        visit_under_binder: impl FnMut(&'a Binder, &'a V) -> Result<V2, E>,
+        visit_under_binder: impl for<'b> FnMut(
+            &'a Binder,
+            &'b V,
+            &'a V,
+        ) -> Result<V2, E>,
     ) -> Result<ValueKind<V2>, E> {
         use crate::semantics::visitor;
         use visitor::ValueKindVisitor;
@@ -485,12 +541,12 @@ impl<V> ValueKind<V> {
     pub(crate) fn map_ref_with_special_handling_of_binders<'a, V2>(
         &'a self,
         mut map_val: impl FnMut(&'a V) -> V2,
-        mut map_under_binder: impl FnMut(&'a Binder, &'a V) -> V2,
+        mut map_under_binder: impl for<'b> FnMut(&'a Binder, &'b V, &'a V) -> V2,
     ) -> ValueKind<V2> {
         use crate::syntax::trivial_result;
         trivial_result(self.traverse_ref_with_special_handling_of_binders(
             |x| Ok(map_val(x)),
-            |l, x| Ok(map_under_binder(l, x)),
+            |l, t, x| Ok(map_under_binder(l, t, x)),
         ))
     }
 
@@ -499,7 +555,7 @@ impl<V> ValueKind<V> {
         &'a self,
         map_val: impl Fn(&'a V) -> V2,
     ) -> ValueKind<V2> {
-        self.map_ref_with_special_handling_of_binders(&map_val, |_, x| {
+        self.map_ref_with_special_handling_of_binders(&map_val, |_, _, x| {
             map_val(x)
         })
     }
diff --git a/dhall/src/semantics/core/var.rs b/dhall/src/semantics/core/var.rs
index 7c2c4de..d7666e3 100644
--- a/dhall/src/semantics/core/var.rs
+++ b/dhall/src/semantics/core/var.rs
@@ -40,19 +40,17 @@ impl AlphaVar {
         // Can't fail since delta is positive
         self.shift(1, &x.into()).unwrap()
     }
+    pub(crate) fn over_binder(&self, x: &AlphaVar) -> Option<Self> {
+        self.shift(-1, x)
+    }
     pub(crate) fn under_multiple_binders(
         &self,
         shift_map: &HashMap<Label, usize>,
-    ) -> Self
-    where
-        Self: Clone,
-    {
-        let mut v = self.clone();
-        for (x, n) in shift_map {
-            // Can't fail since delta is positive
-            v = v.shift(*n as isize, &x.into()).unwrap();
+    ) -> Self {
+        AlphaVar {
+            normal: self.normal.under_multiple_binders(shift_map),
+            alpha: self.alpha.under_multiple_binders(shift_map),
         }
-        v
     }
 }
 
diff --git a/dhall/src/semantics/core/visitor.rs b/dhall/src/semantics/core/visitor.rs
index 2c9ba92..e61a649 100644
--- a/dhall/src/semantics/core/visitor.rs
+++ b/dhall/src/semantics/core/visitor.rs
@@ -10,12 +10,13 @@ pub(crate) trait ValueKindVisitor<'a, V1, V2>: Sized {
     type Error;
 
     fn visit_val(&mut self, val: &'a V1) -> Result<V2, Self::Error>;
-    fn visit_val_under_binder(
+    fn visit_binder(
         mut self,
         _label: &'a Binder,
+        ty: &'a V1,
         val: &'a V1,
-    ) -> Result<V2, Self::Error> {
-        self.visit_val(val)
+    ) -> Result<(V2, V2), Self::Error> {
+        Ok((self.visit_val(ty)?, self.visit_val(val)?))
     }
     fn visit_vec(&mut self, x: &'a [V1]) -> Result<Vec<V2>, Self::Error> {
         x.iter().map(|e| self.visit_val(e)).collect()
@@ -72,10 +73,12 @@ where
     use ValueKind::*;
     Ok(match input {
         Lam(l, t, e) => {
-            Lam(l.clone(), v.visit_val(t)?, v.visit_val_under_binder(l, e)?)
+            let (t, e) = v.visit_binder(l, t, e)?;
+            Lam(l.clone(), t, e)
         }
         Pi(l, t, e) => {
-            Pi(l.clone(), v.visit_val(t)?, v.visit_val_under_binder(l, e)?)
+            let (t, e) = v.visit_binder(l, t, e)?;
+            Pi(l.clone(), t, e)
         }
         AppliedBuiltin(b, xs) => AppliedBuiltin(*b, v.visit_vec(xs)?),
         Var(v) => Var(v.clone()),
@@ -117,18 +120,21 @@ impl<'a, V1, V2, Err, F1, F2> ValueKindVisitor<'a, V1, V2>
 where
     V1: 'a,
     F1: FnMut(&'a V1) -> Result<V2, Err>,
-    F2: FnOnce(&'a Binder, &'a V1) -> Result<V2, Err>,
+    F2: for<'b> FnOnce(&'a Binder, &'b V1, &'a V1) -> Result<V2, Err>,
 {
     type Error = Err;
 
     fn visit_val(&mut self, val: &'a V1) -> Result<V2, Self::Error> {
         (self.visit_val)(val)
     }
-    fn visit_val_under_binder(
-        self,
+    fn visit_binder<'b>(
+        mut self,
         label: &'a Binder,
+        ty: &'a V1,
         val: &'a V1,
-    ) -> Result<V2, Self::Error> {
-        (self.visit_under_binder)(label, val)
+    ) -> Result<(V2, V2), Self::Error> {
+        let val = (self.visit_under_binder)(label, ty, val)?;
+        let ty = (self.visit_val)(ty)?;
+        Ok((ty, val))
     }
 }