Fix all sorts of variable shenanigans

4 files changed, 68 insertions, 117 deletions

diff --git a/dhall/src/semantics/core/value.rs b/dhall/src/semantics/core/value.rs
index 35913cf..42da653 100644
--- a/dhall/src/semantics/core/value.rs
+++ b/dhall/src/semantics/core/value.rs
@@ -303,7 +303,9 @@ impl Value {
             } => TyExprKind::Expr(ExprKind::Lam(
                 binder.to_label(),
                 annot.to_tyexpr(qenv),
-                closure.normalize().to_tyexpr(qenv.insert()),
+                closure
+                    .apply_var(NzVar::new(qenv.size()))
+                    .to_tyexpr(qenv.insert()),
             )),
             ValueKind::PiClosure {
                 binder,
@@ -312,7 +314,9 @@ impl Value {
             } => TyExprKind::Expr(ExprKind::Pi(
                 binder.to_label(),
                 annot.to_tyexpr(qenv),
-                closure.normalize().to_tyexpr(qenv.insert()),
+                closure
+                    .apply_var(NzVar::new(qenv.size()))
+                    .to_tyexpr(qenv.insert()),
             )),
             ValueKind::AppliedBuiltin(b, args, types) => {
                 TyExprKind::Expr(args.into_iter().zip(types.into_iter()).fold(
@@ -692,120 +696,17 @@ impl Closure {
     pub fn apply(&self, val: Value) -> Value {
         self.body.normalize_whnf(&self.env.insert_value(val))
     }
-    pub fn apply_fresh(&self, env: QuoteEnv) -> Value {
+    pub fn apply_var(&self, var: NzVar) -> Value {
         let val = Value::from_kind_and_type(
-            ValueKind::Var(AlphaVar::default(), NzVar::new(env.size())),
+            ValueKind::Var(AlphaVar::default(), var),
             self.arg_ty.clone(),
         );
         self.apply(val)
     }
-    pub fn normalize(&self) -> Value {
-        self.body
-            .normalize_whnf(&self.env.insert_type(self.arg_ty.clone()))
-    }
 }
 
 /// Compare two values for equality modulo alpha/beta-equivalence.
 // TODO: use Rc comparison to shortcut on identical pointers
-fn equiv(val1: &Value, val2: &Value) -> bool {
-    struct ValueWithEnv<'v> {
-        val: &'v Value,
-        env: QuoteEnv,
-    }
-    impl<'v> PartialEq for ValueWithEnv<'v> {
-        fn eq(&self, other: &ValueWithEnv<'v>) -> bool {
-            equiv_with_env(self.env, self.val, other.env, other.val)
-        }
-    }
-    // Push the given context into every subnode of the ValueKind. That way, normal equality of the
-    // resulting value will take into account the context.
-    fn push_context<'v>(
-        env: QuoteEnv,
-        kind: &'v ValueKind<Value>,
-    ) -> ValueKind<ValueWithEnv<'v>> {
-        kind.map_ref_with_special_handling_of_binders(
-            |val| ValueWithEnv { val, env },
-            |_, _, val| ValueWithEnv {
-                val,
-                env: env.insert(),
-            },
-        )
-    }
-
-    fn equiv_with_env<'v>(
-        env1: QuoteEnv,
-        val1: &'v Value,
-        env2: QuoteEnv,
-        val2: &'v Value,
-    ) -> bool {
-        use ValueKind::Var;
-        let kind1 = val1.as_whnf();
-        let kind2 = val2.as_whnf();
-        match (&*kind1, &*kind2) {
-            (Var(_, v1), Var(_, v2)) => {
-                v1 == v2
-                // match (env1.lookup(v1), env2.lookup(v2)) {
-                //     // Both vars were found in the environment, check they point to the same
-                //     // binder.
-                //     // Does that even make any sense with an incomplete environment ?
-                //     // I have no clue what I'm doing anymore.
-                //     (Some(i), Some(j)) => i == j,
-                //     // Both vars point to outside the environment ????
-                //     // (None, None) => v1 == v2,
-                //     _ => false,
-                // }
-            }
-            // (Var(_, v1), Var(_, v2)) => env1.lookup(v1) == env2.lookup(v2),
-            (
-                ValueKind::LamClosure {
-                    annot: a1,
-                    closure: cl1,
-                    ..
-                },
-                ValueKind::LamClosure {
-                    annot: a2,
-                    closure: cl2,
-                    ..
-                },
-            ) => {
-                equiv_with_env(env1, a1, env2, a2)
-                    && equiv_with_env(
-                        env1.insert(),
-                        &cl1.apply_fresh(env1),
-                        env2.insert(),
-                        &cl2.apply_fresh(env2),
-                    )
-            }
-            (
-                ValueKind::PiClosure {
-                    annot: a1,
-                    closure: cl1,
-                    ..
-                },
-                ValueKind::PiClosure {
-                    annot: a2,
-                    closure: cl2,
-                    ..
-                },
-            ) => {
-                equiv_with_env(env1, a1, env2, a2)
-                    && equiv_with_env(
-                        env1.insert(),
-                        &cl1.apply_fresh(env1),
-                        env2.insert(),
-                        &cl2.apply_fresh(env2),
-                    )
-            }
-            (k1, k2) => push_context(env1, k1) == push_context(env2, k2),
-        }
-    }
-
-    // TODO: need to use ambiant env instead of creating new one
-    // Might be possible to generate free variables differently to avoid this
-    equiv_with_env(QuoteEnv::new(), val1, QuoteEnv::new(), val2)
-}
-
-// TODO: use Rc comparison to shortcut on identical pointers
 impl std::cmp::PartialEq for Value {
     fn eq(&self, other: &Self) -> bool {
         *self.as_whnf() == *other.as_whnf()
@@ -815,7 +716,8 @@ impl std::cmp::Eq for Value {}
 
 impl std::cmp::PartialEq for Closure {
     fn eq(&self, other: &Self) -> bool {
-        self.normalize() == other.normalize()
+        let v = NzVar::fresh();
+        self.apply_var(v) == other.apply_var(v)
     }
 }
 impl std::cmp::Eq for Closure {}
diff --git a/dhall/src/semantics/nze/nzexpr.rs b/dhall/src/semantics/nze/nzexpr.rs
index 33fdde3..6559082 100644
--- a/dhall/src/semantics/nze/nzexpr.rs
+++ b/dhall/src/semantics/nze/nzexpr.rs
@@ -81,10 +81,12 @@ pub(crate) struct QuoteEnv {
     size: usize,
 }
 
-// Reverse-debruijn index: counts number of binders from the bottom of the stack.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub(crate) struct NzVar {
-    idx: usize,
+pub(crate) enum NzVar {
+    /// Reverse-debruijn index: counts number of binders from the bottom of the stack.
+    Bound(usize),
+    /// Fake fresh variable generated for expression equality checking.
+    Fresh(usize),
 }
 // TODO: temporary hopefully
 // impl std::cmp::PartialEq for NzVar {
@@ -207,7 +209,9 @@ impl NzEnv {
         let idx = self.items.len() - 1 - var.idx();
         match &self.items[idx] {
             NzEnvItem::Kept(ty) => NzExpr::new(
-                NzExprKind::Var { var: NzVar { idx } },
+                NzExprKind::Var {
+                    var: NzVar::new(idx),
+                },
                 Some(ty.clone()),
             ),
             NzEnvItem::Replaced(x) => x.clone(),
@@ -235,18 +239,32 @@ impl QuoteEnv {
         self.lookup_fallible(var).unwrap()
     }
     pub fn lookup_fallible(&self, var: &NzVar) -> Option<AlphaVar> {
-        let idx = self.size.checked_sub(var.idx + 1)?;
+        let idx = self.size.checked_sub(var.idx() + 1)?;
         Some(AlphaVar::new(V((), idx)))
     }
 }
 
 impl NzVar {
     pub fn new(idx: usize) -> Self {
-        NzVar { idx }
+        NzVar::Bound(idx)
+    }
+    pub fn fresh() -> Self {
+        use std::sync::atomic::{AtomicUsize, Ordering};
+        // Global counter to ensure uniqueness of the generated id.
+        static FRESH_VAR_COUNTER: AtomicUsize = AtomicUsize::new(0);
+        let id = FRESH_VAR_COUNTER.fetch_add(1, Ordering::SeqCst);
+        NzVar::Fresh(id)
     }
     pub fn shift(&self, delta: isize) -> Self {
-        NzVar {
-            idx: (self.idx as isize + delta) as usize,
+        NzVar::new((self.idx() as isize + delta) as usize)
+    }
+    // Panics on a fresh variable.
+    pub fn idx(&self) -> usize {
+        match self {
+            NzVar::Bound(i) => *i,
+            NzVar::Fresh(_) => panic!(
+                "Trying to use a fresh variable outside of equality checking"
+            ),
         }
     }
 }
diff --git a/dhall/src/semantics/tck/tyexpr.rs b/dhall/src/semantics/tck/tyexpr.rs
index a42265d..9e8dc47 100644
--- a/dhall/src/semantics/tck/tyexpr.rs
+++ b/dhall/src/semantics/tck/tyexpr.rs
@@ -12,7 +12,7 @@ use crate::syntax::{ExprKind, Label, Span, V};
 pub(crate) type Type = Value;
 
 // An expression with inferred types at every node and resolved variables.
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 pub(crate) struct TyExpr {
     kind: Box<TyExprKind>,
     ty: Option<Type>,
@@ -114,3 +114,16 @@ fn tyexpr_to_expr<'a>(
         }
     })
 }
+
+impl std::fmt::Debug for TyExpr {
+    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let mut x = fmt.debug_struct("TyExpr");
+        x.field("kind", self.kind());
+        if let Some(ty) = self.ty.as_ref() {
+            x.field("type", &ty);
+        } else {
+            x.field("type", &None::<()>);
+        }
+        x.finish()
+    }
+}
diff --git a/dhall/src/tests.rs b/dhall/src/tests.rs
index 88c09cc..971c48d 100644
--- a/dhall/src/tests.rs
+++ b/dhall/src/tests.rs
@@ -164,6 +164,7 @@ pub fn run_test(test: Test<'_>) -> Result<()> {
             // let expr = parse_file_str(&expr_file_path)?.resolve()?.to_expr();
             // let tyexpr = crate::semantics::nze::nzexpr::typecheck(expr)?;
             // let ty = tyexpr.get_type()?.to_expr();
+            //
             let expr = parse_file_str(&expr_file_path)?.resolve()?.to_expr();
             let ty = crate::semantics::tck::typecheck::typecheck(&expr)?
                 .get_type()?
@@ -173,6 +174,23 @@ pub fn run_test(test: Test<'_>) -> Result<()> {
                 });
             let expected = parse_file_str(&expected_file_path)?.to_expr();
             assert_eq_display!(ty, expected);
+            //
+            // let expr = parse_file_str(&expr_file_path)?.resolve()?.to_expr();
+            // let ty = crate::semantics::tck::typecheck::typecheck(&expr)?
+            //     .get_type()?;
+            // let expected = parse_file_str(&expected_file_path)?.to_expr();
+            // let expected = crate::semantics::tck::typecheck::typecheck(&expected)?
+            //     .normalize_whnf_noenv();
+            // // if ty != expected {
+            // //     assert_eq_display!(ty.to_expr(crate::semantics::phase::ToExprOptions {
+            // //         alpha: false,
+            // //         normalize: true,
+            // //     }), expected.to_expr(crate::semantics::phase::ToExprOptions {
+            // //         alpha: false,
+            // //         normalize: true,
+            // //     }))
+            // // }
+            // assert_eq_pretty!(ty, expected);
         }
         TypeInferenceFailure(file_path) => {
             let mut res =