correctly drop Array::IntoIter

If you called Array::into_iter() you could leak elements in two circumstances:
1. If you did not iterate through all elements before dropping, since they had been converted to `ManuallyDrop`
2. If the Array had multiple references when the iterator was created, but the other references were dropped before the iterator was dropped, then the existence of the iterator would prevent those Rcs from dropping their contents, but then because the iterator is of type ManuallyDrop, it wouldn't drop those elements either

Author: Radvendii

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: 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] {
+        //
+        // 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