pub const BASE_CHUNK: usize = 128;
#[inline]
const fn left_split(len: usize) -> usize {
assert!(
len > BASE_CHUNK,
"left_split: caller must guarantee len > BASE_CHUNK"
);
let mut k = BASE_CHUNK;
while k.saturating_mul(2) < len {
k = k.saturating_mul(2);
}
k
}
#[inline]
fn reduce_block<T, F>(acc: T, items: &[T], combine: &F) -> T
where
T: Copy,
F: Fn(T, T) -> T,
{
let mut out = acc;
for &x in items {
out = combine(out, x);
}
out
}
pub fn pairwise_reduce<T, F>(items: &[T], combine: F, identity: T) -> T
where
T: Copy,
F: Fn(T, T) -> T,
{
reduce_range(items, &combine, identity)
}
fn reduce_range<T, F>(items: &[T], combine: &F, identity: T) -> T
where
T: Copy,
F: Fn(T, T) -> T,
{
let len = items.len();
if len == 0 {
return identity;
}
if len <= BASE_CHUNK {
return reduce_block(items[0], &items[1..], combine);
}
let mid = left_split(len);
let left = reduce_range(&items[..mid], combine, identity);
let right = reduce_range(&items[mid..], combine, identity);
combine(left, right)
}
pub fn pairwise_sum(xs: &[f64]) -> f64 {
pairwise_reduce(xs, |a, b| a + b, 0.0)
}
pub struct StreamingPairwise<T, F>
where
T: Copy,
F: Fn(T, T) -> T,
{
combine: F,
identity: T,
buf: Vec<T>,
forest: Vec<(usize, T)>,
}
impl<T, F> StreamingPairwise<T, F>
where
T: Copy,
F: Fn(T, T) -> T,
{
pub fn new(combine: F, identity: T) -> Self {
Self {
combine,
identity,
buf: Vec::with_capacity(BASE_CHUNK),
forest: Vec::new(),
}
}
pub fn push(&mut self, x: T) {
self.buf.push(x);
if self.buf.len() == BASE_CHUNK {
let block = reduce_block(self.buf[0], &self.buf[1..], &self.combine);
self.buf.clear();
self.absorb(BASE_CHUNK, block);
}
}
pub fn extend_from_slice(&mut self, chunk: &[T]) {
for &x in chunk {
self.push(x);
}
}
fn absorb(&mut self, weight: usize, value: T) {
let mut w = weight;
let mut v = value;
while let Some(&(top_w, top_v)) = self.forest.last() {
if top_w == w {
self.forest.pop();
v = (self.combine)(top_v, v);
w = w.saturating_mul(2);
} else {
break;
}
}
self.forest.push((w, v));
}
pub fn finish(mut self) -> T {
if !self.buf.is_empty() {
let tail = reduce_block(self.buf[0], &self.buf[1..], &self.combine);
let tail_w = self.buf.len();
self.buf.clear();
self.forest.push((tail_w, tail));
}
let mut iter = self.forest.into_iter().rev();
match iter.next() {
None => self.identity,
Some((_, mut acc)) => {
for (_, left) in iter {
acc = (self.combine)(left, acc);
}
acc
}
}
}
}
pub fn pairwise_reduce_chunked<'a, T, F, I>(chunks: I, combine: F, identity: T) -> T
where
T: Copy + 'a,
F: Fn(T, T) -> T,
I: IntoIterator<Item = &'a [T]>,
{
let mut acc = StreamingPairwise::new(combine, identity);
for chunk in chunks {
acc.extend_from_slice(chunk);
}
acc.finish()
}
pub fn pairwise_sum_chunked<'a, I>(chunks: I) -> f64
where
I: IntoIterator<Item = &'a [f64]>,
{
pairwise_reduce_chunked(chunks, |a, b| a + b, 0.0)
}
pub fn par_pairwise_map_reduce<T, M, F>(n: usize, map: M, combine: F, identity: T) -> T
where
T: Copy + Send + Sync,
M: Fn(usize) -> T + Sync,
F: Fn(T, T) -> T + Sync,
{
par_reduce_index_range(0, n, &map, &combine, identity)
}
fn par_reduce_index_range<T, M, F>(lo: usize, hi: usize, map: &M, combine: &F, identity: T) -> T
where
T: Copy + Send + Sync,
M: Fn(usize) -> T + Sync,
F: Fn(T, T) -> T + Sync,
{
let len = hi - lo;
if len == 0 {
return identity;
}
if len <= BASE_CHUNK {
let mut acc = map(lo);
for i in (lo + 1)..hi {
acc = combine(acc, map(i));
}
return acc;
}
let mid = lo + left_split(len);
let (left, right) = rayon::join(
move || par_reduce_index_range(lo, mid, map, combine, identity),
move || par_reduce_index_range(mid, hi, map, combine, identity),
);
combine(left, right)
}
pub fn par_pairwise_sum<M>(n: usize, f: M) -> f64
where
M: Fn(usize) -> f64 + Sync,
{
par_pairwise_map_reduce(n, f, |a, b| a + b, 0.0)
}
pub fn par_deterministic_block_fold<T, B, F>(n: usize, base: B, combine: F) -> Option<T>
where
T: Send,
B: Fn(core::ops::Range<usize>) -> T + Sync,
F: Fn(T, T) -> T + Sync,
{
if n == 0 {
return None;
}
Some(par_block_fold_range(0, n, &base, &combine))
}
pub fn par_deterministic_try_block_fold<T, E, B, F>(
n: usize,
base: B,
combine: F,
) -> Result<Option<T>, E>
where
T: Send,
E: Send,
B: Fn(core::ops::Range<usize>) -> Result<T, E> + Sync,
F: Fn(T, T) -> Result<T, E> + Sync,
{
if n == 0 {
return Ok(None);
}
par_try_block_fold_range(0, n, &base, &combine).map(Some)
}
fn par_try_block_fold_range<T, E, B, F>(lo: usize, hi: usize, base: &B, combine: &F) -> Result<T, E>
where
T: Send,
E: Send,
B: Fn(core::ops::Range<usize>) -> Result<T, E> + Sync,
F: Fn(T, T) -> Result<T, E> + Sync,
{
let len = hi - lo;
if len <= BASE_CHUNK {
return base(lo..hi);
}
let mid = lo + left_split(len);
let (left, right) = rayon::join(
|| par_try_block_fold_range(lo, mid, base, combine),
|| par_try_block_fold_range(mid, hi, base, combine),
);
combine(left?, right?)
}
fn par_block_fold_range<T, B, F>(lo: usize, hi: usize, base: &B, combine: &F) -> T
where
T: Send,
B: Fn(core::ops::Range<usize>) -> T + Sync,
F: Fn(T, T) -> T + Sync,
{
let len = hi - lo;
if len <= BASE_CHUNK {
return base(lo..hi);
}
let mid = lo + left_split(len);
let (left, right) = rayon::join(
|| par_block_fold_range(lo, mid, base, combine),
|| par_block_fold_range(mid, hi, base, combine),
);
combine(left, right)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn left_split_minimal_case() {
assert_eq!(super::left_split(BASE_CHUNK + 1), BASE_CHUNK);
}
#[test]
fn left_split_at_two_blocks() {
assert_eq!(super::left_split(2 * BASE_CHUNK), BASE_CHUNK);
}
#[test]
fn left_split_just_above_two_blocks() {
assert_eq!(super::left_split(2 * BASE_CHUNK + 1), 2 * BASE_CHUNK);
}
#[test]
fn left_split_at_four_blocks() {
assert_eq!(super::left_split(4 * BASE_CHUNK), 2 * BASE_CHUNK);
}
#[test]
fn pairwise_reduce_empty_returns_identity() {
let result = pairwise_reduce::<u64, _>(&[], |a, b| a + b, 99);
assert_eq!(result, 99);
}
#[test]
fn pairwise_reduce_single_element() {
assert_eq!(pairwise_reduce(&[42u64], |a, b| a + b, 0), 42);
}
#[test]
fn pairwise_reduce_small_sum() {
let xs = [1u64, 2, 3, 4, 5];
assert_eq!(pairwise_reduce(&xs, |a, b| a + b, 0), 15);
}
#[test]
fn pairwise_reduce_product() {
let xs = [2u64, 3, 4, 5];
assert_eq!(pairwise_reduce(&xs, |a, b| a * b, 1), 120);
}
#[test]
fn pairwise_sum_empty_is_zero() {
assert_eq!(pairwise_sum(&[]), 0.0);
}
#[test]
fn pairwise_sum_single_element() {
assert_eq!(pairwise_sum(&[3.5f64]), 3.5);
}
#[test]
fn pairwise_sum_small_slice_exact() {
assert_eq!(pairwise_sum(&[1.0f64, 2.0, 3.0, 4.0, 5.0]), 15.0);
}
#[test]
fn pairwise_sum_exactly_base_chunk_elements() {
let xs: Vec<f64> = (1..=BASE_CHUNK as u64).map(|x| x as f64).collect();
let naive: f64 = (1..=BASE_CHUNK as u64).map(|x| x as f64).sum();
assert_eq!(pairwise_sum(&xs), naive);
}
#[test]
fn pairwise_sum_one_above_base_chunk_triggers_split() {
let xs = vec![1.0f64; BASE_CHUNK + 1];
assert_eq!(pairwise_sum(&xs), (BASE_CHUNK + 1) as f64);
}
#[test]
fn pairwise_sum_two_base_chunks() {
let xs = vec![1.0f64; 2 * BASE_CHUNK];
assert_eq!(pairwise_sum(&xs), (2 * BASE_CHUNK) as f64);
}
#[test]
fn streaming_one_at_a_time_matches_whole_slice() {
let xs: Vec<f64> = (0..300).map(|i| i as f64 * 0.1).collect();
let expected = pairwise_sum(&xs);
let mut acc = StreamingPairwise::new(|a: f64, b: f64| a + b, 0.0);
for &x in &xs {
acc.push(x);
}
assert_eq!(acc.finish().to_bits(), expected.to_bits());
}
#[test]
fn chunked_matches_whole_slice_across_chunk_sizes() {
let xs: Vec<f64> = (0..500).map(|i| i as f64).collect();
let expected = pairwise_sum(&xs);
for chunk_size in [1usize, 7, 64, 128, 129, 200, 499, 500] {
let chunks: Vec<&[f64]> = xs.chunks(chunk_size).collect();
let result = pairwise_sum_chunked(chunks);
assert_eq!(
result.to_bits(),
expected.to_bits(),
"chunk_size={chunk_size}"
);
}
}
#[test]
fn pairwise_reduce_chunked_matches_whole_slice() {
let xs: Vec<u64> = (1..=300).collect();
let expected = pairwise_reduce(&xs, |a, b| a + b, 0u64);
let chunks: Vec<&[u64]> = xs.chunks(77).collect();
let result = pairwise_reduce_chunked(chunks, |a, b| a + b, 0u64);
assert_eq!(result, expected);
}
#[test]
fn pairwise_sum_chunked_basic() {
let a = [1.0f64, 2.0, 3.0];
let b = [4.0f64, 5.0];
assert_eq!(pairwise_sum_chunked([a.as_ref(), b.as_ref()]), 15.0);
}
#[test]
fn par_pairwise_sum_bit_identical_to_sequential() {
for n in [
0usize,
1,
5,
BASE_CHUNK,
BASE_CHUNK + 1,
3 * BASE_CHUNK + 17,
5000,
] {
let f = |i: usize| ((i as f64) * 0.7318 - 41.0).sin() * 1e6 / ((i + 1) as f64);
let xs: Vec<f64> = (0..n).map(f).collect();
let expected = pairwise_sum(&xs);
let got = par_pairwise_sum(n, f);
assert_eq!(got.to_bits(), expected.to_bits(), "n={n}");
}
}
#[test]
fn par_pairwise_sum_bit_stable_across_thread_counts() {
let n = 7 * BASE_CHUNK + 13;
let f = |i: usize| (1.0 / ((i + 1) as f64)).ln_1p() * if i % 3 == 0 { -1.0 } else { 1.0 };
let reference = par_pairwise_sum(n, f);
for threads in [1usize, 2, 3, 8] {
let pool = rayon::ThreadPoolBuilder::new()
.num_threads(threads)
.build()
.expect("pool");
let got = pool.install(|| par_pairwise_sum(n, f));
assert_eq!(got.to_bits(), reference.to_bits(), "threads={threads}");
}
}
#[test]
fn par_pairwise_map_reduce_tuple_accumulators() {
let n = 4 * BASE_CHUNK + 3;
let map = |i: usize| {
let x = (i as f64).sqrt() * 0.911;
(x, x * x)
};
let seq: (f64, f64) = {
let xs: Vec<(f64, f64)> = (0..n).map(map).collect();
pairwise_reduce(&xs, |a, b| (a.0 + b.0, a.1 + b.1), (0.0, 0.0))
};
let par = par_pairwise_map_reduce(n, map, |a, b| (a.0 + b.0, a.1 + b.1), (0.0, 0.0));
assert_eq!(par.0.to_bits(), seq.0.to_bits());
assert_eq!(par.1.to_bits(), seq.1.to_bits());
}
#[test]
fn par_deterministic_block_fold_vec_accumulator_thread_invariant() {
let n = 6 * BASE_CHUNK + 41;
let dim = 7usize;
let run = || {
par_deterministic_block_fold(
n,
|range: core::ops::Range<usize>| {
let mut acc = vec![0.0f64; dim];
for i in range {
for (d, slot) in acc.iter_mut().enumerate() {
*slot += ((i * dim + d) as f64).cos() / ((i + 1) as f64);
}
}
acc
},
|mut a: Vec<f64>, b: Vec<f64>| {
for (x, y) in a.iter_mut().zip(&b) {
*x += *y;
}
a
},
)
.expect("n > 0")
};
let reference = run();
for threads in [1usize, 2, 5] {
let pool = rayon::ThreadPoolBuilder::new()
.num_threads(threads)
.build()
.expect("pool");
let got = pool.install(run);
for (g, r) in got.iter().zip(&reference) {
assert_eq!(g.to_bits(), r.to_bits(), "threads={threads}");
}
}
assert!(par_deterministic_block_fold(0, |_| vec![0.0f64; 1], |a, _| a).is_none());
}
#[test]
fn par_deterministic_try_block_fold_matches_infallible_and_propagates_errors() {
let n = 5 * BASE_CHUNK + 7;
let base = |range: core::ops::Range<usize>| -> f64 {
range.map(|i| ((i as f64) * 0.317).sin()).sum()
};
let infallible = par_deterministic_block_fold(n, base, |a, b| a + b).expect("n > 0");
let fallible = par_deterministic_try_block_fold(
n,
|range| Ok::<f64, String>(base(range)),
|a, b| Ok(a + b),
)
.expect("no error")
.expect("n > 0");
assert_eq!(fallible.to_bits(), infallible.to_bits());
let err = par_deterministic_try_block_fold(
n,
|range| {
if range.contains(&(3 * BASE_CHUNK)) {
Err("boom".to_string())
} else {
Ok(0.0f64)
}
},
|a, b| Ok(a + b),
)
.unwrap_err();
assert_eq!(err, "boom");
assert!(
par_deterministic_try_block_fold(0, |_| Ok::<f64, String>(0.0), |a, b| Ok(a + b))
.expect("no error")
.is_none()
);
}
}