Reduce mapped array execution overhead

Cargo.toml

@@ -11,7 +11,7 @@ rustymath = { path = ".", version = ">=0.1.60, <1", default-features = false }
 rustymath-buffer = { path = "buffer", version = ">=0.1.39, <1" }
 pyo3 = ">=0.29"
 pyo3-build-config = ">=0.29"
-rustymath-core = { path = "core", version = ">=0.1.60, <1", default-features = false }
+rustymath-core = { path = "core", version = ">=0.1.61, <1", default-features = false }
 rustymath-kernels = { path = "kernels", version = ">=0.1.59, <1" }
 rustymath-timing = { path = "timing", version = ">=0.1.56, <1" }
 half = { version = ">=2.7.1, <3", features = ["num-traits"] }

buffer/src/selection.rs

@@ -35,9 +35,7 @@ pub fn key_names(obj: &Bound<'_, PyAny>) -> Option<Vec<String>> {
 /// One Python selector, independent of the selectors on the other axes. Keys read by `structure::key_selector`, so a missing key reads the fill, as in BPL.
 pub fn selector(layout: &Layout, axis: usize, obj: &Bound<'_, PyAny>, array: impl FnOnce() -> PyResult<(Selector, bool)>) -> PyResult<(Selector, bool)> {
     let size = *layout.shape().get(axis).ok_or_else(|| PyIndexError::new_err("too many axis selectors"))?;
-    if let Ok(slice) = obj.cast::<PySlice>() {
-        return slice_spec(slice)?.selector(layout, axis).map(|s| (s, false)).map_err(error);
-    }
+    if let Ok(slice) = obj.cast::<PySlice>() { return slice_spec(slice)?.selector(layout, axis).map(|s| (s, false)).map_err(error); }
     if let Some(names) = key_names(obj) {
         let one = obj.is_instance_of::<PyString>();
         let within = (!one).then(|| Layout::from([names.len()]));
@@ -75,8 +73,9 @@ pub fn parts<'py>(index: &Bound<'py, PyAny>) -> PyResult<Vec<IndexPart<Bound<'py
         Ok(if p.is_none() { IndexPart::NewAxis }
         else if p.is(&ellipsis) {
             IndexPart::Rest
-        } else if p.is_instance_of::<PySlice>() && ["start", "stop", "step"].iter().all(|n| p.getattr(*n).is_ok_and(|v| v.is_none())) { IndexPart::All }
-        else if let Ok(slice) = p.cast::<PySlice>() { IndexPart::Slice(Arc::new(slice_spec(slice)?)) }
+        } else if p.is_instance_of::<PySlice>() && ["start", "stop", "step"].iter().all(|n| p.getattr(*n).is_ok_and(|v| v.is_none())) {
+            IndexPart::All
+        } else if let Ok(slice) = p.cast::<PySlice>() { IndexPart::Slice(Arc::new(slice_spec(slice)?)) }
         else { IndexPart::Axis(p) })
     }).collect()
 }

core/src/arc.rs

@@ -70,7 +70,8 @@ pub(crate) fn atan<R: Native, const FUSED: bool>(x: R) -> R {
     (hi + ((u * z).mla::<FUSED>(p, u) + lo)).copysign(x)
 }
 
-pub(crate) fn atan_slice<R: Native>(_: &Atan, _: &crate::recipe::NativeOps<{ crate::FUSED }>, x: [crate::Arg<'_, R>; 1], n: usize, out: &mut [std::mem::MaybeUninit<R>]) -> bool {
+pub(crate) fn atan_slice<R: Native>(
+    _: &Atan, _: &crate::recipe::NativeOps<{ crate::FUSED }>, x: [crate::Arg<'_, R>; 1], n: usize, out: &mut [std::mem::MaybeUninit<R>]) -> bool {
     each!(x, n, out, |[x]| crate::slice::Total(atan::<_, { crate::FUSED }>(x)))
 }
 

core/src/buffer.rs

@@ -787,7 +787,6 @@ pub(crate) fn kept_as<T: Cast>(view: View<'_>, i: usize) -> Result<T, Error> {
     T::cast(view, i).map_err(|_| lost())
 }
 
-
 impl Buffer {
     /// Flattened positions whose item differs from its type's fill (zero, false or a space).
     pub fn nonzero(view: View<'_>) -> Vec<usize> {
@@ -838,11 +837,13 @@ impl Buffer {
     }
     /// The items of `view` at `dtype`. Matching dtypes copy. Otherwise each item converts as [`Source`] reads it, and conversions to a narrower dtype follow these rules. A float converts to an integer by truncating toward zero. A value that doesn't fit its target gives an Overflow error, as do NaN, infinities and anything but exactly 0 or 1 for Booleans. A complex number converts to a real dtype only when its imaginary part is zero, and otherwise gives a Domain error. Characters convert only to characters.
     pub fn cast(view: View<'_>, dtype: DType) -> Result<Self, Error> {
-        crate::convert::Conversion { input: view.dtype(), output: dtype, kind: crate::convert::ConversionKind::Cast }.buffer(view, true, crate::tolerant::Tolerance::EXACT)
+        crate::convert::Conversion { input: view.dtype(), output: dtype, kind: crate::convert::ConversionKind::Cast }.buffer(
+            view, true, crate::tolerant::Tolerance::EXACT)
     }
     /// The items of `view` at `dtype`, where every item keeps its value. Items convert as [`Buffer::cast`] converts them, and a float that becomes a Boolean or an integer must be whole. A float rounds to a float of another width. A Domain error where an item doesn't keep its value.
     pub fn kept(view: View<'_>, dtype: DType) -> Result<Self, Error> {
-        crate::convert::Conversion { input: view.dtype(), output: dtype, kind: crate::convert::ConversionKind::Keep }.buffer(view, true, crate::tolerant::Tolerance::EXACT)
+        crate::convert::Conversion { input: view.dtype(), output: dtype, kind: crate::convert::ConversionKind::Keep }.buffer(
+            view, true, crate::tolerant::Tolerance::EXACT)
     }
 }
 
@@ -853,7 +854,8 @@ impl Buffer {
         crate::convert::Conversion { input: view.dtype(), output: dtype, kind: crate::convert::ConversionKind::Floor }.buffer(view, true, tolerance)
     }
     pub fn cast_results(view: View<'_>, dtype: DType, checked: bool) -> Result<Self, Error> {
-        crate::convert::Conversion { input: view.dtype(), output: dtype, kind: crate::convert::ConversionKind::Truncate }.buffer(view, checked, crate::tolerant::Tolerance::EXACT)
+        crate::convert::Conversion { input: view.dtype(), output: dtype, kind: crate::convert::ConversionKind::Truncate }.buffer(
+            view, checked, crate::tolerant::Tolerance::EXACT)
     }
 }
 

core/src/buffer/output.rs

@@ -59,16 +59,12 @@ impl Scratch {
     }
 }
 
-pub struct Output<'a, 'b> {
-    pub target: MapOutput<'a, 'b>, pub result_dtype: DType, pub checked: bool, pub scratch: Option<&'a mut Scratch>
-}
+pub struct Output<'a, 'b> { pub target: MapOutput<'a, 'b>, pub result_dtype: DType, pub checked: bool, pub scratch: Option<&'a mut Scratch> }
 
 impl Output<'_, '_> {
     /// The block body is erased before output conversion.
     pub fn write<O: OutputItem>(
-        &mut self, n: usize,
-        apply: &(dyn Fn(std::ops::Range<usize>, &mut Slots<'_, O>, Option<&mut Scratch>) -> bool + Sync)
-    ) -> Result<(), Error> {
+        &mut self, n: usize, apply: &(dyn Fn(std::ops::Range<usize>, &mut Slots<'_, O>, Option<&mut Scratch>) -> bool + Sync)) -> Result<(), Error> {
         let checked = self.checked;
         if let MapOutput::Buffer(buffer) = &mut self.target {
             if let Some(values) = O::vec_mut(buffer) { extended(values, n, self.scratch.as_deref_mut(), apply); } else {
@@ -88,8 +84,7 @@ impl Output<'_, '_> {
 
 fn extended<O: Element>(
     values: &mut Vec<O>, n: usize, scratch: Option<&mut Scratch>,
-    apply: &(dyn Fn(std::ops::Range<usize>, &mut Slots<'_, O>, Option<&mut Scratch>) -> bool + Sync)
-) {
+    apply: &(dyn Fn(std::ops::Range<usize>, &mut Slots<'_, O>, Option<&mut Scratch>) -> bool + Sync)) {
     match scratch {
         Some(scratch) => { Slots::append(values, n, |slots| apply(0..n, slots, Some(scratch))); }
         None => { crate::parallel::extended(values, n, crate::CHUNK, &|part, slots| apply(part, slots, None)); }
@@ -98,17 +93,15 @@ fn extended<O: Element>(
 
 fn converted<O: OutputItem>(
     target: &mut MapOutput<'_, '_>, range: std::ops::Range<usize>, checked: bool, mut work: Option<&mut Scratch>,
-    apply: &(dyn Fn(std::ops::Range<usize>, &mut Slots<'_, O>, Option<&mut Scratch>) -> bool + Sync)
-) -> Result<(), Error> {
+    apply: &(dyn Fn(std::ops::Range<usize>, &mut Slots<'_, O>, Option<&mut Scratch>) -> bool + Sync)) -> Result<(), Error> {
     let mut values = work.as_deref_mut().map_or_else(Vec::new, |work| std::mem::take(work.result::<O>()));
     let conversion = crate::convert::Conversion { input: O::DTYPE, output: target.dtype(), kind: crate::convert::ConversionKind::Truncate };
     let result = (|| {
         for at in (range.start..range.end).step_by(crate::CHUNK) {
             let n = crate::CHUNK.min(range.end - at);
             values.clear();
             Slots::append(&mut values, n, |slots| apply(at..at + n, slots, work.as_deref_mut()));
-            if let Err(error) = conversion.append(O::view(&values), checked, crate::tolerant::Tolerance::EXACT,
-                work.as_deref_mut(), target.reborrow()) {
+            if let Err(error) = conversion.append(O::view(&values), checked, crate::tolerant::Tolerance::EXACT, work.as_deref_mut(), target.reborrow()) {
                 target.fill(range.end - at - n);
                 return Err(error);
             }

core/src/complex/bulk.rs

@@ -39,8 +39,7 @@ fn sin_cos<R: Native + Cast>(y: &[R], sin: &mut Vec<R>, cos: &mut Vec<R>) {
 }
 
 /// Real-input cis through the shared real sin/cos slices, with exact quarter turns.
-pub(crate) fn real_cis<R: Native + Cast>(
-    k: &crate::Cis, ops: &NativeOps<{ FUSED }>, [arg]: [Arg<'_, R>; 1], n: usize, out: &mut [MaybeUninit<Complex<R>>]) {
+pub(crate) fn real_cis<R: Native + Cast>(k: &crate::Cis, ops: &NativeOps<{ FUSED }>, [arg]: [Arg<'_, R>; 1], n: usize, out: &mut [MaybeUninit<Complex<R>>]) {
     let ys = match arg {
         Arg::One(y) => {
             let value = k.complex(ops, [c(y, R::zero())]).expect("cis has a complex scalar body");
@@ -62,7 +61,6 @@ pub(crate) fn real_cis<R: Native + Cast>(
     }
 }
 
-
 fn split<K: Formula<N>, R: Native + Cast, const N: usize>(
     k: &K, ops: &NativeOps<{ FUSED }>, args: [Arg<'_, Complex<R>>; N], n: usize, out: &mut [MaybeUninit<Complex<R>>],
     fast: impl FnOnce([(&[R], &[R]); N], &mut Vec<R>, &mut Vec<R>), rare: impl Fn([Complex<R>; N], Complex<R>) -> bool) {

core/src/convert.rs

@@ -91,37 +91,49 @@ impl Conversion {
                         .is_some_and(|(_, valid)| valid != Some(true)))
     }
     /// Converts contiguous storage through the mapped conversion executor.
-    pub(crate) fn buffer(self, view: crate::buffer::View<'_>, checked: bool, tolerance: crate::tolerant::Tolerance) -> Result<crate::buffer::Buffer, crate::error::Error> {
+    pub(crate) fn buffer(
+        self, view: crate::buffer::View<'_>, checked: bool, tolerance: crate::tolerant::Tolerance) -> Result<crate::buffer::Buffer, crate::error::Error> {
         let mut result = crate::buffer::Buffer::with_capacity(self.output, view.len());
         self.append(view, checked, tolerance, None, crate::buffer::MapOutput::Buffer(&mut result))?;
         Ok(result)
     }
     pub(crate) fn append<'a, 'b>(
-        self, view: crate::buffer::View<'_>, checked: bool, tolerance: crate::tolerant::Tolerance,
-        scratch: Option<&'a mut crate::buffer::Scratch>, output: crate::buffer::MapOutput<'a, 'b>
-    ) -> Result<(), crate::error::Error> {
+        self, view: crate::buffer::View<'_>, checked: bool, tolerance: crate::tolerant::Tolerance, scratch: Option<&'a mut crate::buffer::Scratch>,
+        output: crate::buffer::MapOutput<'a, 'b>) -> Result<(), crate::error::Error> {
         let sources = [crate::agreement::Binding::full(view)];
         let plan = crate::agreement::ReadPlan::input(0, view.len());
         self.map_in(crate::agreement::Read { sources: &sources, plan: &plan }, 0..view.len(), checked, tolerance, scratch, output)
     }
     pub fn map_in<'a, 'b>(
         self, input: crate::agreement::Read<'_>, range: std::ops::Range<usize>, checked: bool, tolerance: crate::tolerant::Tolerance,
-        scratch: Option<&'a mut crate::buffer::Scratch>, target: crate::buffer::MapOutput<'a, 'b>
-    ) -> Result<(), crate::error::Error> {
+        scratch: Option<&'a mut crate::buffer::Scratch>, target: crate::buffer::MapOutput<'a, 'b>) -> Result<(), crate::error::Error> {
         let mut output = crate::buffer::Output { target, result_dtype: self.output, checked, scratch };
         macro_rules! write {
-            ($s:ty, $t:ty, $convert:expr) => { written::<$s, $t>(input, range, checked, &mut output, $convert) };
+            ($s:ty, $t:ty, $convert:expr) => {
+                written::<$s, $t>(input, range, checked, &mut output, $convert)
+            };
             (@cast $t:ty) => {{
-                if let crate::buffer::MapOutput::Buffer(buffer) = &mut output.target && let Some(values) = <$t>::vec_mut(buffer) {
+                if let crate::buffer::MapOutput::Buffer(buffer) = &mut output.target
+                    && let Some(values) = <$t>::vec_mut(buffer)
+                {
                     let (at, n) = (values.len(), range.len());
                     if let Err(error) = input.append(range, checked, values) {
-                        crate::error::record(error); values.truncate(at); values.resize(at + n, <$t>::FILL);
+                        crate::error::record(error);
+                        values.truncate(at);
+                        values.resize(at + n, <$t>::FILL);
                     }
                     Ok(())
-                } else { write!($t, $t, Ok) }
+                } else {
+                    write!($t, $t, Ok)
+                }
+            }};
+            (@keep $t:ty) => {
+                kept_mapped::<$t>(self.input, input, range, checked, &mut output)
+            };
+            (@unsupported) => {{
+                output.target.fill(range.len());
+                Err(Error::new(Kind::Domain, "unsupported resolved conversion"))
             }};
-            (@keep $t:ty) => { kept_mapped::<$t>(self.input, input, range, checked, &mut output) };
-            (@unsupported) => {{ output.target.fill(range.len()); Err(Error::new(Kind::Domain, "unsupported resolved conversion")) }};
         }
         crate::error::checked(|| conversion!(self, checked, tolerance, write))
     }
@@ -135,15 +147,14 @@ fn converted_item<T: crate::buffer::Element>(value: Result<T, crate::error::Erro
 #[inline]
 fn written<S: crate::buffer::Cast, T: crate::buffer::OutputItem>(
     input: crate::agreement::Read<'_>, range: std::ops::Range<usize>, checked: bool, output: &mut crate::buffer::Output<'_, '_>,
-    convert: impl Fn(S) -> Result<T, crate::error::Error> + Sync
-) -> Result<(), crate::error::Error> {
+    convert: impl Fn(S) -> Result<T, crate::error::Error> + Sync) -> Result<(), crate::error::Error> {
     output.write::<T>(range.len(), &|part, slots, work| {
         let part = part.start + range.start..part.end + range.start;
         let in_place = input.in_place::<S>(&part);
         let len = if in_place { part.len().max(1) } else { crate::CHUNK };
         let mut local = Vec::new();
         let (mut work, mut copied) = (work, 0);
-        let scratch = match work.as_deref_mut() { Some(work) if !in_place => { let [scratch] = work.inputs::<S, 1>(); scratch }, _ => &mut local };
+        let scratch = match work.as_deref_mut() { Some(work) if !in_place => { let [scratch] = work.inputs::<S, 1>(); scratch } _ => &mut local };
         let ok = part.clone().step_by(len).fold(true, |ok, start| {
             let range = start..(start + len).min(part.end);
             let n = range.len();
@@ -317,15 +328,22 @@ fn kept_mapped<T: crate::buffer::OutputItem + crate::buffer::Cast>(
         let mut local = crate::buffer::Buffer::with_capacity(dtype, 0);
         let scratch = match work { Some(work) => work.buffer(dtype), None => &mut local };
         match input.part_view(part.start + range.start..part.end + range.start, dtype, checked, scratch) {
-            Ok(values) => unsafe { slots.fill(part.len(), |out| {
-                let mut good = true;
-                for (i, slot) in out.iter_mut().enumerate() {
-                    slot.write(match crate::buffer::kept_as::<T>(values, i) {
-                        Ok(value) => value, Err(error) => { crate::error::record(error); good = false; T::FILL }
-                    });
-                }
-                good
-            }) },
+            Ok(values) => unsafe {
+                slots.fill(part.len(), |out| {
+                    let mut good = true;
+                    for (i, slot) in out.iter_mut().enumerate() {
+                        slot.write(match crate::buffer::kept_as::<T>(values, i) {
+                            Ok(value) => value,
+                            Err(error) => {
+                                crate::error::record(error);
+                                good = false;
+                                T::FILL
+                            }
+                        });
+                    }
+                    good
+                })
+            },
             Err(error) => {
                 crate::error::record(error);
                 unsafe { slots.fill(part.len(), |out| { for slot in out { slot.write(T::FILL); } false }) }

core/src/distribution.rs

@@ -66,18 +66,13 @@ trait Formulas {
     #[inline(always)]
     fn quantile_map<F: Float>(&self) -> impl Fn(F) -> F + Copy { move |p| narrowed(p, |p| self.quantile(p)) }
     #[inline(always)]
-    fn pdf_many<F: Native>(&self, x: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool {
-        let pdf = self.pdf_map();
-        each!([x], n, out, |[x]| pdf(x))
-    }
+    fn pdf_many<F: Native>(&self, x: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool { let pdf = self.pdf_map(); each!([x], n, out, |[x]| pdf(x)) }
     #[inline(always)]
-    fn cdf_many<F: Native>(&self, x: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool {
-        let cdf = self.cdf_map();
-        each!([x], n, out, |[x]| cdf(x))
-    }
+    fn cdf_many<F: Native>(&self, x: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool { let cdf = self.cdf_map(); each!([x], n, out, |[x]| cdf(x)) }
     #[inline(always)]
     fn quantile_many<F: Native>(&self, p: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool {
-        let q = self.quantile_map(); each!([p], n, out, |[p]| q(p))
+        let q = self.quantile_map();
+        each!([p], n, out, |[p]| q(p))
     }
 }
 
@@ -648,7 +643,8 @@ impl Formulas for LogNormal {
     fn cdf_many<F: Native>(&self, x: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool {
         let (line, half) = (self.normal().erfc_line(), F::narrow(0.5));
         composed(
-            x, n, out, #[inline(always)] move |x| line(ln(x)), through(Erfc, crate::special::erfc_slice), #[inline(always)] move |x, e| if x <= F::zero() { F::zero() } else { half * e })
+            x, n, out, #[inline(always)] move |x| line(ln(x)), through(Erfc, crate::special::erfc_slice), #[inline(always)]
+            move |x, e| if x <= F::zero() { F::zero() } else { half * e })
     }
     #[inline(always)]
     fn quantile_many<F: Native>(&self, p: Arg<'_, F>, n: usize, out: &mut [MaybeUninit<F>]) -> bool {

core/src/error.rs

@@ -87,9 +87,7 @@ pub fn record(error: Error) {
     });
 }
 
-/// Checked eager entry for Rust and Python callers.
-/// Returns direct failures and, when no slot is active, the first recorded failure. With an
-/// active slot, recorded failures remain for its enclosing operation to collect.
+/// Checked eager entry for Rust and Python callers. Returns direct failures and, when no slot is active, the first recorded failure. With an active slot, recorded failures remain for its enclosing operation to collect.
 pub fn checked<R>(run: impl FnOnce() -> Result<R, Error>) -> Result<R, Error> {
     let (mut run, mut result) = (Some(run), None);
     let recorded = recorded(&mut || result = run.take().map(|run| run()));