rayon/iter/collect/
consumer.rs

1use super::super::plumbing::*;
2use crate::SendPtr;
3use std::marker::PhantomData;
4use std::ptr;
5
6pub(super) struct CollectConsumer<'c, T: Send> {
7    /// See `CollectResult` for explanation of why this is not a slice
8    start: SendPtr<T>,
9    len: usize,
10    marker: PhantomData<&'c mut T>,
11}
12
13impl<T: Send> CollectConsumer<'_, T> {
14    /// Create a collector for `len` items in the unused capacity of the vector.
15    pub(super) fn appender(vec: &mut Vec<T>, len: usize) -> CollectConsumer<'_, T> {
16        let start = vec.len();
17        assert!(vec.capacity() - start >= len);
18
19        // SAFETY: We already made sure to have the additional space allocated.
20        // The pointer is derived from `Vec` directly, not through a `Deref`,
21        // so it has provenance over the whole allocation.
22        unsafe { CollectConsumer::new(vec.as_mut_ptr().add(start), len) }
23    }
24}
25
26impl<'c, T: Send + 'c> CollectConsumer<'c, T> {
27    /// The target memory is considered uninitialized, and will be
28    /// overwritten without reading or dropping existing values.
29    unsafe fn new(start: *mut T, len: usize) -> Self {
30        CollectConsumer {
31            start: SendPtr(start),
32            len,
33            marker: PhantomData,
34        }
35    }
36}
37
38/// CollectResult represents an initialized part of the target slice.
39///
40/// This is a proxy owner of the elements in the slice; when it drops,
41/// the elements will be dropped, unless its ownership is released before then.
42#[must_use]
43pub(super) struct CollectResult<'c, T> {
44    /// This pointer and length has the same representation as a slice,
45    /// but retains the provenance of the entire array so that we can merge
46    /// these regions together in `CollectReducer`.
47    start: SendPtr<T>,
48    total_len: usize,
49    /// The current initialized length after `start`
50    initialized_len: usize,
51    /// Lifetime invariance guarantees that the data flows from consumer to result,
52    /// especially for the `scope_fn` callback in `Collect::with_consumer`.
53    invariant_lifetime: PhantomData<&'c mut &'c mut [T]>,
54}
55
56unsafe impl<'c, T> Send for CollectResult<'c, T> where T: Send {}
57
58impl<'c, T> CollectResult<'c, T> {
59    /// The current length of the collect result
60    pub(super) fn len(&self) -> usize {
61        self.initialized_len
62    }
63
64    /// Release ownership of the slice of elements, and return the length
65    pub(super) fn release_ownership(mut self) -> usize {
66        let ret = self.initialized_len;
67        self.initialized_len = 0;
68        ret
69    }
70}
71
72impl<'c, T> Drop for CollectResult<'c, T> {
73    fn drop(&mut self) {
74        // Drop the first `self.initialized_len` elements, which have been recorded
75        // to be initialized by the folder.
76        unsafe {
77            ptr::drop_in_place(ptr::slice_from_raw_parts_mut(
78                self.start.0,
79                self.initialized_len,
80            ));
81        }
82    }
83}
84
85impl<'c, T: Send + 'c> Consumer<T> for CollectConsumer<'c, T> {
86    type Folder = CollectResult<'c, T>;
87    type Reducer = CollectReducer;
88    type Result = CollectResult<'c, T>;
89
90    fn split_at(self, index: usize) -> (Self, Self, CollectReducer) {
91        let CollectConsumer { start, len, .. } = self;
92
93        // Produce new consumers.
94        // SAFETY: This assert checks that `index` is a valid offset for `start`
95        unsafe {
96            assert!(index <= len);
97            (
98                CollectConsumer::new(start.0, index),
99                CollectConsumer::new(start.0.add(index), len - index),
100                CollectReducer,
101            )
102        }
103    }
104
105    fn into_folder(self) -> Self::Folder {
106        // Create a result/folder that consumes values and writes them
107        // into the region after start. The initial result has length 0.
108        CollectResult {
109            start: self.start,
110            total_len: self.len,
111            initialized_len: 0,
112            invariant_lifetime: PhantomData,
113        }
114    }
115
116    fn full(&self) -> bool {
117        false
118    }
119}
120
121impl<'c, T: Send + 'c> Folder<T> for CollectResult<'c, T> {
122    type Result = Self;
123
124    fn consume(mut self, item: T) -> Self {
125        assert!(
126            self.initialized_len < self.total_len,
127            "too many values pushed to consumer"
128        );
129
130        // SAFETY: The assert above is a bounds check for this write, and we
131        // avoid assignment here so we do not drop an uninitialized T.
132        unsafe {
133            // Write item and increase the initialized length
134            self.start.0.add(self.initialized_len).write(item);
135            self.initialized_len += 1;
136        }
137
138        self
139    }
140
141    fn complete(self) -> Self::Result {
142        // NB: We don't explicitly check that the local writes were complete,
143        // but Collect will assert the total result length in the end.
144        self
145    }
146
147    fn full(&self) -> bool {
148        false
149    }
150}
151
152/// Pretend to be unindexed for `special_collect_into_vec`,
153/// but we should never actually get used that way...
154impl<'c, T: Send + 'c> UnindexedConsumer<T> for CollectConsumer<'c, T> {
155    fn split_off_left(&self) -> Self {
156        unreachable!("CollectConsumer must be indexed!")
157    }
158    fn to_reducer(&self) -> Self::Reducer {
159        CollectReducer
160    }
161}
162
163/// CollectReducer combines adjacent chunks; the result must always
164/// be contiguous so that it is one combined slice.
165pub(super) struct CollectReducer;
166
167impl<'c, T> Reducer<CollectResult<'c, T>> for CollectReducer {
168    fn reduce(
169        self,
170        mut left: CollectResult<'c, T>,
171        right: CollectResult<'c, T>,
172    ) -> CollectResult<'c, T> {
173        // Merge if the CollectResults are adjacent and in left to right order
174        // else: drop the right piece now and total length will end up short in the end,
175        // when the correctness of the collected result is asserted.
176        unsafe {
177            let left_end = left.start.0.add(left.initialized_len);
178            if left_end == right.start.0 {
179                left.total_len += right.total_len;
180                left.initialized_len += right.release_ownership();
181            }
182            left
183        }
184    }
185}
rayon/iter/collect/consumer.rs

rayon/iter/collect/
consumer.rs
