Upgrade rust/crates/rayon to 1.5.1 am: dd2305def6 am: 489816d978 am: 33ba451729 am: 1774cfefb8

Original change: https://android-review.googlesource.com/c/platform/external/rust/crates/rayon/+/1712353

Change-Id: I3d36cf9c98f91edb9e48670feaad490066e5d3e2

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
     }