Skip to content

Commit 82150cb

Browse files
committed
FEAT: Add internal map_collect method that collects to any array storage
This generalizes .map_collect() so that it supports any DataOwned array storage. It is a crate internal method initially - it can be exported later with more experience.
1 parent 00a8d48 commit 82150cb

File tree

1 file changed

+21
-9
lines changed

1 file changed

+21
-9
lines changed

src/zip/mod.rs

+21-9
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
#[macro_use]
1010
mod zipmacro;
1111

12+
#[cfg(feature = "rayon")]
1213
use std::mem::MaybeUninit;
1314
use alloc::vec::Vec;
1415

@@ -811,6 +812,7 @@ where
811812
FoldWhile::Continue(acc)
812813
}
813814

815+
#[cfg(feature = "rayon")]
814816
pub(crate) fn uninitalized_for_current_layout<T>(&self) -> Array<MaybeUninit<T>, D>
815817
{
816818
let is_f = self.prefer_f();
@@ -1079,17 +1081,27 @@ macro_rules! map_impl {
10791081
///
10801082
/// If all inputs are c- or f-order respectively, that is preserved in the output.
10811083
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+
}
10841086

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+
});
10871104
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-
10931105
output.assume_init()
10941106
}
10951107
}

0 commit comments

Comments
 (0)