3 files changed, 52 insertions, 42 deletions

diff --git a/dhall/src/semantics/nze/lazy.rs b/dhall/src/semantics/nze/lazy.rs
index d361313..d3b5c8d 100644
--- a/dhall/src/semantics/nze/lazy.rs
+++ b/dhall/src/semantics/nze/lazy.rs
@@ -35,6 +35,22 @@ where
         let _ = lazy.tgt.set(tgt);
         lazy
     }
+
+    pub fn force(&self) -> &Tgt {
+        self.tgt.get_or_init(|| {
+            let src = self.src.take().unwrap();
+            src.eval()
+        })
+    }
+
+    pub fn get_mut(&mut self) -> &mut Tgt {
+        self.force();
+        self.tgt.get_mut().unwrap()
+    }
+    pub fn into_inner(self) -> Tgt {
+        self.force();
+        self.tgt.into_inner().unwrap()
+    }
 }
 
 impl<Src, Tgt> Deref for Lazy<Src, Tgt>
@@ -43,10 +59,18 @@ where
 {
     type Target = Tgt;
     fn deref(&self) -> &Self::Target {
-        self.tgt.get_or_init(|| {
-            let src = self.src.take().unwrap();
-            src.eval()
-        })
+        self.force()
+    }
+}
+
+/// This implementation evaluates before cloning, because we can't clone the contents of a `Cell`.
+impl<Src, Tgt> Clone for Lazy<Src, Tgt>
+where
+    Src: Eval<Tgt>,
+    Tgt: Clone,
+{
+    fn clone(&self) -> Self {
+        Self::new_completed(self.force().clone())
     }
 }
 
diff --git a/dhall/src/semantics/nze/nir.rs b/dhall/src/semantics/nze/nir.rs
index 12f1b14..7bda836 100644
--- a/dhall/src/semantics/nze/nir.rs
+++ b/dhall/src/semantics/nze/nir.rs
@@ -14,18 +14,13 @@ use crate::syntax::{
 use crate::ToExprOptions;
 
 /// Stores a possibly unevaluated value. Gets (partially) normalized on-demand, sharing computation
-/// automatically. Uses a Rc<RefCell> to share computation.
+/// automatically. Uses a Rc<OnceCell> to share computation.
 /// If you compare for equality two `Nir`s, then equality will be up to alpha-equivalence
 /// (renaming of bound variables) and beta-equivalence (normalization). It will recursively
 /// normalize as needed.
-/// Stands for "Normalized intermediate representation"
+/// Stands for "Normalized Intermediate Representation"
 #[derive(Clone)]
-pub struct Nir(Rc<NirInternal>);
-
-#[derive(Debug)]
-struct NirInternal {
-    kind: lazy::Lazy<Thunk, NirKind>,
-}
+pub struct Nir(Rc<lazy::Lazy<Thunk, NirKind>>);
 
 /// An unevaluated subexpression
 #[derive(Debug, Clone)]
@@ -101,17 +96,17 @@ pub enum NirKind {
 impl Nir {
     /// Construct a Nir from a completely unnormalized expression.
     pub fn new_thunk(env: NzEnv, hir: Hir) -> Nir {
-        NirInternal::from_thunk(Thunk::new(env, hir)).into_nir()
+        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();
-        NirInternal::from_thunk(Thunk::from_partial_expr(env, e)).into_nir()
+        Nir(Rc::new(lazy::Lazy::new(Thunk::from_partial_expr(env, e))))
     }
     /// Make a Nir from a NirKind
     pub fn from_kind(v: NirKind) -> Nir {
-        NirInternal::from_whnf(v).into_nir()
+        Nir(Rc::new(lazy::Lazy::new_completed(v)))
     }
     pub fn from_const(c: Const) -> Self {
         Self::from_kind(NirKind::Const(c))
@@ -135,7 +130,20 @@ impl Nir {
 
     /// This is what you want if you want to pattern-match on the value.
     pub fn kind(&self) -> &NirKind {
-        self.0.kind()
+        &*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 {
+        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 {
+        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 {
@@ -273,26 +281,6 @@ impl Nir {
     }
 }
 
-impl NirInternal {
-    fn from_whnf(k: NirKind) -> Self {
-        NirInternal {
-            kind: lazy::Lazy::new_completed(k),
-        }
-    }
-    fn from_thunk(th: Thunk) -> Self {
-        NirInternal {
-            kind: lazy::Lazy::new(th),
-        }
-    }
-    fn into_nir(self) -> Nir {
-        Nir(Rc::new(self))
-    }
-
-    fn kind(&self) -> &NirKind {
-        &self.kind
-    }
-}
-
 impl NirKind {
     pub fn into_nir(self) -> Nir {
         Nir::from_kind(self)
@@ -448,13 +436,11 @@ impl std::cmp::Eq for Closure {}
 
 impl std::fmt::Debug for Nir {
     fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let vint: &NirInternal = &self.0;
-        let kind = vint.kind();
-        if let NirKind::Const(c) = kind {
+        if let NirKind::Const(c) = self.kind() {
             return write!(fmt, "{:?}", c);
         }
         let mut x = fmt.debug_struct("Nir");
-        x.field("kind", kind);
+        x.field("kind", self.kind());
         x.finish()
     }
 }
diff --git a/dhall/src/semantics/nze/normalize.rs b/dhall/src/semantics/nze/normalize.rs
index d042f3f..62efc5f 100644
--- a/dhall/src/semantics/nze/normalize.rs
+++ b/dhall/src/semantics/nze/normalize.rs
@@ -84,7 +84,7 @@ where
 pub type Ret = NirKind;
 
 pub fn ret_nir(x: Nir) -> Ret {
-    ret_ref(&x)
+    x.into_kind()
 }
 pub fn ret_kind(x: NirKind) -> Ret {
     x
@@ -144,7 +144,7 @@ pub fn normalize_one_layer(expr: ExprKind<Nir>, env: &NzEnv) -> NirKind {
         ExprKind::UnionType(kvs) => {
             ret_kind(UnionType(kvs.into_iter().collect()))
         }
-        ExprKind::Op(ref op) => normalize_operation(op),
+        ExprKind::Op(op) => normalize_operation(op),
         ExprKind::Annot(x, _) => ret_nir(x),
         ExprKind::Assert(x) => ret_kind(Assert(x)),
         ExprKind::Import(..) => {