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.

4 files changed, 66 insertions, 57 deletions

diff --git a/dhall/src/operations/normalization.rs b/dhall/src/operations/normalization.rs
index e9ae825..28375b2 100644
--- a/dhall/src/operations/normalization.rs
+++ b/dhall/src/operations/normalization.rs
@@ -296,62 +296,34 @@ pub fn normalize_operation(opkind: OpKind<Nir>) -> Ret {
             )),
             _ => ret_op(ProjectionByExpr(v, t)),
         },
-        With(record, labels, expr) => {
-            let mut current_nir: Option<Nir> = Some(record.clone());
-            let mut visited: Vec<(&Label, HashMap<Label, Nir>)> = Vec::new();
-            let mut to_create = Vec::new();
-
-            // To be used when an abstract entry is reached
-            let mut abstract_nir = None;
-
-            for label in labels {
-                match current_nir {
-                    None => to_create.push(label.clone()),
-                    Some(nir) => match nir.kind() {
-                        RecordLit(kvs) => {
-                            current_nir = kvs.get(label).cloned();
-                            visited.push((label, kvs.clone()));
-                        }
-                        // Handle partially abstract case
-                        _ => {
-                            abstract_nir = Some(nir);
-                            to_create.push(label.clone());
-                            current_nir = None;
-                        }
-                    },
+        With(mut record, labels, expr) => {
+            let mut labels = labels.into_iter();
+            let mut current = &mut record;
+            // We dig through the current record with the provided labels.
+            while let RecordLit(kvs) = current.kind_mut() {
+                if let Some(label) = labels.next() {
+                    // Get existing entry or insert empty record into it.
+                    let nir = kvs.entry(label).or_insert_with(|| {
+                        Nir::from_kind(RecordLit(HashMap::new()))
+                    });
+                    // Disgusting, but the normal assignment works with -Zpolonius, so this
+                    // is safe. See https://github.com/rust-lang/rust/issues/70255 .
+                    current = unsafe { &mut *(nir as *mut _) };
+                } else {
+                    break;
                 }
             }
 
-            // Create Nir for record bottom up
-            let mut nir = expr.clone();
-
-            match abstract_nir {
-                // No abstract nir, creating singleton records
-                None => {
-                    while let Some(label) = to_create.pop() {
-                        let rec =
-                            RecordLit(once((label.clone(), nir)).collect());
-                        nir = Nir::from_kind(rec);
-                    }
-                }
-                // Abstract nir, creating with op
-                Some(abstract_nir) => {
-                    nir = Nir::from_kind(Op(OpKind::With(
-                        abstract_nir,
-                        to_create,
-                        nir,
-                    )));
-                }
-            }
-
-            // Update visited records
-            while let Some((label, mut kvs)) = visited.pop() {
-                kvs.insert(label.clone(), nir);
-                let rec = RecordLit(kvs);
-                nir = Nir::from_kind(rec);
-            }
+            // If there are still some fields to dig through, we need to create a `with` expression
+            // with the remaining fields.
+            let labels: Vec<_> = labels.collect();
+            *current = if labels.is_empty() {
+                expr
+            } else {
+                Nir::from_kind(Op(OpKind::With(current.clone(), labels, expr)))
+            };
 
-            ret_nir(nir)
+            ret_nir(record)
         }
         Completion(..) => {
             unreachable!("This case should have been handled in resolution")
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