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