correctly drop Array::IntoIter (#1773)

* 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

* eta reduction

Co-authored-by: Yann Hamdaoui <[email protected]>

* update comments and fix eta-reduction error

---------

Co-authored-by: Yann Hamdaoui <[email protected]>

Author: Radvendii

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