Skip to content

Commit

Permalink
correctly drop Array::IntoIter
Browse files Browse the repository at this point in the history
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
  • Loading branch information
Radvendii committed Jan 23, 2024
1 parent 4f84cf0 commit 02a9d1c
Showing 1 changed file with 53 additions and 20 deletions.
73 changes: 53 additions & 20 deletions core/src/term/array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
Expand All @@ -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;
Expand All @@ -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
Expand Down

0 comments on commit 02a9d1c

Please sign in to comment.