4 files changed, 147 insertions, 118 deletions

diff --git a/dhall/src/core/context.rs b/dhall/src/core/context.rs
index 236bebe..328bfdc 100644
--- a/dhall/src/core/context.rs
+++ b/dhall/src/core/context.rs
@@ -33,7 +33,7 @@ impl<T> Context<T> {
         T: Shift + Clone,
     {
         let mut vec = self.0.as_ref().clone();
-        vec.push((x.clone(), CtxItem::Kept(x.into(), t.shift(1, &x.into()))));
+        vec.push((x.clone(), CtxItem::Kept(x.into(), t.under_binder(x))));
         Context(Rc::new(vec))
     }
     pub(crate) fn insert_replaced(&self, x: &Label, th: Thunk, t: T) -> Self
@@ -51,11 +51,10 @@ impl<T> Context<T> {
         let mut var = var.clone();
         let mut shift_map: HashMap<Label, _> = HashMap::new();
         for (l, i) in self.0.iter().rev() {
-            if var == l.into() {
-                return Ok(i.shift0_multiple(&shift_map));
-            } else {
-                var = var.shift(-1, &l.into());
-            }
+            match var.over_binder(l) {
+                None => return Ok(i.under_multiple_binders(&shift_map)),
+                Some(newvar) => var = newvar,
+            };
             if let CtxItem::Kept(_, _) = i {
                 *shift_map.entry(l.clone()).or_insert(0) += 1;
             }
@@ -67,11 +66,11 @@ impl<T> Context<T> {
     /// that the passed variable always makes sense in the context of the passed item.
     /// Once we pass the variable definition, the variable doesn't make sense anymore so we just
     /// copy the remaining items.
-    fn do_with_var(
+    fn do_with_var<E>(
         &self,
         var: &AlphaVar,
-        mut f: impl FnMut(&AlphaVar, &CtxItem<T>) -> CtxItem<T>,
-    ) -> Self
+        mut f: impl FnMut(&AlphaVar, &CtxItem<T>) -> Result<CtxItem<T>, E>,
+    ) -> Result<Self, E>
     where
         T: Clone,
     {
@@ -80,32 +79,32 @@ impl<T> Context<T> {
         let mut var = var.clone();
         let mut iter = self.0.iter().rev();
         for (l, i) in iter.by_ref() {
-            vec.push((l.clone(), f(&var, i)));
+            vec.push((l.clone(), f(&var, i)?));
             if let CtxItem::Kept(_, _) = i {
-                if var == l.into() {
-                    break;
-                } else {
-                    var = var.shift(-1, &l.into());
-                }
+                match var.over_binder(l) {
+                    None => break,
+                    Some(newvar) => var = newvar,
+                };
             }
         }
         for (l, i) in iter {
             vec.push((l.clone(), (*i).clone()));
         }
         vec.reverse();
-        Context(Rc::new(vec))
+        Ok(Context(Rc::new(vec)))
     }
-    fn shift(&self, delta: isize, var: &AlphaVar) -> Self
+    fn shift(&self, delta: isize, var: &AlphaVar) -> Option<Self>
     where
         T: Clone + Shift,
     {
-        self.do_with_var(var, |var, i| i.shift(delta, &var))
+        Some(self.do_with_var(var, |var, i| Ok(i.shift(delta, &var)?))?)
     }
     fn subst_shift(&self, var: &AlphaVar, val: &Typed) -> Self
     where
         T: Clone + Subst<Typed>,
     {
-        self.do_with_var(var, |var, i| i.subst_shift(&var, val))
+        self.do_with_var(var, |var, i| Ok::<_, !>(i.subst_shift(&var, val)))
+            .unwrap()
     }
 }
 
@@ -158,27 +157,27 @@ impl TypecheckContext {
 }
 
 impl<T: Shift> Shift for CtxItem<T> {
-    fn shift(&self, delta: isize, var: &AlphaVar) -> Self {
-        match self {
+    fn shift(&self, delta: isize, var: &AlphaVar) -> Option<Self> {
+        Some(match self {
             CtxItem::Kept(v, t) => {
-                CtxItem::Kept(v.shift(delta, var), t.shift(delta, var))
+                CtxItem::Kept(v.shift(delta, var)?, t.shift(delta, var)?)
             }
             CtxItem::Replaced(e, t) => {
-                CtxItem::Replaced(e.shift(delta, var), t.shift(delta, var))
+                CtxItem::Replaced(e.shift(delta, var)?, t.shift(delta, var)?)
             }
-        }
+        })
     }
 }
 
 impl<T: Clone + Shift> Shift for Context<T> {
-    fn shift(&self, delta: isize, var: &AlphaVar) -> Self {
+    fn shift(&self, delta: isize, var: &AlphaVar) -> Option<Self> {
         self.shift(delta, var)
     }
 }
 
 impl Shift for NormalizationContext {
-    fn shift(&self, delta: isize, var: &AlphaVar) -> Self {
-        NormalizationContext(self.0.shift(delta, var))
+    fn shift(&self, delta: isize, var: &AlphaVar) -> Option<Self> {
+        Some(NormalizationContext(self.0.shift(delta, var)?))
     }
 }
 
@@ -189,12 +188,12 @@ impl<T: Subst<Typed>> Subst<Typed> for CtxItem<T> {
                 e.subst_shift(var, val),
                 t.subst_shift(var, val),
             ),
-            CtxItem::Kept(v, t) if v == var => {
-                CtxItem::Replaced(val.to_thunk(), t.subst_shift(var, val))
-            }
-            CtxItem::Kept(v, t) => {
-                CtxItem::Kept(v.shift(-1, var), t.subst_shift(var, val))
-            }
+            CtxItem::Kept(v, t) => match v.shift(-1, var) {
+                None => {
+                    CtxItem::Replaced(val.to_thunk(), t.subst_shift(var, val))
+                }
+                Some(newvar) => CtxItem::Kept(newvar, t.subst_shift(var, val)),
+            },
         }
     }
 }
diff --git a/dhall/src/core/thunk.rs b/dhall/src/core/thunk.rs
index a02d7ae..51922e1 100644
--- a/dhall/src/core/thunk.rs
+++ b/dhall/src/core/thunk.rs
@@ -264,38 +264,38 @@ impl TypeThunk {
 }
 
 impl Shift for Thunk {
-    fn shift(&self, delta: isize, var: &AlphaVar) -> Self {
-        self.0.borrow().shift(delta, var).into_thunk()
+    fn shift(&self, delta: isize, var: &AlphaVar) -> Option<Self> {
+        Some(self.0.borrow().shift(delta, var)?.into_thunk())
     }
 }
 
 impl Shift for ThunkInternal {
-    fn shift(&self, delta: isize, var: &AlphaVar) -> Self {
-        match self {
+    fn shift(&self, delta: isize, var: &AlphaVar) -> Option<Self> {
+        Some(match self {
             ThunkInternal::Unnormalized(ctx, e) => {
-                ThunkInternal::Unnormalized(ctx.shift(delta, var), e.clone())
+                ThunkInternal::Unnormalized(ctx.shift(delta, var)?, e.clone())
             }
             ThunkInternal::PartialExpr(e) => ThunkInternal::PartialExpr(
-                e.map_ref_with_special_handling_of_binders(
-                    |v| v.shift(delta, var),
-                    |x, v| v.shift(delta, &var.shift(1, &x.into())),
-                    X::clone,
-                    Label::clone,
-                ),
+                e.traverse_ref_with_special_handling_of_binders(
+                    |v| Ok(v.shift(delta, var)?),
+                    |x, v| Ok(v.shift(delta, &var.under_binder(x))?),
+                    |x| Ok(X::clone(x)),
+                    |l| Ok(Label::clone(l)),
+                )?,
             ),
             ThunkInternal::Value(m, v) => {
-                ThunkInternal::Value(*m, v.shift(delta, var))
+                ThunkInternal::Value(*m, v.shift(delta, var)?)
             }
-        }
+        })
     }
 }
 
 impl Shift for TypeThunk {
-    fn shift(&self, delta: isize, var: &AlphaVar) -> Self {
-        match self {
-            TypeThunk::Thunk(th) => TypeThunk::Thunk(th.shift(delta, var)),
-            TypeThunk::Type(t) => TypeThunk::Type(t.shift(delta, var)),
-        }
+    fn shift(&self, delta: isize, var: &AlphaVar) -> Option<Self> {
+        Some(match self {
+            TypeThunk::Thunk(th) => TypeThunk::Thunk(th.shift(delta, var)?),
+            TypeThunk::Type(t) => TypeThunk::Type(t.shift(delta, var)?),
+        })
     }
 }
 
@@ -317,8 +317,8 @@ impl Subst<Typed> for ThunkInternal {
                     |v| v.subst_shift(var, val),
                     |x, v| {
                         v.subst_shift(
-                            &var.shift(1, &x.into()),
-                            &val.shift(1, &x.into()),
+                            &var.under_binder(x),
+                            &val.under_binder(x),
                         )
                     },
                     X::clone,
diff --git a/dhall/src/core/value.rs b/dhall/src/core/value.rs
index 8fa24c7..e91b6bc 100644
--- a/dhall/src/core/value.rs
+++ b/dhall/src/core/value.rs
@@ -98,14 +98,14 @@ impl Value {
             }
             Value::ListConsClosure(a, None) => {
                 // Avoid accidental capture of the new `x` variable
-                let a1 = a.shift(1, &Label::from("x").into());
+                let a1 = a.under_binder(Label::from("x"));
                 let a1 = a1.normalize_to_expr_maybe_alpha(alpha);
                 let a = a.normalize_to_expr_maybe_alpha(alpha);
                 dhall::subexpr!(λ(x : a) -> λ(xs : List a1) -> [ x ] # xs)
             }
             Value::ListConsClosure(n, Some(v)) => {
                 // Avoid accidental capture of the new `xs` variable
-                let v = v.shift(1, &Label::from("xs").into());
+                let v = v.under_binder(Label::from("xs"));
                 let v = v.normalize_to_expr_maybe_alpha(alpha);
                 let a = n.normalize_to_expr_maybe_alpha(alpha);
                 dhall::subexpr!(λ(xs : List a) -> [ v ] # xs)
@@ -322,102 +322,121 @@ impl Value {
 }
 
 impl Shift for Value {
-    fn shift(&self, delta: isize, var: &AlphaVar) -> Self {
-        match self {
+    fn shift(&self, delta: isize, var: &AlphaVar) -> Option<Self> {
+        Some(match self {
             Value::Lam(x, t, e) => Value::Lam(
                 x.clone(),
-                t.shift(delta, var),
-                e.shift(delta, &var.shift(1, &x.into())),
+                t.shift(delta, var)?,
+                e.shift(delta, &var.under_binder(x))?,
             ),
             Value::AppliedBuiltin(b, args) => Value::AppliedBuiltin(
                 *b,
-                args.iter().map(|v| v.shift(delta, var)).collect(),
+                args.iter()
+                    .map(|v| Ok(v.shift(delta, var)?))
+                    .collect::<Result<_, _>>()?,
             ),
             Value::NaturalSuccClosure => Value::NaturalSuccClosure,
             Value::OptionalSomeClosure(tth) => {
-                Value::OptionalSomeClosure(tth.shift(delta, var))
+                Value::OptionalSomeClosure(tth.shift(delta, var)?)
             }
             Value::ListConsClosure(t, v) => Value::ListConsClosure(
-                t.shift(delta, var),
-                v.as_ref().map(|v| v.shift(delta, var)),
+                t.shift(delta, var)?,
+                v.as_ref().map(|v| Ok(v.shift(delta, var)?)).transpose()?,
             ),
             Value::Pi(x, t, e) => Value::Pi(
                 x.clone(),
-                t.shift(delta, var),
-                e.shift(delta, &var.shift(1, &x.into())),
+                t.shift(delta, var)?,
+                e.shift(delta, &var.under_binder(x))?,
             ),
-            Value::Var(v) => Value::Var(v.shift(delta, var)),
+            Value::Var(v) => Value::Var(v.shift(delta, var)?),
             Value::Const(c) => Value::Const(*c),
             Value::BoolLit(b) => Value::BoolLit(*b),
             Value::NaturalLit(n) => Value::NaturalLit(*n),
             Value::IntegerLit(n) => Value::IntegerLit(*n),
             Value::DoubleLit(n) => Value::DoubleLit(*n),
             Value::EmptyOptionalLit(tth) => {
-                Value::EmptyOptionalLit(tth.shift(delta, var))
+                Value::EmptyOptionalLit(tth.shift(delta, var)?)
             }
             Value::NEOptionalLit(th) => {
-                Value::NEOptionalLit(th.shift(delta, var))
+                Value::NEOptionalLit(th.shift(delta, var)?)
             }
             Value::EmptyListLit(tth) => {
-                Value::EmptyListLit(tth.shift(delta, var))
+                Value::EmptyListLit(tth.shift(delta, var)?)
             }
             Value::NEListLit(elts) => Value::NEListLit(
-                elts.iter().map(|v| v.shift(delta, var)).collect(),
+                elts.iter()
+                    .map(|v| Ok(v.shift(delta, var)?))
+                    .collect::<Result<_, _>>()?,
             ),
             Value::RecordLit(kvs) => Value::RecordLit(
                 kvs.iter()
-                    .map(|(k, v)| (k.clone(), v.shift(delta, var)))
-                    .collect(),
+                    .map(|(k, v)| Ok((k.clone(), v.shift(delta, var)?)))
+                    .collect::<Result<_, _>>()?,
             ),
             Value::RecordType(kvs) => Value::RecordType(
                 kvs.iter()
-                    .map(|(k, v)| (k.clone(), v.shift(delta, var)))
-                    .collect(),
+                    .map(|(k, v)| Ok((k.clone(), v.shift(delta, var)?)))
+                    .collect::<Result<_, _>>()?,
             ),
             Value::UnionType(kts) => Value::UnionType(
                 kts.iter()
                     .map(|(k, v)| {
-                        (k.clone(), v.as_ref().map(|v| v.shift(delta, var)))
+                        Ok((
+                            k.clone(),
+                            v.as_ref()
+                                .map(|v| Ok(v.shift(delta, var)?))
+                                .transpose()?,
+                        ))
                     })
-                    .collect(),
+                    .collect::<Result<_, _>>()?,
             ),
             Value::UnionConstructor(x, kts) => Value::UnionConstructor(
                 x.clone(),
                 kts.iter()
                     .map(|(k, v)| {
-                        (k.clone(), v.as_ref().map(|v| v.shift(delta, var)))
+                        Ok((
+                            k.clone(),
+                            v.as_ref()
+                                .map(|v| Ok(v.shift(delta, var)?))
+                                .transpose()?,
+                        ))
                     })
-                    .collect(),
+                    .collect::<Result<_, _>>()?,
             ),
             Value::UnionLit(x, v, kts) => Value::UnionLit(
                 x.clone(),
-                v.shift(delta, var),
+                v.shift(delta, var)?,
                 kts.iter()
                     .map(|(k, v)| {
-                        (k.clone(), v.as_ref().map(|v| v.shift(delta, var)))
+                        Ok((
+                            k.clone(),
+                            v.as_ref()
+                                .map(|v| Ok(v.shift(delta, var)?))
+                                .transpose()?,
+                        ))
                     })
-                    .collect(),
+                    .collect::<Result<_, _>>()?,
             ),
             Value::TextLit(elts) => Value::TextLit(
                 elts.iter()
                     .map(|contents| {
                         use InterpolatedTextContents::{Expr, Text};
-                        match contents {
-                            Expr(th) => Expr(th.shift(delta, var)),
+                        Ok(match contents {
+                            Expr(th) => Expr(th.shift(delta, var)?),
                             Text(s) => Text(s.clone()),
-                        }
+                        })
                     })
-                    .collect(),
+                    .collect::<Result<_, _>>()?,
             ),
-            Value::PartialExpr(e) => {
-                Value::PartialExpr(e.map_ref_with_special_handling_of_binders(
-                    |v| v.shift(delta, var),
-                    |x, v| v.shift(delta, &var.shift(1, &x.into())),
-                    X::clone,
-                    Label::clone,
-                ))
-            }
-        }
+            Value::PartialExpr(e) => Value::PartialExpr(
+                e.traverse_ref_with_special_handling_of_binders(
+                    |v| Ok(v.shift(delta, var)?),
+                    |x, v| Ok(v.shift(delta, &var.under_binder(x))?),
+                    |x| Ok(X::clone(x)),
+                    |l| Ok(Label::clone(l)),
+                )?,
+            ),
+        })
     }
 }
 
@@ -435,8 +454,8 @@ impl Subst<Typed> for Value {
                     |v| v.subst_shift(var, val),
                     |x, v| {
                         v.subst_shift(
-                            &var.shift(1, &x.into()),
-                            &val.shift(1, &x.into()),
+                            &var.under_binder(x),
+                            &val.under_binder(x),
                         )
                     },
                     X::clone,
@@ -456,10 +475,7 @@ impl Subst<Typed> for Value {
             Value::Lam(x, t, e) => Value::Lam(
                 x.clone(),
                 t.subst_shift(var, val),
-                e.subst_shift(
-                    &var.shift(1, &x.into()),
-                    &val.shift(1, &x.into()),
-                ),
+                e.subst_shift(&var.under_binder(x), &val.under_binder(x)),
             ),
             Value::NaturalSuccClosure => Value::NaturalSuccClosure,
             Value::OptionalSomeClosure(tth) => {
@@ -472,13 +488,12 @@ impl Subst<Typed> for Value {
             Value::Pi(x, t, e) => Value::Pi(
                 x.clone(),
                 t.subst_shift(var, val),
-                e.subst_shift(
-                    &var.shift(1, &x.into()),
-                    &val.shift(1, &x.into()),
-                ),
+                e.subst_shift(&var.under_binder(x), &val.under_binder(x)),
             ),
-            Value::Var(v) if v == var => val.to_value().clone(),
-            Value::Var(v) => Value::Var(v.shift(-1, var)),
+            Value::Var(v) => match v.shift(-1, var) {
+                None => val.to_value().clone(),
+                Some(newvar) => Value::Var(newvar),
+            },
             Value::Const(c) => Value::Const(*c),
             Value::BoolLit(b) => Value::BoolLit(*b),
             Value::NaturalLit(n) => Value::NaturalLit(*n),
diff --git a/dhall/src/core/var.rs b/dhall/src/core/var.rs
index ba92172..ea7e55f 100644
--- a/dhall/src/core/var.rs
+++ b/dhall/src/core/var.rs
@@ -17,19 +17,34 @@ pub(crate) struct AlphaVar {
 pub(crate) struct AlphaLabel(Label);
 
 pub(crate) trait Shift: Sized {
-    fn shift(&self, delta: isize, var: &AlphaVar) -> Self;
+    // Shift an expression to move it around binders without changing the meaning of its free
+    // variables. Shift by 1 to move an expression under a binder. Shift by -1 to extract an
+    // expression from under a binder, if the expression does not refer to that bound variable.
+    // Returns None if delta was negative and the variable was free in the expression.
+    fn shift(&self, delta: isize, var: &AlphaVar) -> Option<Self>;
 
-    fn shift0(&self, delta: isize, x: &Label) -> Self {
-        self.shift(delta, &x.into())
+    fn over_binder<T>(&self, x: T) -> Option<Self>
+    where
+        T: Into<AlphaVar>,
+    {
+        self.shift(-1, &x.into())
     }
 
-    fn shift0_multiple(&self, shift_map: &HashMap<Label, isize>) -> Self
+    fn under_binder<T>(&self, x: T) -> Self
+    where
+        T: Into<AlphaVar>,
+    {
+        // Can't fail since delta is positive
+        self.shift(1, &x.into()).unwrap()
+    }
+
+    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 {
-            v = v.shift0(*n, x);
+            v = v.shift(*n as isize, &x.into()).unwrap();
         }
         v
     }
@@ -68,20 +83,20 @@ impl AlphaLabel {
 }
 
 impl Shift for AlphaVar {
-    fn shift(&self, delta: isize, var: &AlphaVar) -> Self {
-        AlphaVar {
-            normal: self.normal.shift(delta, &var.normal),
+    fn shift(&self, delta: isize, var: &AlphaVar) -> Option<Self> {
+        Some(AlphaVar {
+            normal: self.normal.shift(delta, &var.normal)?,
             alpha: match (&self.alpha, &var.alpha) {
-                (Some(x), Some(v)) => Some(x.shift(delta, v)),
+                (Some(x), Some(v)) => Some(x.shift(delta, v)?),
                 _ => None,
             },
-        }
+        })
     }
 }
 
 impl Shift for () {
-    fn shift(&self, _delta: isize, _var: &AlphaVar) -> Self {
-        ()
+    fn shift(&self, _delta: isize, _var: &AlphaVar) -> Option<Self> {
+        Some(())
     }
 }