Normalize `with` by mutation.

This is Cow-style mutation: we clone only what's shared and then mutate
it. This it more legible and more efficient than the immutable version.

3 files changed, 42 insertions, 5 deletions

diff --git a/dhall/src/semantics/nze/lazy.rs b/dhall/src/semantics/nze/lazy.rs
index d361313..550f69a 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,7 @@ where
 {
     type Target = Tgt;
     fn deref(&self) -> &Self::Target {
-        self.tgt.get_or_init(|| {
-            let src = self.src.take().unwrap();
-            src.eval()
-        })
+        self.force()
     }
 }
 
diff --git a/dhall/src/semantics/nze/nir.rs b/dhall/src/semantics/nze/nir.rs
index 12f1b14..95eeba1 100644
--- a/dhall/src/semantics/nze/nir.rs
+++ b/dhall/src/semantics/nze/nir.rs
@@ -138,6 +138,24 @@ impl Nir {
         self.0.kind()
     }
 
+    /// 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 {
+        if Rc::get_mut(&mut self.0).is_none() {
+            // Clone self
+            let kind = self.kind().clone();
+            *self = Nir::from_kind(kind);
+        }
+        Rc::get_mut(&mut self.0).unwrap().kind_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(int) => int.into_kind(),
+            Err(rc) => rc.kind().clone(),
+        }
+    }
+
     pub fn to_type(&self, u: impl Into<Universe>) -> Type {
         Type::new(self.clone(), u.into())
     }
@@ -291,6 +309,12 @@ impl NirInternal {
     fn kind(&self) -> &NirKind {
         &self.kind
     }
+    fn kind_mut(&mut self) -> &mut NirKind {
+        self.kind.get_mut()
+    }
+    fn into_kind(self) -> NirKind {
+        self.kind.into_inner()
+    }
 }
 
 impl NirKind {
diff --git a/dhall/src/semantics/nze/normalize.rs b/dhall/src/semantics/nze/normalize.rs
index 3b40fac..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