Merge pull request #204 from Nadrieril/source-id

3 files changed, 160 insertions, 136 deletions

diff --git a/dhall/src/semantics/nze/env.rs b/dhall/src/semantics/nze/env.rs
index ec99dbe..30fa8af 100644
--- a/dhall/src/semantics/nze/env.rs
+++ b/dhall/src/semantics/nze/env.rs
@@ -1,4 +1,5 @@
 use crate::semantics::{AlphaVar, Nir, NirKind};
+use crate::Ctxt;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum NzVar {
@@ -9,19 +10,20 @@ pub enum NzVar {
 }
 
 #[derive(Debug, Clone)]
-enum EnvItem<Type> {
+enum EnvItem<'cx, T> {
     // Variable is bound with given type
-    Kept(Type),
+    Kept(T),
     // Variable has been replaced by corresponding value
-    Replaced(Nir, Type),
+    Replaced(Nir<'cx>, T),
 }
 
 #[derive(Debug, Clone)]
-pub struct ValEnv<Type> {
-    items: Vec<EnvItem<Type>>,
+pub struct ValEnv<'cx, T> {
+    cx: Ctxt<'cx>,
+    items: Vec<EnvItem<'cx, T>>,
 }
 
-pub type NzEnv = ValEnv<()>;
+pub type NzEnv<'cx> = ValEnv<'cx, ()>;
 
 impl NzVar {
     pub fn new(idx: usize) -> Self {
@@ -46,11 +48,17 @@ impl NzVar {
     }
 }
 
-impl<Type: Clone> ValEnv<Type> {
-    pub fn new() -> Self {
-        ValEnv { items: Vec::new() }
+impl<'cx, T: Clone> ValEnv<'cx, T> {
+    pub fn new(cx: Ctxt<'cx>) -> Self {
+        ValEnv {
+            cx,
+            items: Vec::new(),
+        }
+    }
+    pub fn cx(&self) -> Ctxt<'cx> {
+        self.cx
     }
-    pub fn discard_types(&self) -> ValEnv<()> {
+    pub fn discard_types(&self) -> ValEnv<'cx, ()> {
         let items = self
             .items
             .iter()
@@ -59,27 +67,27 @@ impl<Type: Clone> ValEnv<Type> {
                 EnvItem::Replaced(val, _) => EnvItem::Replaced(val.clone(), ()),
             })
             .collect();
-        ValEnv { items }
+        ValEnv { cx: self.cx, items }
     }
 
-    pub fn insert_type(&self, ty: Type) -> Self {
+    pub fn insert_type(&self, ty: T) -> Self {
         let mut env = self.clone();
         env.items.push(EnvItem::Kept(ty));
         env
     }
-    pub fn insert_value(&self, e: Nir, ty: Type) -> Self {
+    pub fn insert_value(&self, e: Nir<'cx>, ty: T) -> Self {
         let mut env = self.clone();
         env.items.push(EnvItem::Replaced(e, ty));
         env
     }
-    pub fn lookup_val(&self, var: AlphaVar) -> NirKind {
+    pub fn lookup_val(&self, var: AlphaVar) -> NirKind<'cx> {
         let idx = self.items.len() - 1 - var.idx();
         match &self.items[idx] {
             EnvItem::Kept(_) => NirKind::Var(NzVar::new(idx)),
             EnvItem::Replaced(x, _) => x.kind().clone(),
         }
     }
-    pub fn lookup_ty(&self, var: AlphaVar) -> Type {
+    pub fn lookup_ty(&self, var: AlphaVar) -> T {
         let idx = self.items.len() - 1 - var.idx();
         match &self.items[idx] {
             EnvItem::Kept(ty) | EnvItem::Replaced(_, ty) => ty.clone(),
@@ -87,8 +95,8 @@ impl<Type: Clone> ValEnv<Type> {
     }
 }
 
-impl<'a> From<&'a NzEnv> for NzEnv {
-    fn from(x: &'a NzEnv) -> Self {
+impl<'cx, 'a> From<&'a NzEnv<'cx>> for NzEnv<'cx> {
+    fn from(x: &'a NzEnv<'cx>) -> Self {
         x.clone()
     }
 }
diff --git a/dhall/src/semantics/nze/nir.rs b/dhall/src/semantics/nze/nir.rs
index 7bda836..8cf06c5 100644
--- a/dhall/src/semantics/nze/nir.rs
+++ b/dhall/src/semantics/nze/nir.rs
@@ -11,7 +11,7 @@ use crate::semantics::{
 use crate::syntax::{
     Const, Expr, ExprKind, InterpolatedTextContents, Label, NumKind, Span,
 };
-use crate::ToExprOptions;
+use crate::{Ctxt, ToExprOptions};
 
 /// Stores a possibly unevaluated value. Gets (partially) normalized on-demand, sharing computation
 /// automatically. Uses a Rc<OnceCell> to share computation.
@@ -20,31 +20,31 @@ use crate::ToExprOptions;
 /// normalize as needed.
 /// Stands for "Normalized Intermediate Representation"
 #[derive(Clone)]
-pub struct Nir(Rc<lazy::Lazy<Thunk, NirKind>>);
+pub struct Nir<'cx>(Rc<lazy::Lazy<Thunk<'cx>, NirKind<'cx>>>);
 
 /// An unevaluated subexpression
 #[derive(Debug, Clone)]
-pub enum Thunk {
+pub enum Thunk<'cx> {
     /// A completely unnormalized expression.
-    Thunk { env: NzEnv, body: Hir },
+    Thunk { env: NzEnv<'cx>, body: Hir<'cx> },
     /// A partially normalized expression that may need to go through `normalize_one_layer`.
-    PartialExpr { env: NzEnv, expr: ExprKind<Nir> },
+    PartialExpr { expr: ExprKind<Nir<'cx>> },
 }
 
 /// An unevaluated subexpression that takes an argument.
 #[derive(Debug, Clone)]
-pub enum Closure {
+pub enum Closure<'cx> {
     /// Normal closure
-    Closure { env: NzEnv, body: Hir },
+    Closure { env: NzEnv<'cx>, body: Hir<'cx> },
     /// Closure that ignores the argument passed
-    ConstantClosure { body: Nir },
+    ConstantClosure { body: Nir<'cx> },
 }
 
 /// A text literal with interpolations.
 // Invariant: this must not contain interpolations that are themselves TextLits, and contiguous
 // text values must be merged.
 #[derive(Debug, Clone, PartialEq, Eq)]
-pub struct TextLit(Vec<InterpolatedTextContents<Nir>>);
+pub struct TextLit<'cx>(Vec<InterpolatedTextContents<Nir<'cx>>>);
 
 /// This represents a value in Weak Head Normal Form (WHNF). This means that the value is
 /// normalized up to the first constructor, but subexpressions may not be fully normalized.
@@ -54,67 +54,65 @@ pub struct TextLit(Vec<InterpolatedTextContents<Nir>>);
 /// In particular, this means that once we get a `NirKind`, it can be considered immutable, and
 /// we only need to recursively normalize its sub-`Nir`s to get to the NF.
 #[derive(Debug, Clone, PartialEq, Eq)]
-pub enum NirKind {
+pub enum NirKind<'cx> {
     /// Closures
     LamClosure {
         binder: Binder,
-        annot: Nir,
-        closure: Closure,
+        annot: Nir<'cx>,
+        closure: Closure<'cx>,
     },
     PiClosure {
         binder: Binder,
-        annot: Nir,
-        closure: Closure,
+        annot: Nir<'cx>,
+        closure: Closure<'cx>,
     },
-    AppliedBuiltin(BuiltinClosure),
+    AppliedBuiltin(BuiltinClosure<'cx>),
 
     Var(NzVar),
     Const(Const),
     Num(NumKind),
     // Must be a number type, Bool or Text
     BuiltinType(Builtin),
-    TextLit(TextLit),
-    EmptyOptionalLit(Nir),
-    NEOptionalLit(Nir),
-    OptionalType(Nir),
+    TextLit(TextLit<'cx>),
+    EmptyOptionalLit(Nir<'cx>),
+    NEOptionalLit(Nir<'cx>),
+    OptionalType(Nir<'cx>),
     // EmptyListLit(t) means `[] : List t`, not `[] : t`
-    EmptyListLit(Nir),
-    NEListLit(Vec<Nir>),
-    ListType(Nir),
-    RecordLit(HashMap<Label, Nir>),
-    RecordType(HashMap<Label, Nir>),
-    UnionConstructor(Label, HashMap<Label, Option<Nir>>),
-    UnionLit(Label, Nir, HashMap<Label, Option<Nir>>),
-    UnionType(HashMap<Label, Option<Nir>>),
-    Equivalence(Nir, Nir),
-    Assert(Nir),
+    EmptyListLit(Nir<'cx>),
+    NEListLit(Vec<Nir<'cx>>),
+    ListType(Nir<'cx>),
+    RecordLit(HashMap<Label, Nir<'cx>>),
+    RecordType(HashMap<Label, Nir<'cx>>),
+    UnionConstructor(Label, HashMap<Label, Option<Nir<'cx>>>),
+    UnionLit(Label, Nir<'cx>, HashMap<Label, Option<Nir<'cx>>>),
+    UnionType(HashMap<Label, Option<Nir<'cx>>>),
+    Equivalence(Nir<'cx>, Nir<'cx>),
+    Assert(Nir<'cx>),
     /// Invariant: evaluation must not be able to progress with `normalize_operation`.
     /// This is used when an operation couldn't proceed further, for example because of variables.
-    Op(OpKind<Nir>),
+    Op(OpKind<Nir<'cx>>),
 }
 
-impl Nir {
+impl<'cx> Nir<'cx> {
     /// Construct a Nir from a completely unnormalized expression.
-    pub fn new_thunk(env: NzEnv, hir: Hir) -> Nir {
+    pub fn new_thunk(env: NzEnv<'cx>, hir: Hir<'cx>) -> Self {
         Nir(Rc::new(lazy::Lazy::new(Thunk::new(env, hir))))
     }
     /// Construct a Nir from a partially normalized expression that's not in WHNF.
-    pub fn from_partial_expr(e: ExprKind<Nir>) -> Nir {
-        // TODO: env
-        let env = NzEnv::new();
-        Nir(Rc::new(lazy::Lazy::new(Thunk::from_partial_expr(env, e))))
+    pub fn from_partial_expr(e: ExprKind<Self>) -> Self {
+        Nir(Rc::new(lazy::Lazy::new(Thunk::from_partial_expr(e))))
     }
     /// Make a Nir from a NirKind
-    pub fn from_kind(v: NirKind) -> Nir {
+    pub fn from_kind(v: NirKind<'cx>) -> Self {
         Nir(Rc::new(lazy::Lazy::new_completed(v)))
     }
     pub fn from_const(c: Const) -> Self {
         Self::from_kind(NirKind::Const(c))
     }
-    pub fn from_builtin(b: Builtin) -> Self {
-        Self::from_builtin_env(b, &NzEnv::new())
+    pub fn from_builtin(cx: Ctxt<'cx>, b: Builtin) -> Self {
+        Self::from_builtin_env(b, &NzEnv::new(cx))
     }
-    pub fn from_builtin_env(b: Builtin, env: &NzEnv) -> Self {
+    pub fn from_builtin_env(b: Builtin, env: &NzEnv<'cx>) -> Self {
         Nir::from_kind(NirKind::from_builtin_env(b, env.clone()))
     }
     pub fn from_text(txt: impl ToString) -> Self {
@@ -129,51 +127,54 @@ impl Nir {
     }
 
     /// This is what you want if you want to pattern-match on the value.
-    pub fn kind(&self) -> &NirKind {
+    pub fn kind(&self) -> &NirKind<'cx> {
         &*self.0
     }
 
     /// The contents of a `Nir` are immutable and shared. If however we happen to be the sole
     /// owners, we can mutate it directly. Otherwise, this clones the internal value first.
-    pub fn kind_mut(&mut self) -> &mut NirKind {
+    pub fn kind_mut(&mut self) -> &mut NirKind<'cx> {
         Rc::make_mut(&mut self.0).get_mut()
     }
     /// If we are the sole owner of this Nir, we can avoid a clone.
-    pub fn into_kind(self) -> NirKind {
+    pub fn into_kind(self) -> NirKind<'cx> {
         match Rc::try_unwrap(self.0) {
             Ok(lazy) => lazy.into_inner(),
             Err(rc) => (**rc).clone(),
         }
     }
 
-    pub fn to_type(&self, u: impl Into<Universe>) -> Type {
+    pub fn to_type(&self, u: impl Into<Universe>) -> Type<'cx> {
         Type::new(self.clone(), u.into())
     }
     /// Converts a value back to the corresponding AST expression.
-    pub fn to_expr(&self, opts: ToExprOptions) -> Expr {
-        self.to_hir_noenv().to_expr(opts)
+    pub fn to_expr(&self, cx: Ctxt<'cx>, opts: ToExprOptions) -> Expr {
+        self.to_hir_noenv().to_expr(cx, opts)
     }
-    pub fn to_expr_tyenv(&self, tyenv: &TyEnv) -> Expr {
+    pub fn to_expr_tyenv(&self, tyenv: &TyEnv<'cx>) -> Expr {
         self.to_hir(tyenv.as_varenv()).to_expr_tyenv(tyenv)
     }
 
-    pub fn app(&self, v: Nir) -> Nir {
+    pub fn app(&self, v: Self) -> Self {
         Nir::from_kind(self.app_to_kind(v))
     }
-    pub fn app_to_kind(&self, v: Nir) -> NirKind {
+    pub fn app_to_kind(&self, v: Self) -> NirKind<'cx> {
         apply_any(self, v)
     }
 
-    pub fn to_hir(&self, venv: VarEnv) -> Hir {
-        let map_uniontype = |kts: &HashMap<Label, Option<Nir>>| {
-            ExprKind::UnionType(
-                kts.iter()
-                    .map(|(k, v)| {
-                        (k.clone(), v.as_ref().map(|v| v.to_hir(venv)))
-                    })
-                    .collect(),
-            )
-        };
+    pub fn to_hir(&self, venv: VarEnv) -> Hir<'cx> {
+        let map_uniontype =
+            |kts: &HashMap<Label, Option<Nir<'cx>>>| -> ExprKind<Hir<'cx>> {
+                ExprKind::UnionType(
+                    kts.iter()
+                        .map(|(k, v)| {
+                            (k.clone(), v.as_ref().map(|v| v.to_hir(venv)))
+                        })
+                        .collect(),
+                )
+            };
+        let builtin =
+            |b| Hir::new(HirKind::Expr(ExprKind::Builtin(b)), Span::Artificial);
 
         let hir = match self.kind() {
             NirKind::Var(v) => HirKind::Var(venv.lookup(*v)),
@@ -204,21 +205,21 @@ impl Nir {
                 NirKind::BuiltinType(b) => ExprKind::Builtin(*b),
                 NirKind::Num(l) => ExprKind::Num(l.clone()),
                 NirKind::OptionalType(t) => ExprKind::Op(OpKind::App(
-                    Nir::from_builtin(Builtin::Optional).to_hir(venv),
+                    builtin(Builtin::Optional),
                     t.to_hir(venv),
                 )),
                 NirKind::EmptyOptionalLit(n) => ExprKind::Op(OpKind::App(
-                    Nir::from_builtin(Builtin::OptionalNone).to_hir(venv),
+                    builtin(Builtin::OptionalNone),
                     n.to_hir(venv),
                 )),
                 NirKind::NEOptionalLit(n) => ExprKind::SomeLit(n.to_hir(venv)),
                 NirKind::ListType(t) => ExprKind::Op(OpKind::App(
-                    Nir::from_builtin(Builtin::List).to_hir(venv),
+                    builtin(Builtin::List),
                     t.to_hir(venv),
                 )),
                 NirKind::EmptyListLit(n) => ExprKind::EmptyListLit(Hir::new(
                     HirKind::Expr(ExprKind::Op(OpKind::App(
-                        Nir::from_builtin(Builtin::List).to_hir(venv),
+                        builtin(Builtin::List),
                         n.to_hir(venv),
                     ))),
                     Span::Artificial,
@@ -276,52 +277,52 @@ impl Nir {
 
         Hir::new(hir, Span::Artificial)
     }
-    pub fn to_hir_noenv(&self) -> Hir {
+    pub fn to_hir_noenv(&self) -> Hir<'cx> {
         self.to_hir(VarEnv::new())
     }
 }
 
-impl NirKind {
-    pub fn into_nir(self) -> Nir {
+impl<'cx> NirKind<'cx> {
+    pub fn into_nir(self) -> Nir<'cx> {
         Nir::from_kind(self)
     }
 
-    pub fn from_builtin(b: Builtin) -> NirKind {
-        NirKind::from_builtin_env(b, NzEnv::new())
+    pub fn from_builtin(cx: Ctxt<'cx>, b: Builtin) -> Self {
+        NirKind::from_builtin_env(b, NzEnv::new(cx))
     }
-    pub fn from_builtin_env(b: Builtin, env: NzEnv) -> NirKind {
+    pub fn from_builtin_env(b: Builtin, env: NzEnv<'cx>) -> Self {
         BuiltinClosure::new(b, env)
     }
 }
 
-impl Thunk {
-    fn new(env: NzEnv, body: Hir) -> Self {
+impl<'cx> Thunk<'cx> {
+    fn new(env: NzEnv<'cx>, body: Hir<'cx>) -> Self {
         Thunk::Thunk { env, body }
     }
-    fn from_partial_expr(env: NzEnv, expr: ExprKind<Nir>) -> Self {
-        Thunk::PartialExpr { env, expr }
+    fn from_partial_expr(expr: ExprKind<Nir<'cx>>) -> Self {
+        Thunk::PartialExpr { expr }
     }
-    fn eval(self) -> NirKind {
+    fn eval(self) -> NirKind<'cx> {
         match self {
-            Thunk::Thunk { env, body } => normalize_hir(&env, &body),
-            Thunk::PartialExpr { env, expr } => normalize_one_layer(expr, &env),
+            Thunk::Thunk { env, body, .. } => normalize_hir(&env, &body),
+            Thunk::PartialExpr { expr } => normalize_one_layer(expr),
         }
     }
 }
 
-impl Closure {
-    pub fn new(env: &NzEnv, body: Hir) -> Self {
+impl<'cx> Closure<'cx> {
+    pub fn new(env: &NzEnv<'cx>, body: Hir<'cx>) -> Self {
         Closure::Closure {
             env: env.clone(),
             body,
         }
     }
     /// New closure that ignores its argument
-    pub fn new_constant(body: Nir) -> Self {
+    pub fn new_constant(body: Nir<'cx>) -> Self {
         Closure::ConstantClosure { body }
     }
 
-    pub fn apply(&self, val: Nir) -> Nir {
+    pub fn apply(&self, val: Nir<'cx>) -> Nir<'cx> {
         match self {
             Closure::Closure { env, body, .. } => {
                 body.eval(env.insert_value(val, ()))
@@ -329,7 +330,7 @@ impl Closure {
             Closure::ConstantClosure { body, .. } => body.clone(),
         }
     }
-    fn apply_var(&self, var: NzVar) -> Nir {
+    fn apply_var(&self, var: NzVar) -> Nir<'cx> {
         match self {
             Closure::Closure { .. } => {
                 self.apply(Nir::from_kind(NirKind::Var(var)))
@@ -339,13 +340,13 @@ impl Closure {
     }
 
     /// Convert this closure to a Hir expression
-    pub fn to_hir(&self, venv: VarEnv) -> Hir {
+    pub fn to_hir(&self, venv: VarEnv) -> Hir<'cx> {
         self.apply_var(NzVar::new(venv.size()))
             .to_hir(venv.insert())
     }
     /// If the closure variable is free in the closure, return Err. Otherwise, return the value
     /// with that free variable remove.
-    pub fn remove_binder(&self) -> Result<Nir, ()> {
+    pub fn remove_binder(&self) -> Result<Nir<'cx>, ()> {
         match self {
             Closure::Closure { .. } => {
                 let v = NzVar::fresh();
@@ -358,20 +359,20 @@ impl Closure {
     }
 }
 
-impl TextLit {
+impl<'cx> TextLit<'cx> {
     pub fn new(
-        elts: impl Iterator<Item = InterpolatedTextContents<Nir>>,
+        elts: impl Iterator<Item = InterpolatedTextContents<Nir<'cx>>>,
     ) -> Self {
         TextLit(squash_textlit(elts))
     }
-    pub fn interpolate(v: Nir) -> TextLit {
+    pub fn interpolate(v: Nir<'cx>) -> Self {
         TextLit(vec![InterpolatedTextContents::Expr(v)])
     }
-    pub fn from_text(s: String) -> TextLit {
+    pub fn from_text(s: String) -> Self {
         TextLit(vec![InterpolatedTextContents::Text(s)])
     }
 
-    pub fn concat(&self, other: &TextLit) -> TextLit {
+    pub fn concat(&self, other: &Self) -> Self {
         TextLit::new(self.iter().chain(other.iter()).cloned())
     }
     pub fn is_empty(&self) -> bool {
@@ -379,7 +380,7 @@ impl TextLit {
     }
     /// If the literal consists of only one interpolation and not text, return the interpolated
     /// value.
-    pub fn as_single_expr(&self) -> Option<&Nir> {
+    pub fn as_single_expr(&self) -> Option<&Nir<'cx>> {
         use InterpolatedTextContents::Expr;
         if let [Expr(v)] = self.0.as_slice() {
             Some(v)
@@ -398,43 +399,45 @@ impl TextLit {
             None
         }
     }
-    pub fn iter(&self) -> impl Iterator<Item = &InterpolatedTextContents<Nir>> {
+    pub fn iter(
+        &self,
+    ) -> impl Iterator<Item = &InterpolatedTextContents<Nir<'cx>>> {
         self.0.iter()
     }
 }
 
-impl lazy::Eval<NirKind> for Thunk {
-    fn eval(self) -> NirKind {
+impl<'cx> lazy::Eval<NirKind<'cx>> for Thunk<'cx> {
+    fn eval(self) -> NirKind<'cx> {
         self.eval()
     }
 }
 
 /// Compare two values for equality modulo alpha/beta-equivalence.
-impl std::cmp::PartialEq for Nir {
+impl<'cx> std::cmp::PartialEq for Nir<'cx> {
     fn eq(&self, other: &Self) -> bool {
         Rc::ptr_eq(&self.0, &other.0) || self.kind() == other.kind()
     }
 }
-impl std::cmp::Eq for Nir {}
+impl<'cx> std::cmp::Eq for Nir<'cx> {}
 
-impl std::cmp::PartialEq for Thunk {
+impl<'cx> std::cmp::PartialEq for Thunk<'cx> {
     fn eq(&self, _other: &Self) -> bool {
         unreachable!(
             "Trying to compare thunks but we should only compare WHNFs"
         )
     }
 }
-impl std::cmp::Eq for Thunk {}
+impl<'cx> std::cmp::Eq for Thunk<'cx> {}
 
-impl std::cmp::PartialEq for Closure {
+impl<'cx> std::cmp::PartialEq for Closure<'cx> {
     fn eq(&self, other: &Self) -> bool {
         let v = NzVar::fresh();
         self.apply_var(v) == other.apply_var(v)
     }
 }
-impl std::cmp::Eq for Closure {}
+impl<'cx> std::cmp::Eq for Closure<'cx> {}
 
-impl std::fmt::Debug for Nir {
+impl<'cx> std::fmt::Debug for Nir<'cx> {
     fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         if let NirKind::Const(c) = self.kind() {
             return write!(fmt, "{:?}", c);
diff --git a/dhall/src/semantics/nze/normalize.rs b/dhall/src/semantics/nze/normalize.rs
index 62efc5f..59710d1 100644
--- a/dhall/src/semantics/nze/normalize.rs
+++ b/dhall/src/semantics/nze/normalize.rs
@@ -5,7 +5,7 @@ use crate::semantics::NzEnv;
 use crate::semantics::{Binder, Closure, Hir, HirKind, Nir, NirKind, TextLit};
 use crate::syntax::{ExprKind, InterpolatedTextContents};
 
-pub fn apply_any(f: &Nir, a: Nir) -> NirKind {
+pub fn apply_any<'cx>(f: &Nir<'cx>, a: Nir<'cx>) -> NirKind<'cx> {
     match f.kind() {
         NirKind::LamClosure { closure, .. } => closure.apply(a).kind().clone(),
         NirKind::AppliedBuiltin(closure) => closure.apply(a),
@@ -16,16 +16,16 @@ pub fn apply_any(f: &Nir, a: Nir) -> NirKind {
     }
 }
 
-pub fn squash_textlit(
-    elts: impl Iterator<Item = InterpolatedTextContents<Nir>>,
-) -> Vec<InterpolatedTextContents<Nir>> {
+pub fn squash_textlit<'cx>(
+    elts: impl Iterator<Item = InterpolatedTextContents<Nir<'cx>>>,
+) -> Vec<InterpolatedTextContents<Nir<'cx>>> {
     use std::mem::replace;
     use InterpolatedTextContents::{Expr, Text};
 
-    fn inner(
-        elts: impl Iterator<Item = InterpolatedTextContents<Nir>>,
+    fn inner<'cx>(
+        elts: impl Iterator<Item = InterpolatedTextContents<Nir<'cx>>>,
         crnt_str: &mut String,
-        ret: &mut Vec<InterpolatedTextContents<Nir>>,
+        ret: &mut Vec<InterpolatedTextContents<Nir<'cx>>>,
     ) {
         for contents in elts {
             match contents {
@@ -81,22 +81,22 @@ where
     kvs
 }
 
-pub type Ret = NirKind;
+pub type Ret<'cx> = NirKind<'cx>;
 
-pub fn ret_nir(x: Nir) -> Ret {
+pub fn ret_nir<'cx>(x: Nir<'cx>) -> Ret<'cx> {
     x.into_kind()
 }
-pub fn ret_kind(x: NirKind) -> Ret {
+pub fn ret_kind<'cx>(x: NirKind<'cx>) -> Ret<'cx> {
     x
 }
-pub fn ret_ref(x: &Nir) -> Ret {
+pub fn ret_ref<'cx>(x: &Nir<'cx>) -> Ret<'cx> {
     x.kind().clone()
 }
-pub fn ret_op(x: OpKind<Nir>) -> Ret {
+pub fn ret_op<'cx>(x: OpKind<Nir<'cx>>) -> Ret<'cx> {
     NirKind::Op(x)
 }
 
-pub fn normalize_one_layer(expr: ExprKind<Nir>, env: &NzEnv) -> NirKind {
+pub fn normalize_one_layer<'cx>(expr: ExprKind<Nir<'cx>>) -> NirKind<'cx> {
     use NirKind::{
         Assert, Const, NEListLit, NEOptionalLit, Num, RecordLit, RecordType,
         UnionType,
@@ -106,15 +106,13 @@ pub fn normalize_one_layer(expr: ExprKind<Nir>, env: &NzEnv) -> NirKind {
         ExprKind::Var(..)
         | ExprKind::Lam(..)
         | ExprKind::Pi(..)
-        | ExprKind::Let(..) => {
+        | ExprKind::Let(..)
+        | ExprKind::Builtin(..) => {
             unreachable!("This case should have been handled in normalize_hir")
         }
 
         ExprKind::Const(c) => ret_kind(Const(c)),
         ExprKind::Num(l) => ret_kind(Num(l)),
-        ExprKind::Builtin(b) => {
-            ret_kind(NirKind::from_builtin_env(b, env.clone()))
-        }
         ExprKind::TextLit(elts) => {
             let tlit = TextLit::new(elts.into_iter());
             // Simplify bare interpolation
@@ -154,10 +152,22 @@ pub fn normalize_one_layer(expr: ExprKind<Nir>, env: &NzEnv) -> NirKind {
 }
 
 /// Normalize Hir into WHNF
-pub fn normalize_hir(env: &NzEnv, hir: &Hir) -> NirKind {
+pub fn normalize_hir<'cx>(env: &NzEnv<'cx>, hir: &Hir<'cx>) -> NirKind<'cx> {
     match hir.kind() {
+        HirKind::MissingVar(..) => unreachable!("ruled out by typechecking"),
         HirKind::Var(var) => env.lookup_val(*var),
-        HirKind::Import(hir, _) => normalize_hir(env, hir),
+        HirKind::Import(import) => {
+            let typed = env.cx()[import].unwrap_result();
+            normalize_hir(env, &typed.hir)
+        }
+        HirKind::ImportAlternative(alt, left, right) => {
+            let hir = if env.cx()[alt].unwrap_selected() {
+                left
+            } else {
+                right
+            };
+            normalize_hir(env, hir)
+        }
         HirKind::Expr(ExprKind::Lam(binder, annot, body)) => {
             let annot = annot.eval(env);
             NirKind::LamClosure {
@@ -178,9 +188,12 @@ pub fn normalize_hir(env: &NzEnv, hir: &Hir) -> NirKind {
             let val = val.eval(env);
             body.eval(env.insert_value(val, ())).kind().clone()
         }
+        HirKind::Expr(ExprKind::Builtin(b)) => {
+            NirKind::from_builtin_env(*b, env.clone())
+        }
         HirKind::Expr(e) => {
             let e = e.map_ref(|hir| hir.eval(env));
-            normalize_one_layer(e, env)
+            normalize_one_layer(e)
         }
     }
 }