Reimplement basic tck/nze with proper environments

Inspired from dhall_haskell

2 files changed, 401 insertions, 0 deletions

diff --git a/dhall/src/semantics/nze/mod.rs b/dhall/src/semantics/nze/mod.rs
new file mode 100644
index 0000000..ef8918f
--- /dev/null
+++ b/dhall/src/semantics/nze/mod.rs
@@ -0,0 +1,2 @@
+pub mod nzexpr;
+// pub(crate) use nzexpr::*;
diff --git a/dhall/src/semantics/nze/nzexpr.rs b/dhall/src/semantics/nze/nzexpr.rs
new file mode 100644
index 0000000..2f27dcb
--- /dev/null
+++ b/dhall/src/semantics/nze/nzexpr.rs
@@ -0,0 +1,399 @@
+#![allow(dead_code)]
+use crate::semantics::core::var::AlphaVar;
+use crate::semantics::phase::typecheck::rc;
+use crate::semantics::phase::{Normalized, NormalizedExpr, ToExprOptions};
+use crate::syntax::{Expr, ExprKind, Label, V};
+
+pub(crate) type Type = NzExpr;
+#[derive(Debug, Clone)]
+pub(crate) struct TypeError;
+pub(crate) type Binder = Label;
+
+// An expression with inferred types at every node and resolved variables.
+#[derive(Debug, Clone)]
+pub(crate) struct TyExpr {
+    kind: Box<TyExprKind>,
+    ty: Option<Type>,
+}
+
+#[derive(Debug, Clone)]
+pub(crate) enum TyExprKind {
+    Var(AlphaVar),
+    Expr(ExprKind<TyExpr, Normalized>),
+}
+
+#[derive(Debug, Clone)]
+pub(crate) struct NzExpr {
+    kind: Box<NzExprKind>,
+    ty: Option<Box<Type>>,
+}
+
+#[derive(Debug, Clone)]
+pub(crate) enum NzExprKind {
+    Var {
+        var: NzVar,
+    },
+    Lam {
+        binder: Binder,
+        annot: NzExpr,
+        env: TyEnv,
+        body: TyExpr,
+    },
+    Pi {
+        binder: Binder,
+        annot: NzExpr,
+        env: TyEnv,
+        body: TyExpr,
+    },
+    Expr(ExprKind<NzExpr, Normalized>),
+    // Unevaled {
+    //     env: TyEnv,
+    //     expr: TyExprKind,
+    // },
+}
+
+#[derive(Debug, Clone)]
+enum TyEnvItem {
+    // Variable is bound with given type
+    Kept(Type),
+    // Variable has been replaced by corresponding value
+    Replaced(NzExpr),
+}
+
+#[derive(Debug, Clone)]
+pub(crate) struct TyEnv {
+    env: Vec<(Binder, TyEnvItem)>,
+}
+
+#[derive(Debug, Clone, Copy)]
+pub(crate) struct QuoteEnv {
+    size: usize,
+}
+
+// Reverse-debruijn index: counts number of binders from the bottom of the stack.
+#[derive(Debug, Clone, Copy)]
+pub(crate) struct NzVar {
+    idx: usize,
+}
+
+impl TyEnv {
+    pub fn new() -> Self {
+        TyEnv { env: Vec::new() }
+    }
+    pub fn to_quoteenv(&self) -> QuoteEnv {
+        QuoteEnv {
+            size: self.env.len(),
+        }
+    }
+    fn with_vec(&self, vec: Vec<(Binder, TyEnvItem)>) -> Self {
+        TyEnv { env: vec }
+    }
+
+    pub fn insert_type(&self, x: &Binder, t: NzExpr) -> Self {
+        let mut vec = self.env.clone();
+        vec.push((x.clone(), TyEnvItem::Kept(t)));
+        self.with_vec(vec)
+    }
+    pub fn insert_value(&self, x: &Binder, e: NzExpr) -> Self {
+        let mut vec = self.env.clone();
+        vec.push((x.clone(), TyEnvItem::Replaced(e)));
+        self.with_vec(vec)
+    }
+    pub fn lookup_var(&self, var: &V<Binder>) -> Option<TyExpr> {
+        let V(name, idx) = var;
+        let (idx, (_, item)) = self
+            .env
+            .iter()
+            .rev()
+            .enumerate()
+            .filter(|(_, (x, _))| x == name)
+            .nth(*idx)?;
+        let ty = match item {
+            TyEnvItem::Kept(ty) => ty.clone(),
+            TyEnvItem::Replaced(e) => e.get_type().unwrap(),
+        };
+        Some(TyExpr::new(
+            TyExprKind::Var(AlphaVar::new(V((), idx))),
+            Some(ty),
+        ))
+    }
+    pub fn lookup_val(&self, var: &AlphaVar) -> NzExpr {
+        let idx = self.env.len() - 1 - var.idx();
+        match &self.env[idx].1 {
+            TyEnvItem::Kept(ty) => NzExpr::new(
+                NzExprKind::Var { var: NzVar { idx } },
+                Some(ty.clone()),
+            ),
+            TyEnvItem::Replaced(x) => x.clone(),
+        }
+    }
+    pub fn size(&self) -> usize {
+        self.env.len()
+    }
+}
+
+impl QuoteEnv {
+    pub fn new() -> Self {
+        QuoteEnv { size: 0 }
+    }
+    pub fn insert(&self) -> Self {
+        QuoteEnv {
+            size: self.size + 1,
+        }
+    }
+    pub fn lookup(&self, var: &NzVar) -> AlphaVar {
+        AlphaVar::new(V((), self.size - var.idx - 1))
+    }
+}
+
+impl TyExpr {
+    pub fn new(kind: TyExprKind, ty: Option<Type>) -> Self {
+        TyExpr {
+            kind: Box::new(kind),
+            ty,
+        }
+    }
+
+    pub fn kind(&self) -> &TyExprKind {
+        self.kind.as_ref()
+    }
+    pub fn get_type(&self) -> Result<Type, TypeError> {
+        match &self.ty {
+            Some(t) => Ok(t.clone()),
+            None => Err(TypeError),
+        }
+    }
+
+    /// Converts a value back to the corresponding AST expression.
+    pub fn to_expr(&self, opts: ToExprOptions) -> NormalizedExpr {
+        fn tyexpr_to_expr<'a>(
+            tyexpr: &'a TyExpr,
+            opts: ToExprOptions,
+            env: &mut Vec<&'a Binder>,
+        ) -> NormalizedExpr {
+            rc(match tyexpr.kind() {
+                TyExprKind::Var(var) if opts.alpha => {
+                    ExprKind::Var(V("_".into(), var.idx()))
+                }
+                TyExprKind::Var(var) => {
+                    let name = env[env.len() - 1 - var.idx()];
+                    let idx = env
+                        .iter()
+                        .rev()
+                        .take(var.idx())
+                        .filter(|l| **l == name)
+                        .count();
+                    ExprKind::Var(V(name.clone(), idx))
+                }
+                TyExprKind::Expr(e) => {
+                    let e = e.map_ref_maybe_binder(|l, tye| {
+                        if let Some(l) = l {
+                            env.push(l);
+                        }
+                        let e = tyexpr_to_expr(tye, opts, env);
+                        if let Some(_) = l {
+                            env.pop();
+                        }
+                        e
+                    });
+
+                    match e {
+                        ExprKind::Lam(_, t, e) if opts.alpha => {
+                            ExprKind::Lam("_".into(), t, e)
+                        }
+                        ExprKind::Pi(_, t, e) if opts.alpha => {
+                            ExprKind::Pi("_".into(), t, e)
+                        }
+                        e => e,
+                    }
+                }
+            })
+        }
+
+        tyexpr_to_expr(self, opts, &mut Vec::new())
+    }
+
+    pub fn normalize_nf(&self, env: &TyEnv) -> NzExpr {
+        let kind = match self.kind() {
+            TyExprKind::Var(var) => return env.lookup_val(var),
+            TyExprKind::Expr(ExprKind::Lam(binder, annot, body)) => {
+                NzExprKind::Lam {
+                    binder: binder.clone(),
+                    annot: annot.normalize_nf(env),
+                    env: env.clone(),
+                    body: body.clone(),
+                }
+            }
+            TyExprKind::Expr(ExprKind::Pi(binder, annot, body)) => {
+                NzExprKind::Pi {
+                    binder: binder.clone(),
+                    annot: annot.normalize_nf(env),
+                    env: env.clone(),
+                    body: body.clone(),
+                }
+            }
+            TyExprKind::Expr(e) => {
+                let e = e.map_ref(|tye| tye.normalize_nf(env));
+                match e {
+                    ExprKind::App(f, arg) => match f.kind() {
+                        NzExprKind::Lam {
+                            binder, env, body, ..
+                        } => {
+                            return body
+                                .normalize_nf(&env.insert_value(binder, arg))
+                        }
+                        _ => NzExprKind::Expr(ExprKind::App(f, arg)),
+                    },
+                    _ => NzExprKind::Expr(e),
+                }
+            }
+        };
+        let ty = self.ty.clone();
+        NzExpr::new(kind, ty)
+    }
+
+    pub fn normalize(&self) -> NzExpr {
+        self.normalize_nf(&TyEnv::new())
+    }
+}
+
+impl NzExpr {
+    pub fn new(kind: NzExprKind, ty: Option<Type>) -> Self {
+        NzExpr {
+            kind: Box::new(kind),
+            ty: ty.map(Box::new),
+        }
+    }
+
+    pub fn kind(&self) -> &NzExprKind {
+        self.kind.as_ref()
+    }
+    pub fn get_type(&self) -> Result<Type, TypeError> {
+        match &self.ty {
+            Some(t) => Ok(t.as_ref().clone()),
+            None => Err(TypeError),
+        }
+    }
+
+    pub fn quote(&self, qenv: QuoteEnv) -> TyExpr {
+        let ty = self.ty.as_ref().map(|t| (**t).clone());
+        let kind = match self.kind.as_ref() {
+            NzExprKind::Var { var } => TyExprKind::Var(qenv.lookup(var)),
+            NzExprKind::Lam {
+                binder,
+                annot,
+                env,
+                body,
+            } => TyExprKind::Expr(ExprKind::Lam(
+                binder.clone(),
+                annot.quote(qenv),
+                body.normalize_nf(&env.insert_type(binder, annot.clone()))
+                    .quote(qenv.insert()),
+            )),
+            NzExprKind::Pi {
+                binder,
+                annot,
+                env,
+                body,
+            } => TyExprKind::Expr(ExprKind::Pi(
+                binder.clone(),
+                annot.quote(qenv),
+                body.normalize_nf(&env.insert_type(binder, annot.clone()))
+                    .quote(qenv.insert()),
+            )),
+            NzExprKind::Expr(e) => {
+                TyExprKind::Expr(e.map_ref(|nze| nze.quote(qenv)))
+            } // NzExprKind::Unevaled { env, expr } => {
+              //     return TyExpr::new(expr.clone(), ty.clone())
+              //         .normalize_nf(env)
+              //         .quote(qenv)
+              // }
+        };
+        TyExpr::new(kind, ty)
+    }
+    pub fn to_expr(&self) -> NormalizedExpr {
+        self.quote(QuoteEnv::new()).to_expr(ToExprOptions {
+            alpha: false,
+            normalize: false,
+        })
+    }
+}
+
+/// Type-check an expression and return the expression alongside its type if type-checking
+/// succeeded, or an error if type-checking failed.
+fn type_with(env: &TyEnv, e: Expr<Normalized>) -> Result<TyExpr, TypeError> {
+    match e.as_ref() {
+        ExprKind::Var(var) => match env.lookup_var(&var) {
+            Some(x) => Ok(x),
+            None => Err(TypeError),
+        },
+        ExprKind::Lam(binder, annot, body) => {
+            let annot = type_with(env, annot.clone())?;
+            let annot_nf = annot.normalize_nf(env);
+            let body = type_with(
+                &env.insert_type(&binder, annot_nf.clone()),
+                body.clone(),
+            )?;
+            let ty = NzExpr::new(
+                NzExprKind::Pi {
+                    binder: binder.clone(),
+                    annot: annot_nf,
+                    env: env.clone(),
+                    body: body.get_type()?.quote(env.to_quoteenv()),
+                },
+                None, // TODO: function_check
+            );
+            Ok(TyExpr::new(
+                TyExprKind::Expr(ExprKind::Lam(binder.clone(), annot, body)),
+                Some(ty),
+            ))
+        }
+        ExprKind::Pi(binder, annot, body) => {
+            let annot = type_with(env, annot.clone())?;
+            let annot_nf = annot.normalize_nf(env);
+            let body = type_with(
+                &env.insert_type(&binder, annot_nf.clone()),
+                body.clone(),
+            )?;
+            Ok(TyExpr::new(
+                TyExprKind::Expr(ExprKind::Pi(binder.clone(), annot, body)),
+                None, // TODO: function_check
+            ))
+        }
+        ExprKind::App(f, arg) => {
+            let f = type_with(env, f.clone())?;
+            let arg = type_with(env, arg.clone())?;
+            let tf = f.get_type()?;
+            let (binder, _expected_arg_ty, body_env, body_ty) = match tf.kind()
+            {
+                NzExprKind::Pi {
+                    binder,
+                    annot,
+                    env,
+                    body,
+                } => (binder, annot, env, body),
+                _ => return Err(TypeError),
+            };
+            // if arg.get_type()? != *expected_arg_ty {
+            //     return Err(TypeError);
+            // }
+
+            let arg_nf = arg.normalize_nf(env);
+            let ty =
+                body_ty.normalize_nf(&body_env.insert_value(&binder, arg_nf));
+
+            Ok(TyExpr::new(
+                TyExprKind::Expr(ExprKind::App(f, arg)),
+                Some(ty),
+            ))
+        }
+        e => Ok(TyExpr::new(
+            TyExprKind::Expr(e.traverse_ref(|e| type_with(env, e.clone()))?),
+            None,
+        )),
+    }
+}
+
+pub(crate) fn typecheck(e: Expr<Normalized>) -> Result<TyExpr, TypeError> {
+    type_with(&TyEnv::new(), e)
+}