correctly drop Array::IntoIter

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