diff options
Diffstat (limited to 'src/iter/collect/consumer.rs')
-rw-r--r-- | src/iter/collect/consumer.rs | 72 |
1 files changed, 36 insertions, 36 deletions
diff --git a/src/iter/collect/consumer.rs b/src/iter/collect/consumer.rs index 689f29c..3a8eea0 100644 --- a/src/iter/collect/consumer.rs +++ b/src/iter/collect/consumer.rs @@ -1,26 +1,18 @@ use super::super::plumbing::*; use std::marker::PhantomData; +use std::mem::MaybeUninit; use std::ptr; use std::slice; pub(super) struct CollectConsumer<'c, T: Send> { /// A slice covering the target memory, not yet initialized! - target: &'c mut [T], -} - -pub(super) struct CollectFolder<'c, T: Send> { - /// The folder writes into `result` and must extend the result - /// up to exactly this number of elements. - final_len: usize, - - /// The current written-to part of our slice of the target - result: CollectResult<'c, T>, + target: &'c mut [MaybeUninit<T>], } impl<'c, T: Send + 'c> CollectConsumer<'c, T> { /// The target memory is considered uninitialized, and will be /// overwritten without reading or dropping existing values. - pub(super) fn new(target: &'c mut [T]) -> Self { + pub(super) fn new(target: &'c mut [MaybeUninit<T>]) -> Self { CollectConsumer { target } } } @@ -31,8 +23,12 @@ impl<'c, T: Send + 'c> CollectConsumer<'c, T> { /// the elements will be dropped, unless its ownership is released before then. #[must_use] pub(super) struct CollectResult<'c, T> { - start: *mut T, + /// A slice covering the target memory, initialized up to our separate `len`. + target: &'c mut [MaybeUninit<T>], + /// The current initialized length in `target` len: usize, + /// Lifetime invariance guarantees that the data flows from consumer to result, + /// especially for the `scope_fn` callback in `Collect::with_consumer`. invariant_lifetime: PhantomData<&'c mut &'c mut [T]>, } @@ -57,13 +53,15 @@ impl<'c, T> Drop for CollectResult<'c, T> { // Drop the first `self.len` elements, which have been recorded // to be initialized by the folder. unsafe { - ptr::drop_in_place(slice::from_raw_parts_mut(self.start, self.len)); + // TODO: use `MaybeUninit::slice_as_mut_ptr` + let start = self.target.as_mut_ptr() as *mut T; + ptr::drop_in_place(slice::from_raw_parts_mut(start, self.len)); } } } impl<'c, T: Send + 'c> Consumer<T> for CollectConsumer<'c, T> { - type Folder = CollectFolder<'c, T>; + type Folder = CollectResult<'c, T>; type Reducer = CollectReducer; type Result = CollectResult<'c, T>; @@ -80,16 +78,13 @@ impl<'c, T: Send + 'c> Consumer<T> for CollectConsumer<'c, T> { ) } - fn into_folder(self) -> CollectFolder<'c, T> { - // Create a folder that consumes values and writes them + fn into_folder(self) -> Self::Folder { + // Create a result/folder that consumes values and writes them // into target. The initial result has length 0. - CollectFolder { - final_len: self.target.len(), - result: CollectResult { - start: self.target.as_mut_ptr(), - len: 0, - invariant_lifetime: PhantomData, - }, + CollectResult { + target: self.target, + len: 0, + invariant_lifetime: PhantomData, } } @@ -98,19 +93,19 @@ impl<'c, T: Send + 'c> Consumer<T> for CollectConsumer<'c, T> { } } -impl<'c, T: Send + 'c> Folder<T> for CollectFolder<'c, T> { - type Result = CollectResult<'c, T>; +impl<'c, T: Send + 'c> Folder<T> for CollectResult<'c, T> { + type Result = Self; - fn consume(mut self, item: T) -> CollectFolder<'c, T> { - if self.result.len >= self.final_len { - panic!("too many values pushed to consumer"); - } + fn consume(mut self, item: T) -> Self { + let dest = self + .target + .get_mut(self.len) + .expect("too many values pushed to consumer"); - // Compute target pointer and write to it, and - // extend the current result by one element + // Write item and increase the initialized length unsafe { - self.result.start.add(self.result.len).write(item); - self.result.len += 1; + dest.as_mut_ptr().write(item); + self.len += 1; } self @@ -119,7 +114,7 @@ impl<'c, T: Send + 'c> Folder<T> for CollectFolder<'c, T> { fn complete(self) -> Self::Result { // NB: We don't explicitly check that the local writes were complete, // but Collect will assert the total result length in the end. - self.result + self } fn full(&self) -> bool { @@ -151,8 +146,13 @@ impl<'c, T> Reducer<CollectResult<'c, T>> for CollectReducer { // Merge if the CollectResults are adjacent and in left to right order // else: drop the right piece now and total length will end up short in the end, // when the correctness of the collected result is asserted. - if left.start.wrapping_add(left.len) == right.start { - left.len += right.release_ownership(); + let left_end = left.target[left.len..].as_ptr(); + if left_end == right.target.as_ptr() { + let len = left.len + right.release_ownership(); + unsafe { + left.target = slice::from_raw_parts_mut(left.target.as_mut_ptr(), len); + } + left.len = len; } left } |