|
9 | 9 | #[macro_use]
|
10 | 10 | mod zipmacro;
|
11 | 11 |
|
| 12 | +#[cfg(feature = "rayon")] |
12 | 13 | use std::mem::MaybeUninit;
|
13 | 14 | use alloc::vec::Vec;
|
14 | 15 |
|
@@ -811,6 +812,7 @@ where
|
811 | 812 | FoldWhile::Continue(acc)
|
812 | 813 | }
|
813 | 814 |
|
| 815 | + #[cfg(feature = "rayon")] |
814 | 816 | pub(crate) fn uninitalized_for_current_layout<T>(&self) -> Array<MaybeUninit<T>, D>
|
815 | 817 | {
|
816 | 818 | let is_f = self.prefer_f();
|
@@ -1079,17 +1081,27 @@ macro_rules! map_impl {
|
1079 | 1081 | ///
|
1080 | 1082 | /// If all inputs are c- or f-order respectively, that is preserved in the output.
|
1081 | 1083 | pub fn map_collect<R>(self, f: impl FnMut($($p::Item,)* ) -> R) -> Array<R, D> {
|
1082 |
| - // Make uninit result |
1083 |
| - let mut output = self.uninitalized_for_current_layout::<R>(); |
| 1084 | + self.map_collect_owned(f) |
| 1085 | + } |
1084 | 1086 |
|
1085 |
| - // Use partial to counts the number of filled elements, and can drop the right |
1086 |
| - // number of elements on unwinding (if it happens during apply/collect). |
| 1087 | + pub(crate) fn map_collect_owned<S, R>(self, f: impl FnMut($($p::Item,)* ) -> R) |
| 1088 | + -> ArrayBase<S, D> |
| 1089 | + where S: DataOwned<Elem = R> |
| 1090 | + { |
| 1091 | + // safe because: all elements are written before the array is completed |
| 1092 | + |
| 1093 | + let shape = self.dimension.clone().set_f(self.prefer_f()); |
| 1094 | + let output = <ArrayBase<S, D>>::build_uninit(shape, |output| { |
| 1095 | + // Use partial to count the number of filled elements, and can drop the right |
| 1096 | + // number of elements on unwinding (if it happens during apply/collect). |
| 1097 | + unsafe { |
| 1098 | + let output_view = output.cast::<R>(); |
| 1099 | + self.and(output_view) |
| 1100 | + .collect_with_partial(f) |
| 1101 | + .release_ownership(); |
| 1102 | + } |
| 1103 | + }); |
1087 | 1104 | unsafe {
|
1088 |
| - let output_view = output.raw_view_mut().cast::<R>(); |
1089 |
| - self.and(output_view) |
1090 |
| - .collect_with_partial(f) |
1091 |
| - .release_ownership(); |
1092 |
| - |
1093 | 1105 | output.assume_init()
|
1094 | 1106 | }
|
1095 | 1107 | }
|
|
0 commit comments