1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
use super::super::plumbing::*;
use crate::SendPtr;
use std::marker::PhantomData;
use std::ptr;
use std::slice;

pub(super) struct CollectConsumer<'c, T: Send> {
    /// See `CollectResult` for explanation of why this is not a slice
    start: SendPtr<T>,
    len: usize,
    marker: PhantomData<&'c mut T>,
}

impl<T: Send> CollectConsumer<'_, T> {
    /// Create a collector for `len` items in the unused capacity of the vector.
    pub(super) fn appender(vec: &mut Vec<T>, len: usize) -> CollectConsumer<'_, T> {
        let start = vec.len();
        assert!(vec.capacity() - start >= len);

        // SAFETY: We already made sure to have the additional space allocated.
        // The pointer is derived from `Vec` directly, not through a `Deref`,
        // so it has provenance over the whole allocation.
        unsafe { CollectConsumer::new(vec.as_mut_ptr().add(start), len) }
    }
}

impl<'c, T: Send + 'c> CollectConsumer<'c, T> {
    /// The target memory is considered uninitialized, and will be
    /// overwritten without reading or dropping existing values.
    unsafe fn new(start: *mut T, len: usize) -> Self {
        CollectConsumer {
            start: SendPtr(start),
            len,
            marker: PhantomData,
        }
    }
}

/// CollectResult represents an initialized part of the target slice.
///
/// This is a proxy owner of the elements in the slice; when it drops,
/// the elements will be dropped, unless its ownership is released before then.
#[must_use]
pub(super) struct CollectResult<'c, T> {
    /// This pointer and length has the same representation as a slice,
    /// but retains the provenance of the entire array so that we can merge
    /// these regions together in `CollectReducer`.
    start: SendPtr<T>,
    total_len: usize,
    /// The current initialized length after `start`
    initialized_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]>,
}

unsafe impl<'c, T> Send for CollectResult<'c, T> where T: Send {}

impl<'c, T> CollectResult<'c, T> {
    /// The current length of the collect result
    pub(super) fn len(&self) -> usize {
        self.initialized_len
    }

    /// Release ownership of the slice of elements, and return the length
    pub(super) fn release_ownership(mut self) -> usize {
        let ret = self.initialized_len;
        self.initialized_len = 0;
        ret
    }
}

impl<'c, T> Drop for CollectResult<'c, T> {
    fn drop(&mut self) {
        // Drop the first `self.initialized_len` elements, which have been recorded
        // to be initialized by the folder.
        unsafe {
            ptr::drop_in_place(slice::from_raw_parts_mut(
                self.start.0,
                self.initialized_len,
            ));
        }
    }
}

impl<'c, T: Send + 'c> Consumer<T> for CollectConsumer<'c, T> {
    type Folder = CollectResult<'c, T>;
    type Reducer = CollectReducer;
    type Result = CollectResult<'c, T>;

    fn split_at(self, index: usize) -> (Self, Self, CollectReducer) {
        let CollectConsumer { start, len, .. } = self;

        // Produce new consumers.
        // SAFETY: This assert checks that `index` is a valid offset for `start`
        unsafe {
            assert!(index <= len);
            (
                CollectConsumer::new(start.0, index),
                CollectConsumer::new(start.0.add(index), len - index),
                CollectReducer,
            )
        }
    }

    fn into_folder(self) -> Self::Folder {
        // Create a result/folder that consumes values and writes them
        // into the region after start. The initial result has length 0.
        CollectResult {
            start: self.start,
            total_len: self.len,
            initialized_len: 0,
            invariant_lifetime: PhantomData,
        }
    }

    fn full(&self) -> bool {
        false
    }
}

impl<'c, T: Send + 'c> Folder<T> for CollectResult<'c, T> {
    type Result = Self;

    fn consume(mut self, item: T) -> Self {
        assert!(
            self.initialized_len < self.total_len,
            "too many values pushed to consumer"
        );

        // SAFETY: The assert above is a bounds check for this write, and we
        // avoid assignment here so we do not drop an uninitialized T.
        unsafe {
            // Write item and increase the initialized length
            self.start.0.add(self.initialized_len).write(item);
            self.initialized_len += 1;
        }

        self
    }

    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
    }

    fn full(&self) -> bool {
        false
    }
}

/// Pretend to be unindexed for `special_collect_into_vec`,
/// but we should never actually get used that way...
impl<'c, T: Send + 'c> UnindexedConsumer<T> for CollectConsumer<'c, T> {
    fn split_off_left(&self) -> Self {
        unreachable!("CollectConsumer must be indexed!")
    }
    fn to_reducer(&self) -> Self::Reducer {
        CollectReducer
    }
}

/// CollectReducer combines adjacent chunks; the result must always
/// be contiguous so that it is one combined slice.
pub(super) struct CollectReducer;

impl<'c, T> Reducer<CollectResult<'c, T>> for CollectReducer {
    fn reduce(
        self,
        mut left: CollectResult<'c, T>,
        right: CollectResult<'c, T>,
    ) -> CollectResult<'c, T> {
        // 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.
        unsafe {
            let left_end = left.start.0.add(left.initialized_len);
            if left_end == right.start.0 {
                left.total_len += right.total_len;
                left.initialized_len += right.release_ownership();
            }
            left
        }
    }
}