Merge pull request #131 from Nadrieril/hir

Decouple main expression types

5 files changed, 509 insertions, 564 deletions

diff --git a/dhall/src/semantics/tck/env.rs b/dhall/src/semantics/tck/env.rs
index 1fc711c..17b3cfe 100644
--- a/dhall/src/semantics/tck/env.rs
+++ b/dhall/src/semantics/tck/env.rs
@@ -1,5 +1,5 @@
-use crate::semantics::{AlphaVar, NzEnv, NzVar, Type, Value};
-use crate::syntax::{Label, V};
+use crate::semantics::{AlphaVar, NameEnv, Nir, NzEnv, NzVar, Type, ValEnv};
+use crate::syntax::Label;
 
 /// Environment for indexing variables.
 #[derive(Debug, Clone, Copy)]
@@ -7,23 +7,20 @@ pub(crate) struct VarEnv {
     size: usize,
 }
 
-/// Environment for resolving names.
-#[derive(Debug, Clone)]
-pub(crate) struct NameEnv {
-    names: Vec<Label>,
-}
-
 /// Environment for typing expressions.
 #[derive(Debug, Clone)]
 pub(crate) struct TyEnv {
     names: NameEnv,
-    items: NzEnv,
+    items: ValEnv<Type>,
 }
 
 impl VarEnv {
     pub fn new() -> Self {
         VarEnv { size: 0 }
     }
+    pub fn from_size(size: usize) -> Self {
+        VarEnv { size }
+    }
     pub fn size(&self) -> usize {
         self.size
     }
@@ -41,84 +38,42 @@ impl VarEnv {
     }
 }
 
-impl NameEnv {
-    pub fn new() -> Self {
-        NameEnv { names: Vec::new() }
-    }
-    pub fn as_varenv(&self) -> VarEnv {
-        VarEnv {
-            size: self.names.len(),
-        }
-    }
-
-    pub fn insert(&self, x: &Label) -> Self {
-        let mut env = self.clone();
-        env.insert_mut(x);
-        env
-    }
-    pub fn insert_mut(&mut self, x: &Label) {
-        self.names.push(x.clone())
-    }
-    pub fn remove_mut(&mut self) {
-        self.names.pop();
-    }
-
-    pub fn unlabel_var(&self, var: &V) -> Option<AlphaVar> {
-        let V(name, idx) = var;
-        let (idx, _) = self
-            .names
-            .iter()
-            .rev()
-            .enumerate()
-            .filter(|(_, n)| *n == name)
-            .nth(*idx)?;
-        Some(AlphaVar::new(idx))
-    }
-    pub fn label_var(&self, var: &AlphaVar) -> V {
-        let name = &self.names[self.names.len() - 1 - var.idx()];
-        let idx = self
-            .names
-            .iter()
-            .rev()
-            .take(var.idx())
-            .filter(|n| *n == name)
-            .count();
-        V(name.clone(), idx)
-    }
-}
-
 impl TyEnv {
     pub fn new() -> Self {
         TyEnv {
             names: NameEnv::new(),
-            items: NzEnv::new(),
+            items: ValEnv::new(),
         }
     }
     pub fn as_varenv(&self) -> VarEnv {
         self.names.as_varenv()
     }
-    pub fn as_nzenv(&self) -> &NzEnv {
-        &self.items
+    pub fn to_nzenv(&self) -> NzEnv {
+        self.items.discard_types()
     }
     pub fn as_nameenv(&self) -> &NameEnv {
         &self.names
     }
 
-    pub fn insert_type(&self, x: &Label, t: Type) -> Self {
+    pub fn insert_type(&self, x: &Label, ty: Type) -> Self {
         TyEnv {
             names: self.names.insert(x),
-            items: self.items.insert_type(t),
+            items: self.items.insert_type(ty),
         }
     }
-    pub fn insert_value(&self, x: &Label, e: Value) -> Self {
+    pub fn insert_value(&self, x: &Label, e: Nir, ty: Type) -> Self {
         TyEnv {
             names: self.names.insert(x),
-            items: self.items.insert_value(e),
+            items: self.items.insert_value(e, ty),
         }
     }
-    pub fn lookup(&self, var: &V) -> Option<(AlphaVar, Type)> {
-        let var = self.names.unlabel_var(var)?;
-        let ty = self.items.lookup_ty(&var);
-        Some((var, ty))
+    pub fn lookup(&self, var: &AlphaVar) -> Type {
+        self.items.lookup_ty(&var)
+    }
+}
+
+impl<'a> From<&'a TyEnv> for NzEnv {
+    fn from(x: &'a TyEnv) -> Self {
+        x.to_nzenv()
     }
 }
diff --git a/dhall/src/semantics/tck/mod.rs b/dhall/src/semantics/tck/mod.rs
index 1df2a48..93c8f48 100644
--- a/dhall/src/semantics/tck/mod.rs
+++ b/dhall/src/semantics/tck/mod.rs
@@ -1,6 +1,6 @@
 pub mod env;
-pub mod tyexpr;
+pub mod tir;
 pub mod typecheck;
 pub(crate) use env::*;
-pub(crate) use tyexpr::*;
+pub(crate) use tir::*;
 pub(crate) use typecheck::*;
diff --git a/dhall/src/semantics/tck/tir.rs b/dhall/src/semantics/tck/tir.rs
new file mode 100644
index 0000000..aeb7bf9
--- /dev/null
+++ b/dhall/src/semantics/tck/tir.rs
@@ -0,0 +1,175 @@
+use crate::error::{ErrorBuilder, TypeError};
+use crate::semantics::{mkerr, Hir, Nir, NirKind, NzEnv, TyEnv, VarEnv};
+use crate::syntax::{Builtin, Const, Span};
+use crate::NormalizedExpr;
+
+/// The type of a type. 0 is `Type`, 1 is `Kind`, etc...
+#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Default)]
+pub(crate) struct Universe(u8);
+
+/// An expression representing a type
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(crate) struct Type {
+    val: Nir,
+    univ: Universe,
+}
+
+/// A hir expression plus its inferred type.
+/// Stands for "Typed intermediate representation"
+#[derive(Debug, Clone)]
+pub(crate) struct Tir<'hir> {
+    hir: &'hir Hir,
+    ty: Type,
+}
+
+impl Universe {
+    pub fn from_const(c: Const) -> Self {
+        Universe(match c {
+            Const::Type => 0,
+            Const::Kind => 1,
+            Const::Sort => 2,
+        })
+    }
+    pub fn as_const(self) -> Option<Const> {
+        match self.0 {
+            0 => Some(Const::Type),
+            1 => Some(Const::Kind),
+            2 => Some(Const::Sort),
+            _ => None,
+        }
+    }
+    pub fn next(self) -> Self {
+        Universe(self.0 + 1)
+    }
+}
+
+impl Type {
+    pub fn new(val: Nir, univ: Universe) -> Self {
+        Type { val, univ }
+    }
+    /// Creates a new Type and infers its universe by re-typechecking its value.
+    /// TODO: ideally avoid this function altogether. Would need to store types in RecordType and
+    /// PiClosure.
+    pub fn new_infer_universe(
+        env: &TyEnv,
+        val: Nir,
+    ) -> Result<Self, TypeError> {
+        let c = val.to_hir(env.as_varenv()).typecheck(env)?.ty().as_const();
+        let u = match c {
+            Some(c) => c.to_universe(),
+            None => unreachable!(
+                "internal type error: this is not a type: {:?}",
+                val
+            ),
+        };
+        Ok(Type::new(val, u))
+    }
+    pub fn from_const(c: Const) -> Self {
+        Self::new(Nir::from_const(c), c.to_universe().next())
+    }
+    pub fn from_builtin(b: Builtin) -> Self {
+        use Builtin::*;
+        match b {
+            Bool | Natural | Integer | Double | Text => {}
+            _ => unreachable!("this builtin is not a type: {}", b),
+        }
+
+        Self::new(Nir::from_builtin(b), Universe::from_const(Const::Type))
+    }
+
+    /// Get the type of this type
+    pub fn ty(&self) -> Universe {
+        self.univ
+    }
+
+    pub fn to_nir(&self) -> Nir {
+        self.val.clone()
+    }
+    pub fn as_nir(&self) -> &Nir {
+        &self.val
+    }
+    pub fn into_nir(self) -> Nir {
+        self.val
+    }
+    pub fn as_const(&self) -> Option<Const> {
+        self.val.as_const()
+    }
+    pub fn kind(&self) -> &NirKind {
+        self.val.kind()
+    }
+
+    pub fn to_hir(&self, venv: VarEnv) -> Hir {
+        self.val.to_hir(venv)
+    }
+    pub fn to_expr_tyenv(&self, tyenv: &TyEnv) -> NormalizedExpr {
+        self.val.to_hir(tyenv.as_varenv()).to_expr_tyenv(tyenv)
+    }
+}
+
+impl<'hir> Tir<'hir> {
+    pub fn from_hir(hir: &'hir Hir, ty: Type) -> Self {
+        Tir { hir, ty }
+    }
+
+    pub fn span(&self) -> Span {
+        self.as_hir().span()
+    }
+    pub fn ty(&self) -> &Type {
+        &self.ty
+    }
+
+    pub fn to_hir(&self) -> Hir {
+        self.as_hir().clone()
+    }
+    pub fn as_hir(&self) -> &Hir {
+        &self.hir
+    }
+    pub fn to_expr_tyenv(&self, env: &TyEnv) -> NormalizedExpr {
+        self.as_hir().to_expr_tyenv(env)
+    }
+
+    /// Eval the Tir. It will actually get evaluated only as needed on demand.
+    pub fn eval(&self, env: impl Into<NzEnv>) -> Nir {
+        self.as_hir().eval(env.into())
+    }
+    pub fn ensure_is_type(&self, env: &TyEnv) -> Result<(), TypeError> {
+        if self.ty().as_const().is_none() {
+            return mkerr(
+                ErrorBuilder::new(format!(
+                    "Expected a type, found: `{}`",
+                    self.to_expr_tyenv(env),
+                ))
+                .span_err(
+                    self.span(),
+                    format!(
+                        "this has type: `{}`",
+                        self.ty().to_expr_tyenv(env)
+                    ),
+                )
+                .help(format!(
+                    "An expression in type position must have type `Type`, \
+                     `Kind` or `Sort`",
+                ))
+                .format(),
+            );
+        }
+        Ok(())
+    }
+    /// Evaluate to a Type.
+    pub fn eval_to_type(&self, env: &TyEnv) -> Result<Type, TypeError> {
+        self.ensure_is_type(env)?;
+        Ok(Type::new(
+            self.eval(env),
+            self.ty()
+                .as_const()
+                .expect("case handled in ensure_is_type")
+                .to_universe(),
+        ))
+    }
+}
+
+impl From<Const> for Universe {
+    fn from(x: Const) -> Universe {
+        Universe::from_const(x)
+    }
+}
diff --git a/dhall/src/semantics/tck/tyexpr.rs b/dhall/src/semantics/tck/tyexpr.rs
deleted file mode 100644
index 1886646..0000000
--- a/dhall/src/semantics/tck/tyexpr.rs
+++ /dev/null
@@ -1,137 +0,0 @@
-use crate::error::{TypeError, TypeMessage};
-use crate::semantics::{rc, NameEnv, NzEnv, TyEnv, Value};
-use crate::syntax::{ExprKind, Span, V};
-use crate::Normalized;
-use crate::{NormalizedExpr, ToExprOptions};
-
-pub(crate) type Type = Value;
-
-/// Stores an alpha-normalized variable.
-#[derive(Debug, Clone, Copy)]
-pub struct AlphaVar {
-    idx: usize,
-}
-
-#[derive(Debug, Clone)]
-pub(crate) enum TyExprKind {
-    Var(AlphaVar),
-    // Forbidden ExprKind variants: Var, Import, Embed
-    Expr(ExprKind<TyExpr, Normalized>),
-}
-
-// An expression with inferred types at every node and resolved variables.
-#[derive(Clone)]
-pub(crate) struct TyExpr {
-    kind: Box<TyExprKind>,
-    ty: Option<Type>,
-    span: Span,
-}
-
-impl AlphaVar {
-    pub(crate) fn new(idx: usize) -> Self {
-        AlphaVar { idx }
-    }
-    pub(crate) fn idx(&self) -> usize {
-        self.idx
-    }
-}
-
-impl TyExpr {
-    pub fn new(kind: TyExprKind, ty: Option<Type>, span: Span) -> Self {
-        TyExpr {
-            kind: Box::new(kind),
-            ty,
-            span,
-        }
-    }
-
-    pub fn kind(&self) -> &TyExprKind {
-        &*self.kind
-    }
-    pub fn span(&self) -> Span {
-        self.span.clone()
-    }
-    pub fn get_type(&self) -> Result<Type, TypeError> {
-        match &self.ty {
-            Some(t) => Ok(t.clone()),
-            None => Err(TypeError::new(TypeMessage::Sort)),
-        }
-    }
-
-    /// Converts a value back to the corresponding AST expression.
-    pub fn to_expr(&self, opts: ToExprOptions) -> NormalizedExpr {
-        tyexpr_to_expr(self, opts, &mut NameEnv::new())
-    }
-    pub fn to_expr_tyenv(&self, env: &TyEnv) -> NormalizedExpr {
-        let opts = ToExprOptions {
-            normalize: true,
-            alpha: false,
-        };
-        let mut env = env.as_nameenv().clone();
-        tyexpr_to_expr(self, opts, &mut env)
-    }
-
-    /// Eval the TyExpr. It will actually get evaluated only as needed on demand.
-    pub fn eval(&self, env: &NzEnv) -> Value {
-        Value::new_thunk(env, self.clone())
-    }
-    /// Eval a closed TyExpr (i.e. without free variables). It will actually get evaluated only as
-    /// needed on demand.
-    pub fn eval_closed_expr(&self) -> Value {
-        self.eval(&NzEnv::new())
-    }
-    /// Eval a closed TyExpr fully and recursively;
-    pub fn rec_eval_closed_expr(&self) -> Value {
-        let mut val = self.eval_closed_expr();
-        val.normalize_mut();
-        val
-    }
-}
-
-fn tyexpr_to_expr<'a>(
-    tyexpr: &'a TyExpr,
-    opts: ToExprOptions,
-    env: &mut NameEnv,
-) -> NormalizedExpr {
-    rc(match tyexpr.kind() {
-        TyExprKind::Var(v) if opts.alpha => {
-            ExprKind::Var(V("_".into(), v.idx()))
-        }
-        TyExprKind::Var(v) => ExprKind::Var(env.label_var(v)),
-        TyExprKind::Expr(e) => {
-            let e = e.map_ref_maybe_binder(|l, tye| {
-                if let Some(l) = l {
-                    env.insert_mut(l);
-                }
-                let e = tyexpr_to_expr(tye, opts, env);
-                if let Some(_) = l {
-                    env.remove_mut();
-                }
-                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,
-            }
-        }
-    })
-}
-
-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/semantics/tck/typecheck.rs b/dhall/src/semantics/tck/typecheck.rs
index dd9a8fa..972326b 100644
--- a/dhall/src/semantics/tck/typecheck.rs
+++ b/dhall/src/semantics/tck/typecheck.rs
@@ -5,34 +5,44 @@ use std::collections::HashMap;
 use crate::error::{ErrorBuilder, TypeError, TypeMessage};
 use crate::semantics::merge_maps;
 use crate::semantics::{
-    type_of_builtin, Binder, BuiltinClosure, Closure, TyEnv, TyExpr,
-    TyExprKind, Type, Value, ValueKind,
+    type_of_builtin, Binder, BuiltinClosure, Closure, Hir, HirKind, Nir,
+    NirKind, Tir, TyEnv, Type,
 };
 use crate::syntax::{
-    BinOp, Builtin, Const, Expr, ExprKind, InterpolatedTextContents, Span,
+    BinOp, Builtin, Const, ExprKind, InterpolatedTextContents, LitKind, Span,
 };
-use crate::Normalized;
 
-fn type_of_recordtype<'a>(
-    span: Span,
-    tys: impl Iterator<Item = Cow<'a, TyExpr>>,
-) -> Result<Type, TypeError> {
-    let span_err = |msg: &str| {
-        mkerr(
-            ErrorBuilder::new(msg.to_string())
-                .span_err(span.clone(), msg.to_string())
-                .format(),
-        )
+fn check_rectymerge(
+    span: &Span,
+    env: &TyEnv,
+    x: Nir,
+    y: Nir,
+) -> Result<(), TypeError> {
+    let kts_x = match x.kind() {
+        NirKind::RecordType(kts) => kts,
+        _ => {
+            return mk_span_err(
+                span.clone(),
+                "RecordTypeMergeRequiresRecordType",
+            )
+        }
     };
-    // An empty record type has type Type
-    let mut k = Const::Type;
-    for t in tys {
-        match t.get_type()?.as_const() {
-            Some(c) => k = max(k, c),
-            None => return span_err("InvalidFieldType"),
+    let kts_y = match y.kind() {
+        NirKind::RecordType(kts) => kts,
+        _ => {
+            return mk_span_err(
+                span.clone(),
+                "RecordTypeMergeRequiresRecordType",
+            )
+        }
+    };
+    for (k, tx) in kts_x {
+        if let Some(ty) = kts_y.get(k) {
+            // TODO: store Type in RecordType ?
+            check_rectymerge(span, env, tx.clone(), ty.clone())?;
         }
     }
-    Ok(Value::from_const(k))
+    Ok(())
 }
 
 fn function_check(a: Const, b: Const) -> Const {
@@ -43,90 +53,62 @@ fn function_check(a: Const, b: Const) -> Const {
     }
 }
 
-fn mkerr<T, S: ToString>(x: S) -> Result<T, TypeError> {
-    Err(TypeError::new(TypeMessage::Custom(x.to_string())))
+pub fn mkerr<T, S: ToString>(msg: S) -> Result<T, TypeError> {
+    Err(TypeError::new(TypeMessage::Custom(msg.to_string())))
+}
+
+pub fn mk_span_err<T, S: ToString>(span: Span, msg: S) -> Result<T, TypeError> {
+    mkerr(
+        ErrorBuilder::new(msg.to_string())
+            .span_err(span, msg.to_string())
+            .format(),
+    )
 }
 
 /// When all sub-expressions have been typed, check the remaining toplevel
 /// layer.
 fn type_one_layer(
     env: &TyEnv,
-    ekind: ExprKind<TyExpr, Normalized>,
+    ekind: ExprKind<Tir<'_>>,
     span: Span,
-) -> Result<TyExpr, TypeError> {
-    let span_err = |msg: &str| {
-        mkerr(
-            ErrorBuilder::new(msg.to_string())
-                .span_err(span.clone(), msg.to_string())
-                .format(),
-        )
-    };
+) -> Result<Type, TypeError> {
+    let span_err = |msg: &str| mk_span_err(span.clone(), msg);
 
     let ty = match &ekind {
-        ExprKind::Import(..) => unreachable!(
-            "There should remain no imports in a resolved expression"
-        ),
+        ExprKind::Import(..) | ExprKind::Completion(..) => {
+            unreachable!("This case should have been handled in resolution")
+        }
         ExprKind::Var(..)
         | ExprKind::Const(Const::Sort)
-        | ExprKind::Embed(..) => unreachable!(), // Handled in type_with
-
-        ExprKind::Lam(binder, annot, body) => {
-            let body_ty = body.get_type()?;
-            let body_ty = body_ty.to_tyexpr(env.as_varenv().insert());
-            let pi_ekind = ExprKind::Pi(binder.clone(), annot.clone(), body_ty);
-            type_one_layer(env, pi_ekind, Span::Artificial)?
-                .eval(env.as_nzenv())
+        | ExprKind::Lam(..)
+        | ExprKind::Pi(..)
+        | ExprKind::Let(..)
+        | ExprKind::Annot(..) => {
+            unreachable!("This case should have been handled in type_with")
         }
-        ExprKind::Pi(_, annot, body) => {
-            let ks = match annot.get_type()?.as_const() {
-                Some(k) => k,
-                _ => {
-                    return mkerr(
-                        ErrorBuilder::new(format!(
-                            "Invalid input type: `{}`",
-                            annot.get_type()?.to_expr_tyenv(env),
-                        ))
-                        .span_err(
-                            annot.span(),
-                            format!(
-                                "this has type: `{}`",
-                                annot.get_type()?.to_expr_tyenv(env)
-                            ),
-                        )
-                        .help(format!(
-                            "The input type of a function must have type \
-                             `Type`, `Kind` or `Sort`",
-                        ))
-                        .format(),
-                    );
-                }
-            };
-            let kt = match body.get_type()?.as_const() {
-                Some(k) => k,
-                _ => return span_err("Invalid output type"),
-            };
 
-            Value::from_const(function_check(ks, kt))
-        }
-        ExprKind::Let(_, _, _, body) => body.get_type()?,
-
-        ExprKind::Const(Const::Type) => Value::from_const(Const::Kind),
-        ExprKind::Const(Const::Kind) => Value::from_const(Const::Sort),
+        ExprKind::Const(Const::Type) => Type::from_const(Const::Kind),
+        ExprKind::Const(Const::Kind) => Type::from_const(Const::Sort),
         ExprKind::Builtin(b) => {
-            let t_expr = type_of_builtin(*b);
-            let t_tyexpr = type_with(env, &t_expr)?;
-            t_tyexpr.eval(env.as_nzenv())
-        }
-        ExprKind::BoolLit(_) => Value::from_builtin(Builtin::Bool),
-        ExprKind::NaturalLit(_) => Value::from_builtin(Builtin::Natural),
-        ExprKind::IntegerLit(_) => Value::from_builtin(Builtin::Integer),
-        ExprKind::DoubleLit(_) => Value::from_builtin(Builtin::Double),
+            let t_hir = type_of_builtin(*b);
+            typecheck(&t_hir)?.eval_to_type(env)?
+        }
+        ExprKind::Lit(LitKind::Bool(_)) => Type::from_builtin(Builtin::Bool),
+        ExprKind::Lit(LitKind::Natural(_)) => {
+            Type::from_builtin(Builtin::Natural)
+        }
+        ExprKind::Lit(LitKind::Integer(_)) => {
+            Type::from_builtin(Builtin::Integer)
+        }
+        ExprKind::Lit(LitKind::Double(_)) => {
+            Type::from_builtin(Builtin::Double)
+        }
         ExprKind::TextLit(interpolated) => {
-            let text_type = Value::from_builtin(Builtin::Text);
+            let text_type = Type::from_builtin(Builtin::Text);
             for contents in interpolated.iter() {
                 use InterpolatedTextContents::Expr;
                 if let Expr(x) = contents {
-                    if x.get_type()? != text_type {
+                    if *x.ty() != text_type {
                         return span_err("InvalidTextInterpolation");
                     }
                 }
@@ -134,9 +116,9 @@ fn type_one_layer(
             text_type
         }
         ExprKind::EmptyListLit(t) => {
-            let t = t.eval(env.as_nzenv());
-            match &*t.kind() {
-                ValueKind::AppliedBuiltin(BuiltinClosure {
+            let t = t.eval_to_type(env)?;
+            match t.kind() {
+                NirKind::AppliedBuiltin(BuiltinClosure {
                     b: Builtin::List,
                     args,
                     ..
@@ -149,49 +131,57 @@ fn type_one_layer(
             let mut iter = xs.iter();
             let x = iter.next().unwrap();
             for y in iter {
-                if x.get_type()? != y.get_type()? {
+                if x.ty() != y.ty() {
                     return span_err("InvalidListElement");
                 }
             }
-            let t = x.get_type()?;
-            if t.get_type()?.as_const() != Some(Const::Type) {
+            if x.ty().ty().as_const() != Some(Const::Type) {
                 return span_err("InvalidListType");
             }
 
-            Value::from_builtin(Builtin::List).app(t)
+            let t = x.ty().to_nir();
+            Nir::from_builtin(Builtin::List).app(t).to_type(Const::Type)
         }
         ExprKind::SomeLit(x) => {
-            let t = x.get_type()?;
-            if t.get_type()?.as_const() != Some(Const::Type) {
+            if x.ty().ty().as_const() != Some(Const::Type) {
                 return span_err("InvalidOptionalType");
             }
 
-            Value::from_builtin(Builtin::Optional).app(t)
+            let t = x.ty().to_nir();
+            Nir::from_builtin(Builtin::Optional)
+                .app(t)
+                .to_type(Const::Type)
         }
         ExprKind::RecordLit(kvs) => {
             use std::collections::hash_map::Entry;
             let mut kts = HashMap::new();
+            // An empty record type has type Type
+            let mut k = Const::Type;
             for (x, v) in kvs {
                 // Check for duplicated entries
                 match kts.entry(x.clone()) {
                     Entry::Occupied(_) => {
                         return span_err("RecordTypeDuplicateField")
                     }
-                    Entry::Vacant(e) => e.insert(v.get_type()?),
+                    Entry::Vacant(e) => e.insert(v.ty().to_nir()),
                 };
+
+                // Check that the fields have a valid kind
+                match v.ty().ty().as_const() {
+                    Some(c) => k = max(k, c),
+                    None => return span_err("InvalidFieldType"),
+                }
             }
 
-            let ty = type_of_recordtype(
-                span.clone(),
-                kts.iter()
-                    .map(|(_, t)| Cow::Owned(t.to_tyexpr(env.as_varenv()))),
-            )?;
-            Value::from_kind_and_type(ValueKind::RecordType(kts), ty)
+            Nir::from_kind(NirKind::RecordType(kts)).to_type(k)
         }
         ExprKind::RecordType(kts) => {
             use std::collections::hash_map::Entry;
             let mut seen_fields = HashMap::new();
-            for (x, _) in kts {
+            // An empty record type has type Type
+            let mut k = Const::Type;
+
+            for (x, t) in kts {
                 // Check for duplicated entries
                 match seen_fields.entry(x.clone()) {
                     Entry::Occupied(_) => {
@@ -199,12 +189,15 @@ fn type_one_layer(
                     }
                     Entry::Vacant(e) => e.insert(()),
                 };
+
+                // Check the type is a Const and compute final type
+                match t.ty().as_const() {
+                    Some(c) => k = max(k, c),
+                    None => return span_err("InvalidFieldType"),
+                }
             }
 
-            type_of_recordtype(
-                span.clone(),
-                kts.iter().map(|(_, t)| Cow::Borrowed(t)),
-            )?
+            Type::from_const(k)
         }
         ExprKind::UnionType(kts) => {
             use std::collections::hash_map::Entry;
@@ -213,7 +206,7 @@ fn type_one_layer(
             let mut k = None;
             for (x, t) in kts {
                 if let Some(t) = t {
-                    match (k, t.get_type()?.as_const()) {
+                    match (k, t.ty().as_const()) {
                         (None, Some(k2)) => k = Some(k2),
                         (Some(k1), Some(k2)) if k1 == k2 => {}
                         _ => return span_err("InvalidFieldType"),
@@ -231,80 +224,52 @@ fn type_one_layer(
             // an union type with only unary variants also has type Type
             let k = k.unwrap_or(Const::Type);
 
-            Value::from_const(k)
+            Type::from_const(k)
         }
         ExprKind::Field(scrut, x) => {
-            match &*scrut.get_type()?.kind() {
-                ValueKind::RecordType(kts) => match kts.get(&x) {
-                    Some(tth) => tth.clone(),
+            match scrut.ty().kind() {
+                NirKind::RecordType(kts) => match kts.get(&x) {
+                    Some(val) => Type::new_infer_universe(env, val.clone())?,
                     None => return span_err("MissingRecordField"),
                 },
-                // TODO: branch here only when scrut.get_type() is a Const
-                _ => {
-                    let scrut_nf = scrut.eval(env.as_nzenv());
-                    match scrut_nf.kind() {
-                        ValueKind::UnionType(kts) => match kts.get(x) {
+                NirKind::Const(_) => {
+                    let scrut = scrut.eval_to_type(env)?;
+                    match scrut.kind() {
+                        NirKind::UnionType(kts) => match kts.get(x) {
                             // Constructor has type T -> < x: T, ... >
                             Some(Some(ty)) => {
-                                // Can't fail because uniontypes must have type Const(_).
-                                let kt = scrut.get_type()?.as_const().unwrap();
-                                // The type of the field must be Const smaller than `kt`, thus the
-                                // function type has type `kt`.
-                                let pi_ty = Value::from_const(kt);
-
-                                Value::from_kind_and_type(
-                                    ValueKind::PiClosure {
-                                        binder: Binder::new(x.clone()),
-                                        annot: ty.clone(),
-                                        closure: Closure::new_constant(
-                                            scrut_nf,
-                                        ),
-                                    },
-                                    pi_ty,
-                                )
+                                Nir::from_kind(NirKind::PiClosure {
+                                    binder: Binder::new(x.clone()),
+                                    annot: ty.clone(),
+                                    closure: Closure::new_constant(
+                                        scrut.to_nir(),
+                                    ),
+                                })
+                                .to_type(scrut.ty())
                             }
-                            Some(None) => scrut_nf,
+                            Some(None) => scrut,
                             None => return span_err("MissingUnionField"),
                         },
                         _ => return span_err("NotARecord"),
                     }
-                } // _ => span_err("NotARecord"),
-            }
-        }
-        ExprKind::Annot(x, t) => {
-            let t = t.eval(env.as_nzenv());
-            let x_ty = x.get_type()?;
-            if x_ty != t {
-                return span_err(&format!(
-                    "annot mismatch: ({} : {}) : {}",
-                    x.to_expr_tyenv(env),
-                    x_ty.to_tyexpr(env.as_varenv()).to_expr_tyenv(env),
-                    t.to_tyexpr(env.as_varenv()).to_expr_tyenv(env)
-                ));
-                // return span_err(format!(
-                //     "annot mismatch: {} != {}",
-                //     x_ty.to_tyexpr(env.as_varenv()).to_expr_tyenv(env),
-                //     t.to_tyexpr(env.as_varenv()).to_expr_tyenv(env)
-                // ));
-                // return span_err(format!("annot mismatch: {:#?} : {:#?}", x, t,));
+                }
+                _ => return span_err("NotARecord"),
             }
-            x_ty
         }
         ExprKind::Assert(t) => {
-            let t = t.eval(env.as_nzenv());
-            match &*t.kind() {
-                ValueKind::Equivalence(x, y) if x == y => {}
-                ValueKind::Equivalence(..) => {
-                    return span_err("AssertMismatch")
-                }
+            let t = t.eval_to_type(env)?;
+            match t.kind() {
+                NirKind::Equivalence(x, y) if x == y => {}
+                NirKind::Equivalence(..) => return span_err("AssertMismatch"),
                 _ => return span_err("AssertMustTakeEquivalence"),
             }
             t
         }
         ExprKind::App(f, arg) => {
-            match f.get_type()?.kind() {
-                ValueKind::PiClosure { annot, closure, .. } => {
-                    if arg.get_type()? != *annot {
+            match f.ty().kind() {
+                // TODO: store Type in closure
+                NirKind::PiClosure { annot, closure, .. } => {
+                    if arg.ty().as_nir() != annot {
                         return mkerr(
                             ErrorBuilder::new(format!(
                                 "wrong type of function argument"
@@ -320,25 +285,25 @@ fn type_one_layer(
                                 arg.span(),
                                 format!(
                                     "but this has type: {}",
-                                    arg.get_type()?.to_expr_tyenv(env),
+                                    arg.ty().to_expr_tyenv(env),
                                 ),
                             )
                             .note(format!(
                                 "expected type `{}`\n   found type `{}`",
                                 annot.to_expr_tyenv(env),
-                                arg.get_type()?.to_expr_tyenv(env),
+                                arg.ty().to_expr_tyenv(env),
                             ))
                             .format(),
                         );
                     }
 
-                    let arg_nf = arg.eval(env.as_nzenv());
-                    closure.apply(arg_nf)
+                    let arg_nf = arg.eval(env);
+                    Type::new_infer_universe(env, closure.apply(arg_nf))?
                 }
                 _ => return mkerr(
                     ErrorBuilder::new(format!(
                         "expected function, found `{}`",
-                        f.get_type()?.to_expr_tyenv(env)
+                        f.ty().to_expr_tyenv(env)
                     ))
                     .span_err(
                         f.span(),
@@ -349,33 +314,30 @@ fn type_one_layer(
             }
         }
         ExprKind::BoolIf(x, y, z) => {
-            if *x.get_type()?.kind() != ValueKind::from_builtin(Builtin::Bool) {
+            if *x.ty().kind() != NirKind::from_builtin(Builtin::Bool) {
                 return span_err("InvalidPredicate");
             }
-            if y.get_type()?.get_type()?.as_const() != Some(Const::Type) {
-                return span_err("IfBranchMustBeTerm");
-            }
-            if z.get_type()?.get_type()?.as_const() != Some(Const::Type) {
+            if y.ty().ty().as_const() != Some(Const::Type) {
                 return span_err("IfBranchMustBeTerm");
             }
-            if y.get_type()? != z.get_type()? {
+            if y.ty() != z.ty() {
                 return span_err("IfBranchMismatch");
             }
 
-            y.get_type()?
+            y.ty().clone()
         }
         ExprKind::BinOp(BinOp::RightBiasedRecordMerge, x, y) => {
-            let x_type = x.get_type()?;
-            let y_type = y.get_type()?;
+            let x_type = x.ty();
+            let y_type = y.ty();
 
             // Extract the LHS record type
             let kts_x = match x_type.kind() {
-                ValueKind::RecordType(kts) => kts,
+                NirKind::RecordType(kts) => kts,
                 _ => return span_err("MustCombineRecord"),
             };
             // Extract the RHS record type
             let kts_y = match y_type.kind() {
-                ValueKind::RecordType(kts) => kts,
+                NirKind::RecordType(kts) => kts,
                 _ => return span_err("MustCombineRecord"),
             };
 
@@ -385,80 +347,59 @@ fn type_one_layer(
                 Ok(r_t.clone())
             })?;
 
-            // Construct the final record type
-            let ty = type_of_recordtype(
-                span.clone(),
-                kts.iter()
-                    .map(|(_, t)| Cow::Owned(t.to_tyexpr(env.as_varenv()))),
-            )?;
-            Value::from_kind_and_type(ValueKind::RecordType(kts), ty)
+            let u = max(x.ty().ty(), y.ty().ty());
+            Nir::from_kind(NirKind::RecordType(kts)).to_type(u)
         }
         ExprKind::BinOp(BinOp::RecursiveRecordMerge, x, y) => {
-            let ekind = ExprKind::BinOp(
-                BinOp::RecursiveRecordTypeMerge,
-                x.get_type()?.to_tyexpr(env.as_varenv()),
-                y.get_type()?.to_tyexpr(env.as_varenv()),
+            check_rectymerge(&span, env, x.ty().to_nir(), y.ty().to_nir())?;
+
+            let hir = Hir::new(
+                HirKind::Expr(ExprKind::BinOp(
+                    BinOp::RecursiveRecordTypeMerge,
+                    x.ty().to_hir(env.as_varenv()),
+                    y.ty().to_hir(env.as_varenv()),
+                )),
+                span.clone(),
             );
-            type_one_layer(env, ekind, Span::Artificial)?.eval(env.as_nzenv())
+            let x_u = x.ty().ty();
+            let y_u = y.ty().ty();
+            Type::new(hir.eval(env), max(x_u, y_u))
         }
         ExprKind::BinOp(BinOp::RecursiveRecordTypeMerge, x, y) => {
-            let x_val = x.eval(env.as_nzenv());
-            let y_val = y.eval(env.as_nzenv());
-            let kts_x = match x_val.kind() {
-                ValueKind::RecordType(kts) => kts,
-                _ => return span_err("RecordTypeMergeRequiresRecordType"),
-            };
-            let kts_y = match y_val.kind() {
-                ValueKind::RecordType(kts) => kts,
-                _ => return span_err("RecordTypeMergeRequiresRecordType"),
-            };
-            for (k, tx) in kts_x {
-                if let Some(ty) = kts_y.get(k) {
-                    type_one_layer(
-                        env,
-                        ExprKind::BinOp(
-                            BinOp::RecursiveRecordTypeMerge,
-                            tx.to_tyexpr(env.as_varenv()),
-                            ty.to_tyexpr(env.as_varenv()),
-                        ),
-                        Span::Artificial,
-                    )?;
-                }
-            }
+            check_rectymerge(&span, env, x.eval(env), y.eval(env))?;
 
             // A RecordType's type is always a const
-            let xk = x.get_type()?.as_const().unwrap();
-            let yk = y.get_type()?.as_const().unwrap();
-            Value::from_const(max(xk, yk))
+            let xk = x.ty().as_const().unwrap();
+            let yk = y.ty().as_const().unwrap();
+            Type::from_const(max(xk, yk))
         }
         ExprKind::BinOp(BinOp::ListAppend, l, r) => {
-            let l_ty = l.get_type()?;
-            match &*l_ty.kind() {
-                ValueKind::AppliedBuiltin(BuiltinClosure {
+            match l.ty().kind() {
+                NirKind::AppliedBuiltin(BuiltinClosure {
                     b: Builtin::List,
                     ..
                 }) => {}
                 _ => return span_err("BinOpTypeMismatch"),
             }
 
-            if l_ty != r.get_type()? {
+            if l.ty() != r.ty() {
                 return span_err("BinOpTypeMismatch");
             }
 
-            l_ty
+            l.ty().clone()
         }
         ExprKind::BinOp(BinOp::Equivalence, l, r) => {
-            if l.get_type()? != r.get_type()? {
+            if l.ty() != r.ty() {
                 return span_err("EquivalenceTypeMismatch");
             }
-            if l.get_type()?.get_type()?.as_const() != Some(Const::Type) {
+            if l.ty().ty().as_const() != Some(Const::Type) {
                 return span_err("EquivalenceArgumentsMustBeTerms");
             }
 
-            Value::from_const(Const::Type)
+            Type::from_const(Const::Type)
         }
         ExprKind::BinOp(o, l, r) => {
-            let t = Value::from_builtin(match o {
+            let t = Type::from_builtin(match o {
                 BinOp::BoolAnd
                 | BinOp::BoolOr
                 | BinOp::BoolEQ
@@ -473,27 +414,27 @@ fn type_one_layer(
                 BinOp::ImportAlt => unreachable!("ImportAlt leftover in tck"),
             });
 
-            if l.get_type()? != t {
+            if *l.ty() != t {
                 return span_err("BinOpTypeMismatch");
             }
 
-            if r.get_type()? != t {
+            if *r.ty() != t {
                 return span_err("BinOpTypeMismatch");
             }
 
             t
         }
         ExprKind::Merge(record, union, type_annot) => {
-            let record_type = record.get_type()?;
+            let record_type = record.ty();
             let handlers = match record_type.kind() {
-                ValueKind::RecordType(kts) => kts,
+                NirKind::RecordType(kts) => kts,
                 _ => return span_err("Merge1ArgMustBeRecord"),
             };
 
-            let union_type = union.get_type()?;
+            let union_type = union.ty();
             let variants = match union_type.kind() {
-                ValueKind::UnionType(kts) => Cow::Borrowed(kts),
-                ValueKind::AppliedBuiltin(BuiltinClosure {
+                NirKind::UnionType(kts) => Cow::Borrowed(kts),
+                NirKind::AppliedBuiltin(BuiltinClosure {
                     b: Builtin::Optional,
                     args,
                     ..
@@ -509,10 +450,10 @@ fn type_one_layer(
 
             let mut inferred_type = None;
             for (x, handler_type) in handlers {
-                let handler_return_type = match variants.get(x) {
+                let handler_return_type: Type = match variants.get(x) {
                     // Union alternative with type
                     Some(Some(variant_type)) => match handler_type.kind() {
-                        ValueKind::PiClosure { closure, annot, .. } => {
+                        NirKind::PiClosure { closure, annot, .. } => {
                             if variant_type != annot {
                                 return mkerr(
                                     ErrorBuilder::new(format!(
@@ -543,8 +484,11 @@ fn type_one_layer(
                                 );
                             }
 
+                            // TODO: this actually doesn't check anything yet
                             match closure.remove_binder() {
-                                Ok(v) => v,
+                                Ok(v) => {
+                                    Type::new_infer_universe(env, v.clone())?
+                                }
                                 Err(()) => {
                                     return span_err(
                                         "MergeReturnTypeIsDependent",
@@ -588,7 +532,9 @@ fn type_one_layer(
                         }
                     },
                     // Union alternative without type
-                    Some(None) => handler_type.clone(),
+                    Some(None) => {
+                        Type::new_infer_universe(env, handler_type.clone())?
+                    }
                     None => return span_err("MergeHandlerMissingVariant"),
                 };
                 match &inferred_type {
@@ -606,8 +552,10 @@ fn type_one_layer(
                 }
             }
 
-            let type_annot =
-                type_annot.as_ref().map(|t| t.eval(env.as_nzenv()));
+            let type_annot = type_annot
+                .as_ref()
+                .map(|t| t.eval_to_type(env))
+                .transpose()?;
             match (inferred_type, type_annot) {
                 (Some(t1), Some(t2)) => {
                     if t1 != t2 {
@@ -621,9 +569,12 @@ fn type_one_layer(
             }
         }
         ExprKind::ToMap(record, annot) => {
-            let record_t = record.get_type()?;
+            if record.ty().ty().as_const() != Some(Const::Type) {
+                return span_err("`toMap` only accepts records of type `Type`");
+            }
+            let record_t = record.ty();
             let kts = match record_t.kind() {
-                ValueKind::RecordType(kts) => kts,
+                NirKind::RecordType(kts) => kts,
                 _ => {
                     return span_err("The argument to `toMap` must be a record")
                 }
@@ -638,12 +589,12 @@ fn type_one_layer(
                          annotation",
                     );
                 };
-                let annot_val = annot.eval(env.as_nzenv());
+                let annot_val = annot.eval_to_type(env)?;
 
                 let err_msg = "The type of `toMap x` must be of the form \
                                `List { mapKey : Text, mapValue : T }`";
                 let arg = match annot_val.kind() {
-                    ValueKind::AppliedBuiltin(BuiltinClosure {
+                    NirKind::AppliedBuiltin(BuiltinClosure {
                         b: Builtin::List,
                         args,
                         ..
@@ -651,14 +602,14 @@ fn type_one_layer(
                     _ => return span_err(err_msg),
                 };
                 let kts = match arg.kind() {
-                    ValueKind::RecordType(kts) => kts,
+                    NirKind::RecordType(kts) => kts,
                     _ => return span_err(err_msg),
                 };
                 if kts.len() != 2 {
                     return span_err(err_msg);
                 }
                 match kts.get(&"mapKey".into()) {
-                    Some(t) if *t == Value::from_builtin(Builtin::Text) => {}
+                    Some(t) if *t == Nir::from_builtin(Builtin::Text) => {}
                     _ => return span_err(err_msg),
                 }
                 match kts.get(&"mapValue".into()) {
@@ -668,11 +619,6 @@ fn type_one_layer(
                 annot_val
             } else {
                 let entry_type = kts.iter().next().unwrap().1.clone();
-                if entry_type.get_type()?.as_const() != Some(Const::Type) {
-                    return span_err(
-                        "`toMap` only accepts records of type `Type`",
-                    );
-                }
                 for (_, t) in kts.iter() {
                     if *t != entry_type {
                         return span_err(
@@ -682,16 +628,13 @@ fn type_one_layer(
                 }
 
                 let mut kts = HashMap::new();
-                kts.insert("mapKey".into(), Value::from_builtin(Builtin::Text));
+                kts.insert("mapKey".into(), Nir::from_builtin(Builtin::Text));
                 kts.insert("mapValue".into(), entry_type);
-                let output_type = Value::from_builtin(Builtin::List).app(
-                    Value::from_kind_and_type(
-                        ValueKind::RecordType(kts),
-                        Value::from_const(Const::Type),
-                    ),
-                );
+                let output_type: Type = Nir::from_builtin(Builtin::List)
+                    .app(Nir::from_kind(NirKind::RecordType(kts)))
+                    .to_type(Const::Type);
                 if let Some(annot) = annot {
-                    let annot_val = annot.eval(env.as_nzenv());
+                    let annot_val = annot.eval_to_type(env)?;
                     if output_type != annot_val {
                         return span_err("Annotation mismatch");
                     }
@@ -700,9 +643,9 @@ fn type_one_layer(
             }
         }
         ExprKind::Projection(record, labels) => {
-            let record_type = record.get_type()?;
+            let record_type = record.ty();
             let kts = match record_type.kind() {
-                ValueKind::RecordType(kts) => kts,
+                NirKind::RecordType(kts) => kts,
                 _ => return span_err("ProjectionMustBeRecord"),
             };
 
@@ -722,21 +665,21 @@ fn type_one_layer(
                 };
             }
 
-            Value::from_kind_and_type(
-                ValueKind::RecordType(new_kts),
-                record_type.get_type()?,
-            )
+            Type::new_infer_universe(
+                env,
+                Nir::from_kind(NirKind::RecordType(new_kts)),
+            )?
         }
         ExprKind::ProjectionByExpr(record, selection) => {
-            let record_type = record.get_type()?;
+            let record_type = record.ty();
             let rec_kts = match record_type.kind() {
-                ValueKind::RecordType(kts) => kts,
+                NirKind::RecordType(kts) => kts,
                 _ => return span_err("ProjectionMustBeRecord"),
             };
 
-            let selection_val = selection.eval(env.as_nzenv());
+            let selection_val = selection.eval_to_type(env)?;
             let sel_kts = match selection_val.kind() {
-                ValueKind::RecordType(kts) => kts,
+                NirKind::RecordType(kts) => kts,
                 _ => return span_err("ProjectionByExprTakesRecordType"),
             };
 
@@ -753,110 +696,119 @@ fn type_one_layer(
 
             selection_val
         }
-        ExprKind::Completion(ty, compl) => {
-            let ty_field_default = type_one_layer(
-                env,
-                ExprKind::Field(ty.clone(), "default".into()),
-                span.clone(),
-            )?;
-            let ty_field_type = type_one_layer(
-                env,
-                ExprKind::Field(ty.clone(), "Type".into()),
-                span.clone(),
-            )?;
-            let merge = type_one_layer(
-                env,
-                ExprKind::BinOp(
-                    BinOp::RightBiasedRecordMerge,
-                    ty_field_default,
-                    compl.clone(),
-                ),
-                span.clone(),
-            )?;
-            return type_one_layer(
-                env,
-                ExprKind::Annot(merge, ty_field_type),
-                span.clone(),
-            );
-        }
     };
 
-    Ok(TyExpr::new(TyExprKind::Expr(ekind), Some(ty), span))
+    Ok(ty)
 }
 
-/// `type_with` typechecks an expressio in the provided environment.
-pub(crate) fn type_with(
+/// `type_with` typechecks an expression in the provided environment. Optionally pass an annotation
+/// to compare with.
+pub(crate) fn type_with<'hir>(
     env: &TyEnv,
-    expr: &Expr<Normalized>,
-) -> Result<TyExpr, TypeError> {
-    let (tyekind, ty) = match expr.as_ref() {
-        ExprKind::Var(var) => match env.lookup(&var) {
-            Some((v, ty)) => (TyExprKind::Var(v), Some(ty)),
-            None => {
-                return mkerr(
-                    ErrorBuilder::new(format!("unbound variable `{}`", var))
-                        .span_err(expr.span(), "not found in this scope")
-                        .format(),
-                )
-            }
-        },
-        ExprKind::Const(Const::Sort) => {
-            (TyExprKind::Expr(ExprKind::Const(Const::Sort)), None)
-        }
-        ExprKind::Embed(p) => {
-            return Ok(p.clone().into_value().to_tyexpr_noenv())
-        }
-        ekind => {
-            let ekind = match ekind {
-                ExprKind::Lam(binder, annot, body) => {
-                    let annot = type_with(env, annot)?;
-                    let annot_nf = annot.eval(env.as_nzenv());
-                    let body_env = env.insert_type(binder, annot_nf);
-                    let body = type_with(&body_env, body)?;
-                    ExprKind::Lam(binder.clone(), annot, body)
-                }
-                ExprKind::Pi(binder, annot, body) => {
-                    let annot = type_with(env, annot)?;
-                    let annot_nf = annot.eval(env.as_nzenv());
-                    let body_env = env.insert_type(binder, annot_nf);
-                    let body = type_with(&body_env, body)?;
-                    ExprKind::Pi(binder.clone(), annot, body)
-                }
-                ExprKind::Let(binder, annot, val, body) => {
-                    let val = if let Some(t) = annot {
-                        t.rewrap(ExprKind::Annot(val.clone(), t.clone()))
-                    } else {
-                        val.clone()
-                    };
-                    let val = type_with(env, &val)?;
-                    val.get_type()?; // Ensure val is not Sort
-                    let val_nf = val.eval(&env.as_nzenv());
-                    let body_env = env.insert_value(&binder, val_nf);
-                    let body = type_with(&body_env, body)?;
-                    ExprKind::Let(binder.clone(), None, val, body)
+    hir: &'hir Hir,
+    annot: Option<Type>,
+) -> Result<Tir<'hir>, TypeError> {
+    let tir = match hir.kind() {
+        HirKind::Var(var) => Tir::from_hir(hir, env.lookup(var)),
+        HirKind::Import(_, ty) => Tir::from_hir(hir, ty.clone()),
+        HirKind::Expr(ExprKind::Var(_)) => {
+            unreachable!("Hir should contain no unresolved variables")
+        }
+        HirKind::Expr(ExprKind::Const(Const::Sort)) => {
+            return mk_span_err(hir.span(), "Sort does not have a type")
+        }
+        HirKind::Expr(ExprKind::Annot(x, t)) => {
+            let t = match t.kind() {
+                HirKind::Expr(ExprKind::Const(Const::Sort)) => {
+                    Type::from_const(Const::Sort)
                 }
-                _ => ekind.traverse_ref(|e| type_with(env, e))?,
+                _ => type_with(env, t, None)?.eval_to_type(env)?,
             };
-            return type_one_layer(env, ekind, expr.span());
+            type_with(env, x, Some(t))?
+        }
+
+        HirKind::Expr(ExprKind::Lam(binder, annot, body)) => {
+            let annot = type_with(env, annot, None)?;
+            let annot_nf = annot.eval_to_type(env)?;
+            let body_env = env.insert_type(binder, annot_nf);
+            let body = type_with(&body_env, body, None)?;
+
+            let u_annot = annot.ty().as_const().unwrap();
+            let u_body = match body.ty().ty().as_const() {
+                Some(k) => k,
+                _ => return mk_span_err(hir.span(), "Invalid output type"),
+            };
+            let u = function_check(u_annot, u_body).to_universe();
+            let ty_hir = Hir::new(
+                HirKind::Expr(ExprKind::Pi(
+                    binder.clone(),
+                    annot.to_hir(),
+                    body.ty().to_hir(body_env.as_varenv()),
+                )),
+                hir.span(),
+            );
+            let ty = Type::new(ty_hir.eval(env), u);
+
+            Tir::from_hir(hir, ty)
+        }
+        HirKind::Expr(ExprKind::Pi(binder, annot, body)) => {
+            let annot = type_with(env, annot, None)?;
+            let annot_val = annot.eval_to_type(env)?;
+            let body_env = env.insert_type(binder, annot_val);
+            let body = type_with(&body_env, body, None)?;
+            body.ensure_is_type(env)?;
+
+            let ks = annot.ty().as_const().unwrap();
+            let kt = body.ty().as_const().unwrap();
+            let ty = Type::from_const(function_check(ks, kt));
+            Tir::from_hir(hir, ty)
+        }
+        HirKind::Expr(ExprKind::Let(binder, annot, val, body)) => {
+            let val_annot = annot
+                .as_ref()
+                .map(|t| Ok(type_with(env, t, None)?.eval_to_type(env)?))
+                .transpose()?;
+            let val = type_with(env, &val, val_annot)?;
+            let val_nf = val.eval(env);
+            let body_env = env.insert_value(&binder, val_nf, val.ty().clone());
+            let body = type_with(&body_env, body, None)?;
+            let ty = body.ty().clone();
+            Tir::from_hir(hir, ty)
+        }
+        HirKind::Expr(ekind) => {
+            let ekind = ekind.traverse_ref(|e| type_with(env, e, None))?;
+            let ty = type_one_layer(env, ekind, hir.span())?;
+            Tir::from_hir(hir, ty)
         }
     };
 
-    Ok(TyExpr::new(tyekind, ty, expr.span()))
+    if let Some(annot) = annot {
+        if *tir.ty() != annot {
+            return mk_span_err(
+                hir.span(),
+                &format!(
+                    "annot mismatch: {} != {}",
+                    tir.ty().to_expr_tyenv(env),
+                    annot.to_expr_tyenv(env)
+                ),
+            );
+        }
+    }
+
+    Ok(tir)
 }
 
 /// Typecheck an expression and return the expression annotated with types if type-checking
 /// succeeded, or an error if type-checking failed.
-pub(crate) fn typecheck(e: &Expr<Normalized>) -> Result<TyExpr, TypeError> {
-    let res = type_with(&TyEnv::new(), e)?;
-    // Ensure that the inferred type exists (i.e. this is not Sort)
-    res.get_type()?;
-    Ok(res)
+pub(crate) fn typecheck<'hir>(hir: &'hir Hir) -> Result<Tir<'hir>, TypeError> {
+    type_with(&TyEnv::new(), hir, None)
 }
 
 /// Like `typecheck`, but additionally checks that the expression's type matches the provided type.
-pub(crate) fn typecheck_with(
-    expr: &Expr<Normalized>,
-    ty: Expr<Normalized>,
-) -> Result<TyExpr, TypeError> {
-    typecheck(&expr.rewrap(ExprKind::Annot(expr.clone(), ty)))
+pub(crate) fn typecheck_with<'hir>(
+    hir: &'hir Hir,
+    ty: Hir,
+) -> Result<Tir<'hir>, TypeError> {
+    let ty = typecheck(&ty)?.eval_to_type(&TyEnv::new())?;
+    type_with(&TyEnv::new(), hir, Some(ty))
 }