From b652cbe230f05c45dd825c1ba43b737939dd7199 Mon Sep 17 00:00:00 2001
From: Taeer Bar-Yam <taeer.bar-yam@tweag.io>
Date: Mon, 6 Nov 2023 14:28:10 -0500
Subject: [PATCH] correctly drop Array::IntoIter

---
 core/src/term/array.rs | 73 ++++++++++++++++++++++++++++++------------
 1 file changed, 53 insertions(+), 20 deletions(-)

diff --git a/core/src/term/array.rs b/core/src/term/array.rs
index bcc9d99d17..24b203a675 100644
--- a/core/src/term/array.rs
+++ b/core/src/term/array.rs
@@ -157,29 +157,50 @@ impl IntoIterator for Array {
         // to the inner slice, then we can:
         // - Drop the elements outside our inner view,
         // - Move out the elements inside out inner view.
+        // - Drop the rest of the elements when we're dropped
         // Otherwise, we clone everything.
-
-        unsafe {
-            let mut inner = transmute::<Rc<[RichTerm]>, Rc<[ManuallyDrop<RichTerm>]>>(self.inner);
-            let idx = self.start;
-            let end = self.end;
-
-            if let Some(slice) = Rc::get_mut(&mut inner) {
-                for term in &mut slice[..idx] {
+        //
+        // If we start as a shared reference, we could become the only reference
+        // later, but we clone everything anyways, so we don't end up in an
+        // invalid in-between state where some terms have been freed manually
+        // and others haven't.
+
+        let inner = if Rc::strong_count(&self.inner) != 1 || Rc::weak_count(&self.inner) != 0 {
+            IntoIterInner::Shared(self.inner)
+        } else {
+            unsafe {
+                let mut inner =
+                    transmute::<Rc<[RichTerm]>, Rc<[ManuallyDrop<RichTerm>]>>(self.inner);
+                let slice =
+                    Rc::get_mut(&mut inner).expect("non-unique Rc after checking for uniqueness");
+                for term in &mut slice[..self.start] {
                     ManuallyDrop::drop(term)
                 }
-                for term in &mut slice[end..] {
+                for term in &mut slice[self.end..] {
                     ManuallyDrop::drop(term)
                 }
-            }
 
-            IntoIter { inner, idx, end }
+                IntoIterInner::Owned(inner)
+            }
+        };
+        IntoIter {
+            inner,
+            idx: self.start,
+            end: self.end,
         }
     }
 }
 
+enum IntoIterInner {
+    Shared(Rc<[RichTerm]>),
+    // This shouldn't really be an Rc. It should only ever have one reference.
+    // There may be a way to rewrite this once unsized locals are stabilized.
+    // But for now, there is no good way to repackage the inner array, so we
+    // keep it in an Rc.
+    Owned(Rc<[ManuallyDrop<RichTerm>]>),
+}
 pub struct IntoIter {
-    inner: Rc<[ManuallyDrop<RichTerm>]>,
+    inner: IntoIterInner,
     idx: usize,
     end: usize,
 }
@@ -191,14 +212,14 @@ impl Iterator for IntoIter {
         if self.idx == self.end {
             None
         } else {
-            let term = match Rc::get_mut(&mut self.inner) {
-                None => self.inner[..self.end]
-                    .get(self.idx)
-                    .cloned()
-                    .map(ManuallyDrop::into_inner),
-                Some(slice) => slice[..self.end]
-                    .get_mut(self.idx)
-                    .map(|t| unsafe { ManuallyDrop::take(t) }),
+            let term = match &mut self.inner {
+                IntoIterInner::Shared(inner) => inner.get(self.idx).cloned(),
+                IntoIterInner::Owned(inner) => unsafe {
+                    Rc::get_mut(inner)
+                        .expect("non-unique Rc after checking for uniqueness")
+                        .get_mut(self.idx)
+                        .map(|t| ManuallyDrop::take(t))
+                },
             };
 
             self.idx += 1;
@@ -207,6 +228,18 @@ impl Iterator for IntoIter {
     }
 }
 
+impl Drop for IntoIter {
+    fn drop(&mut self) {
+        // drop the rest of the items we haven't iterated through
+        if let IntoIterInner::Owned(inner) = &mut self.inner {
+            let inner = Rc::get_mut(inner).expect("non-unique Rc after checking for uniqueness");
+            for term in &mut inner[self.idx..self.end] {
+                unsafe { ManuallyDrop::drop(term) }
+            }
+        }
+    }
+}
+
 impl ExactSizeIterator for IntoIter {
     fn len(&self) -> usize {
         self.end - self.idx