correctly drop Array::IntoIter

core/src/term/array.rs

@@ -157,29 +157,45 @@ 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: Rc<[ManuallyDrop<RichTerm>]> = transmute(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] {
+        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,22 +207,40 @@ 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]
+            let term = match &mut self.inner {
+                IntoIterInner::Shared(inner) => inner.get(self.idx).cloned(),
+                IntoIterInner::Owned(inner) => Rc::get_mut(inner)
+                    .expect("non-unique Rc after checking for uniqueness")
                     .get_mut(self.idx)
+                    // SAFETY: We already checcked that we have the only
+                    // reference to inner, and after this we increment idx, and
+                    // we never access elements with indexes less than idx
                     .map(|t| unsafe { ManuallyDrop::take(t) }),
             };
-
             self.idx += 1;
             term
         }
     }
 }
 
+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] {
+                // SAFETY: Wwe already checked that we have the only reference
+                // to inner, and once we are done dropping the remaining
+                // elements, `inner` will get dropped and there will be no
+                // remaining references.
+                unsafe {
+                    let _ = ManuallyDrop::take(term);
+                }
+            }
+        }
+    }
+}
+
 impl ExactSizeIterator for IntoIter {
     fn len(&self) -> usize {
         self.end - self.idx