Implement basic env-based normalization for Value-based TyExpr

5 files changed, 271 insertions, 109 deletions

diff --git a/dhall/src/semantics/core/value.rs b/dhall/src/semantics/core/value.rs
index a7c207d..8beed24 100644
--- a/dhall/src/semantics/core/value.rs
+++ b/dhall/src/semantics/core/value.rs
@@ -5,7 +5,7 @@ use std::rc::Rc;
 use crate::error::{TypeError, TypeMessage};
 use crate::semantics::core::var::{AlphaVar, Binder};
 use crate::semantics::nze::{NzVar, QuoteEnv};
-use crate::semantics::phase::normalize::{apply_any, normalize_whnf};
+use crate::semantics::phase::normalize::{apply_any, normalize_whnf, NzEnv};
 use crate::semantics::phase::typecheck::{builtin_to_value, const_to_value};
 use crate::semantics::phase::{
     Normalized, NormalizedExpr, ToExprOptions, Typed,
@@ -53,11 +53,27 @@ pub(crate) enum Form {
     NF,
 }
 
+#[derive(Debug, Clone)]
+pub(crate) struct Closure {
+    env: NzEnv,
+    body: TyExpr,
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub(crate) enum ValueKind<Value> {
     /// Closures
     Lam(Binder, Value, Value),
     Pi(Binder, Value, Value),
+    LamClosure {
+        binder: Binder,
+        annot: Value,
+        closure: Closure,
+    },
+    PiClosure {
+        binder: Binder,
+        annot: Value,
+        closure: Closure,
+    },
     // Invariant: in whnf, the evaluation must not be able to progress further.
     AppliedBuiltin(Builtin, Vec<Value>),
 
@@ -267,128 +283,136 @@ impl 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()
-    //         }
-    //     }
-    // }
 
     pub fn to_tyexpr(&self, qenv: QuoteEnv) -> TyExpr {
-        let self_kind = self.as_kind();
-        let self_ty = self.as_internal().ty.clone();
-        let self_span = self.as_internal().span.clone();
-        if let ValueKind::Var(v, _w) = &*self_kind {
-            return TyExpr::new(
+        let tye = match &*self.as_kind() {
+            ValueKind::Var(v, _w) => {
+                TyExprKind::Var(*v)
                 // TODO: Only works when we don't normalize under lambdas
-                // TyExprKind::Var(qenv.lookup(w)),
-                TyExprKind::Var(*v),
-                self_ty,
-                self_span,
-            );
-        }
-
-        // TODO: properly recover intermediate types; `None` is a time-bomb here.
-        let self_kind = self_kind.map_ref_with_special_handling_of_binders(
-            |v| v.to_tyexpr(qenv),
-            |_, _, v| v.to_tyexpr(qenv.insert()),
-        );
-        let expr = match self_kind {
-            ValueKind::Var(..) => unreachable!(),
-            ValueKind::Lam(x, t, e) => ExprKind::Lam(x.to_label(), t, e),
-            ValueKind::Pi(x, t, e) => ExprKind::Pi(x.to_label(), t, e),
-            ValueKind::AppliedBuiltin(b, args) => {
-                args.into_iter().fold(ExprKind::Builtin(b), |acc, v| {
-                    ExprKind::App(
-                        TyExpr::new(
-                            TyExprKind::Expr(acc),
-                            None, // TODO
-                            Span::Artificial,
-                        ),
-                        v,
-                    )
-                })
+                // TyExprKind::Var(qenv.lookup(w))
             }
-            ValueKind::Const(c) => ExprKind::Const(c),
-            ValueKind::BoolLit(b) => ExprKind::BoolLit(b),
-            ValueKind::NaturalLit(n) => ExprKind::NaturalLit(n),
-            ValueKind::IntegerLit(n) => ExprKind::IntegerLit(n),
-            ValueKind::DoubleLit(n) => ExprKind::DoubleLit(n),
-            ValueKind::EmptyOptionalLit(n) => ExprKind::App(
-                builtin_to_value(Builtin::OptionalNone).to_tyexpr(qenv),
-                n,
-            ),
-            ValueKind::NEOptionalLit(n) => ExprKind::SomeLit(n),
-            ValueKind::EmptyListLit(n) => ExprKind::EmptyListLit(TyExpr::new(
-                TyExprKind::Expr(ExprKind::App(
-                    builtin_to_value(Builtin::List).to_tyexpr(qenv),
-                    n,
-                )),
-                Some(const_to_value(Const::Type)),
-                Span::Artificial,
+            ValueKind::LamClosure {
+                binder,
+                annot,
+                closure,
+            } => TyExprKind::Expr(ExprKind::Lam(
+                binder.to_label(),
+                annot.to_tyexpr(qenv),
+                closure.apply_ty(annot.clone()).to_tyexpr(qenv.insert()),
             )),
-            ValueKind::NEListLit(elts) => ExprKind::NEListLit(elts),
-            ValueKind::RecordLit(kvs) => {
-                ExprKind::RecordLit(kvs.into_iter().collect())
-            }
-            ValueKind::RecordType(kts) => {
-                ExprKind::RecordType(kts.into_iter().collect())
-            }
-            ValueKind::UnionType(kts) => {
-                ExprKind::UnionType(kts.into_iter().collect())
-            }
-            ValueKind::UnionConstructor(l, kts, t) => ExprKind::Field(
-                TyExpr::new(
-                    TyExprKind::Expr(ExprKind::UnionType(
-                        kts.into_iter().collect(),
-                    )),
-                    Some(t),
-                    Span::Artificial,
-                ),
-                l,
-            ),
-            ValueKind::UnionLit(l, v, kts, uniont, ctort) => ExprKind::App(
-                TyExpr::new(
-                    TyExprKind::Expr(ExprKind::Field(
+            ValueKind::PiClosure {
+                binder,
+                annot,
+                closure,
+            } => TyExprKind::Expr(ExprKind::Pi(
+                binder.to_label(),
+                annot.to_tyexpr(qenv),
+                closure.apply_ty(annot.clone()).to_tyexpr(qenv.insert()),
+            )),
+            self_kind => {
+                // TODO: properly recover intermediate types; `None` is a time-bomb here.
+                let self_kind = self_kind
+                    .map_ref_with_special_handling_of_binders(
+                        |v| v.to_tyexpr(qenv),
+                        |_, _, v| v.to_tyexpr(qenv.insert()),
+                    );
+                let expr = match self_kind {
+                    ValueKind::Var(..)
+                    | ValueKind::LamClosure { .. }
+                    | ValueKind::PiClosure { .. } => unreachable!(),
+                    ValueKind::Lam(x, t, e) => {
+                        ExprKind::Lam(x.to_label(), t, e)
+                    }
+                    ValueKind::Pi(x, t, e) => ExprKind::Pi(x.to_label(), t, e),
+                    ValueKind::AppliedBuiltin(b, args) => {
+                        args.into_iter().fold(ExprKind::Builtin(b), |acc, v| {
+                            ExprKind::App(
+                                TyExpr::new(
+                                    TyExprKind::Expr(acc),
+                                    None, // TODO
+                                    Span::Artificial,
+                                ),
+                                v,
+                            )
+                        })
+                    }
+                    ValueKind::Const(c) => ExprKind::Const(c),
+                    ValueKind::BoolLit(b) => ExprKind::BoolLit(b),
+                    ValueKind::NaturalLit(n) => ExprKind::NaturalLit(n),
+                    ValueKind::IntegerLit(n) => ExprKind::IntegerLit(n),
+                    ValueKind::DoubleLit(n) => ExprKind::DoubleLit(n),
+                    ValueKind::EmptyOptionalLit(n) => ExprKind::App(
+                        builtin_to_value(Builtin::OptionalNone).to_tyexpr(qenv),
+                        n,
+                    ),
+                    ValueKind::NEOptionalLit(n) => ExprKind::SomeLit(n),
+                    ValueKind::EmptyListLit(n) => {
+                        ExprKind::EmptyListLit(TyExpr::new(
+                            TyExprKind::Expr(ExprKind::App(
+                                builtin_to_value(Builtin::List).to_tyexpr(qenv),
+                                n,
+                            )),
+                            Some(const_to_value(Const::Type)),
+                            Span::Artificial,
+                        ))
+                    }
+                    ValueKind::NEListLit(elts) => ExprKind::NEListLit(elts),
+                    ValueKind::RecordLit(kvs) => {
+                        ExprKind::RecordLit(kvs.into_iter().collect())
+                    }
+                    ValueKind::RecordType(kts) => {
+                        ExprKind::RecordType(kts.into_iter().collect())
+                    }
+                    ValueKind::UnionType(kts) => {
+                        ExprKind::UnionType(kts.into_iter().collect())
+                    }
+                    ValueKind::UnionConstructor(l, kts, t) => ExprKind::Field(
                         TyExpr::new(
                             TyExprKind::Expr(ExprKind::UnionType(
                                 kts.into_iter().collect(),
                             )),
-                            Some(uniont),
+                            Some(t),
                             Span::Artificial,
                         ),
                         l,
-                    )),
-                    Some(ctort),
-                    Span::Artificial,
-                ),
-                v,
-            ),
-            ValueKind::TextLit(elts) => {
-                ExprKind::TextLit(elts.into_iter().collect())
-            }
-            ValueKind::Equivalence(x, y) => {
-                ExprKind::BinOp(BinOp::Equivalence, x, y)
+                    ),
+                    ValueKind::UnionLit(l, v, kts, uniont, ctort) => {
+                        ExprKind::App(
+                            TyExpr::new(
+                                TyExprKind::Expr(ExprKind::Field(
+                                    TyExpr::new(
+                                        TyExprKind::Expr(ExprKind::UnionType(
+                                            kts.into_iter().collect(),
+                                        )),
+                                        Some(uniont),
+                                        Span::Artificial,
+                                    ),
+                                    l,
+                                )),
+                                Some(ctort),
+                                Span::Artificial,
+                            ),
+                            v,
+                        )
+                    }
+                    ValueKind::TextLit(elts) => {
+                        ExprKind::TextLit(elts.into_iter().collect())
+                    }
+                    ValueKind::Equivalence(x, y) => {
+                        ExprKind::BinOp(BinOp::Equivalence, x, y)
+                    }
+                    ValueKind::PartialExpr(e) => e,
+                };
+                TyExprKind::Expr(expr)
             }
-            ValueKind::PartialExpr(e) => e,
         };
 
-        TyExpr::new(TyExprKind::Expr(expr), self_ty, self_span)
+        let self_ty = self.as_internal().ty.clone();
+        let self_span = self.as_internal().span.clone();
+        TyExpr::new(tye, self_ty, self_span)
+    }
+    pub fn to_tyexpr_noenv(&self) -> TyExpr {
+        self.to_tyexpr(QuoteEnv::new())
     }
 }
 
@@ -502,6 +526,10 @@ impl ValueKind<Value> {
                 t.normalize_mut();
                 e.normalize_mut();
             }
+            ValueKind::LamClosure { annot, .. }
+            | ValueKind::PiClosure { annot, .. } => {
+                annot.normalize_mut();
+            }
             ValueKind::AppliedBuiltin(_, args) => {
                 for x in args.iter_mut() {
                     x.normalize_mut();
@@ -620,6 +648,21 @@ impl<V> ValueKind<V> {
     }
 }
 
+impl Closure {
+    pub fn new(env: &NzEnv, body: TyExpr) -> Self {
+        Closure {
+            env: env.clone(),
+            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))
+    }
+}
+
 // TODO: use Rc comparison to shortcut on identical pointers
 impl std::cmp::PartialEq for Value {
     fn eq(&self, other: &Self) -> bool {
@@ -628,6 +671,14 @@ 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")
+    }
+}
+impl std::cmp::Eq for Closure {}
+
 impl std::fmt::Debug for Value {
     fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let vint: &ValueInternal = &self.as_internal();
diff --git a/dhall/src/semantics/core/visitor.rs b/dhall/src/semantics/core/visitor.rs
index 5a04747..ee44ed7 100644
--- a/dhall/src/semantics/core/visitor.rs
+++ b/dhall/src/semantics/core/visitor.rs
@@ -80,6 +80,24 @@ where
             let (t, e) = v.visit_binder(l, t, e)?;
             Pi(l.clone(), t, e)
         }
+        LamClosure {
+            binder,
+            annot,
+            closure,
+        } => LamClosure {
+            binder: binder.clone(),
+            annot: v.visit_val(annot)?,
+            closure: closure.clone(),
+        },
+        PiClosure {
+            binder,
+            annot,
+            closure,
+        } => PiClosure {
+            binder: binder.clone(),
+            annot: v.visit_val(annot)?,
+            closure: closure.clone(),
+        },
         AppliedBuiltin(b, xs) => AppliedBuiltin(*b, v.visit_vec(xs)?),
         Var(v, w) => Var(v.clone(), *w),
         Const(k) => Const(*k),
diff --git a/dhall/src/semantics/nze/nzexpr.rs b/dhall/src/semantics/nze/nzexpr.rs
index 1256ea0..723c895 100644
--- a/dhall/src/semantics/nze/nzexpr.rs
+++ b/dhall/src/semantics/nze/nzexpr.rs
@@ -200,6 +200,7 @@ impl NzEnv {
         env
     }
     pub fn lookup_val(&self, var: &AlphaVar) -> NzExpr {
+        // TODO: use VarEnv
         let idx = self.items.len() - 1 - var.idx();
         match &self.items[idx] {
             NzEnvItem::Kept(ty) => NzExpr::new(
@@ -211,6 +212,7 @@ impl NzEnv {
     }
 }
 
+// TODO: rename to VarEnv
 impl QuoteEnv {
     pub fn new() -> Self {
         QuoteEnv { size: 0 }
diff --git a/dhall/src/semantics/phase/normalize.rs b/dhall/src/semantics/phase/normalize.rs
index f0a6a8c..7fd76df 100644
--- a/dhall/src/semantics/phase/normalize.rs
+++ b/dhall/src/semantics/phase/normalize.rs
@@ -1,9 +1,12 @@
 use std::collections::HashMap;
 use std::convert::TryInto;
 
+use crate::semantics::nze::NzVar;
 use crate::semantics::phase::typecheck::{rc, typecheck};
 use crate::semantics::phase::Normalized;
-use crate::semantics::{AlphaVar, Value, ValueKind};
+use crate::semantics::{
+    AlphaVar, Binder, Closure, TyExpr, TyExprKind, Value, ValueKind,
+};
 use crate::syntax;
 use crate::syntax::Const::Type;
 use crate::syntax::{
@@ -786,3 +789,83 @@ pub(crate) fn normalize_whnf(
         v => v,
     }
 }
+
+#[derive(Debug, Clone)]
+enum NzEnvItem {
+    // Variable is bound with given type
+    Kept(Value),
+    // Variable has been replaced by corresponding value
+    Replaced(Value),
+}
+
+#[derive(Debug, Clone)]
+pub(crate) struct NzEnv {
+    items: Vec<NzEnvItem>,
+}
+
+impl NzEnv {
+    pub fn new() -> Self {
+        NzEnv { items: Vec::new() }
+    }
+
+    pub fn insert_type(&self, t: Value) -> Self {
+        let mut env = self.clone();
+        env.items.push(NzEnvItem::Kept(t));
+        env
+    }
+    pub fn insert_value(&self, e: Value) -> Self {
+        let mut env = self.clone();
+        env.items.push(NzEnvItem::Replaced(e));
+        env
+    }
+    pub fn lookup_val(&self, var: &AlphaVar) -> Value {
+        let idx = self.items.len() - 1 - var.idx();
+        match &self.items[idx] {
+            NzEnvItem::Kept(ty) => Value::from_kind_and_type_whnf(
+                ValueKind::Var(var.clone(), NzVar::new(idx)),
+                ty.clone(),
+            ),
+            NzEnvItem::Replaced(x) => x.clone(),
+        }
+    }
+}
+
+/// Normalize a TyExpr into WHNF
+pub(crate) fn normalize_tyexpr_whnf(e: &TyExpr, env: &NzEnv) -> Value {
+    let kind = match e.kind() {
+        TyExprKind::Var(var) => return env.lookup_val(var),
+        TyExprKind::Expr(ExprKind::Lam(binder, annot, body)) => {
+            ValueKind::LamClosure {
+                binder: Binder::new(binder.clone()),
+                annot: annot.normalize_whnf(env),
+                closure: Closure::new(env, body.clone()),
+            }
+        }
+        TyExprKind::Expr(ExprKind::Pi(binder, annot, body)) => {
+            ValueKind::PiClosure {
+                binder: Binder::new(binder.clone()),
+                annot: annot.normalize_whnf(env),
+                closure: Closure::new(env, body.clone()),
+            }
+        }
+        TyExprKind::Expr(e) => {
+            let e = e.map_ref(|tye| tye.normalize_whnf(env));
+            match e {
+                ExprKind::App(f, arg) => {
+                    let f_borrow = f.as_whnf();
+                    match &*f_borrow {
+                        ValueKind::LamClosure { closure, .. } => {
+                            return closure.apply(arg)
+                        }
+                        _ => {
+                            drop(f_borrow);
+                            ValueKind::PartialExpr(ExprKind::App(f, arg))
+                        }
+                    }
+                }
+                _ => ValueKind::PartialExpr(e),
+            }
+        }
+    };
+    Value::from_kind_and_type_whnf(kind, e.get_type().unwrap())
+}
diff --git a/dhall/src/semantics/tck/tyexpr.rs b/dhall/src/semantics/tck/tyexpr.rs
index 983f3f8..c8be3a3 100644
--- a/dhall/src/semantics/tck/tyexpr.rs
+++ b/dhall/src/semantics/tck/tyexpr.rs
@@ -1,6 +1,7 @@
 #![allow(dead_code)]
 use crate::error::{TypeError, TypeMessage};
 use crate::semantics::core::var::AlphaVar;
+use crate::semantics::phase::normalize::{normalize_tyexpr_whnf, NzEnv};
 use crate::semantics::phase::typecheck::rc;
 use crate::semantics::phase::Normalized;
 use crate::semantics::phase::{NormalizedExpr, ToExprOptions};
@@ -51,6 +52,13 @@ impl TyExpr {
     pub fn to_value(&self) -> Value {
         todo!()
     }
+
+    pub fn normalize_whnf(&self, env: &NzEnv) -> Value {
+        normalize_tyexpr_whnf(self, env)
+    }
+    pub fn normalize_whnf_noenv(&self) -> Value {
+        normalize_tyexpr_whnf(self, &NzEnv::new())
+    }
 }
 
 fn tyexpr_to_expr<'a>(