Enable comparing Closures for equality

3 files changed, 20 insertions, 14 deletions

diff --git a/dhall/src/semantics/core/value.rs b/dhall/src/semantics/core/value.rs
index 11d8bd8..a19cb02 100644
--- a/dhall/src/semantics/core/value.rs
+++ b/dhall/src/semantics/core/value.rs
@@ -55,6 +55,7 @@ pub(crate) enum Form {
 
 #[derive(Debug, Clone)]
 pub(crate) struct Closure {
+    arg_ty: Value,
     env: NzEnv,
     body: TyExpr,
 }
@@ -302,7 +303,7 @@ impl Value {
             } => TyExprKind::Expr(ExprKind::Lam(
                 binder.to_label(),
                 annot.to_tyexpr(qenv),
-                closure.apply_ty(annot.clone()).to_tyexpr(qenv.insert()),
+                closure.normalize().to_tyexpr(qenv.insert()),
             )),
             ValueKind::PiClosure {
                 binder,
@@ -311,7 +312,7 @@ impl Value {
             } => TyExprKind::Expr(ExprKind::Pi(
                 binder.to_label(),
                 annot.to_tyexpr(qenv),
-                closure.apply_ty(annot.clone()).to_tyexpr(qenv.insert()),
+                closure.normalize().to_tyexpr(qenv.insert()),
             )),
             ValueKind::AppliedBuiltin(b, args, types) => {
                 TyExprKind::Expr(args.into_iter().zip(types.into_iter()).fold(
@@ -681,17 +682,19 @@ impl<V> ValueKind<V> {
 }
 
 impl Closure {
-    pub fn new(env: &NzEnv, body: TyExpr) -> Self {
+    pub fn new(arg_ty: Value, env: &NzEnv, body: TyExpr) -> Self {
         Closure {
+            arg_ty,
             env: env.clone(),
-            body: body,
+            body,
         }
     }
     pub fn apply(&self, val: Value) -> Value {
         self.body.normalize_whnf(&self.env.insert_value(val))
     }
-    pub fn apply_ty(&self, ty: Value) -> Value {
-        self.body.normalize_whnf(&self.env.insert_type(ty))
+    pub fn normalize(&self) -> Value {
+        self.body
+            .normalize_whnf(&self.env.insert_type(self.arg_ty.clone()))
     }
 }
 
@@ -703,10 +706,9 @@ impl std::cmp::PartialEq for Value {
 }
 impl std::cmp::Eq for Value {}
 
-// TODO: panics
 impl std::cmp::PartialEq for Closure {
-    fn eq(&self, _other: &Self) -> bool {
-        panic!("comparing closures for equality seems like a bad idea")
+    fn eq(&self, other: &Self) -> bool {
+        self.normalize() == other.normalize()
     }
 }
 impl std::cmp::Eq for Closure {}
diff --git a/dhall/src/semantics/phase/normalize.rs b/dhall/src/semantics/phase/normalize.rs
index 066a004..7f547d7 100644
--- a/dhall/src/semantics/phase/normalize.rs
+++ b/dhall/src/semantics/phase/normalize.rs
@@ -884,17 +884,20 @@ pub(crate) fn normalize_tyexpr_whnf(tye: &TyExpr, env: &NzEnv) -> Value {
     let kind = match tye.kind() {
         TyExprKind::Var(var) => return env.lookup_val(var),
         TyExprKind::Expr(ExprKind::Lam(binder, annot, body)) => {
+            let annot = annot.normalize_whnf(env);
             ValueKind::LamClosure {
                 binder: Binder::new(binder.clone()),
-                annot: annot.normalize_whnf(env),
-                closure: Closure::new(env, body.clone()),
+                annot: annot.clone(),
+                closure: Closure::new(annot, env, body.clone()),
             }
         }
         TyExprKind::Expr(ExprKind::Pi(binder, annot, body)) => {
+            let annot = annot.normalize_whnf(env);
+            let closure = Closure::new(annot.clone(), env, body.clone());
             ValueKind::PiClosure {
                 binder: Binder::new(binder.clone()),
-                annot: annot.normalize_whnf(env),
-                closure: Closure::new(env, body.clone()),
+                annot,
+                closure,
             }
         }
         TyExprKind::Expr(ExprKind::Let(_, None, val, body)) => {
diff --git a/dhall/src/semantics/tck/typecheck.rs b/dhall/src/semantics/tck/typecheck.rs
index 046c999..52c7ac7 100644
--- a/dhall/src/semantics/tck/typecheck.rs
+++ b/dhall/src/semantics/tck/typecheck.rs
@@ -163,8 +163,9 @@ fn type_with(
             let ty = Value::from_kind_and_type(
                 ValueKind::PiClosure {
                     binder: Binder::new(binder.clone()),
-                    annot: annot_nf,
+                    annot: annot_nf.clone(),
                     closure: Closure::new(
+                        annot_nf,
                         env.as_nzenv(),
                         body.get_type()?.to_tyexpr(env.as_quoteenv().insert()),
                     ),