1 files changed, 34 insertions, 23 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())
     }