book_scratch_pool/
book_scratch_pool.rs

1//! Temporary buffers through Mnemosyne's scratch pool.
2//!
3//! [`ScratchPool<T>`] provides a reusable pool of up to
4//! [`MAX_POOL_SLOTS`] aligned buffers for a single element type.  The
5//! pool is designed for `thread_local!` storage; it is `Send` but not `Sync`.
6//! Construction is zero-allocation: the backing buffer is only mapped when
7//! `with_scratch` is first called.
8//!
9//! Up to `MAX_POOL_SLOTS` (4) nested borrows are supported without growing the
10//! pool, which covers recursive FFT twiddle computation, nested solver
11//! residuals, and similar patterns common in kwavers and apollo.
12
13extern crate mnemosyne;
14
15use mnemosyne::{MAX_POOL_SLOTS, ScratchBank, ScratchPool};
16
17fn dot_product(a: &[f64], b: &[f64]) -> f64 {
18    assert_eq!(a.len(), b.len());
19    a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
20}
21
22fn main() {
23    // Construct a pool locally and use it directly (not thread_local! here
24    // to avoid the clippy::missing_const_for_thread_local false positive on
25    // 1.97.0 — see ATLAS-MNEMOSYNE-CI-1).  In production callers store the
26    // pool in thread_local! storage so it persists between calls.
27    let f64_pool: ScratchPool<f64> = ScratchPool::new();
28    let f32_pool: ScratchPool<f32> = ScratchPool::new();
29
30    // Simple single-level scratch use: compute a dot product via a temp buffer.
31    f64_pool.with_scratch(1024, |scratch| {
32        for (i, slot) in scratch.iter_mut().enumerate() {
33            *slot = i as f64;
34        }
35        let partial: f64 = scratch.iter().sum();
36        println!(
37            "sum of 0..1024 via scratch: {} (expected {})",
38            partial,
39            (0..1024usize).sum::<usize>() as f64,
40        );
41        assert_eq!(partial, (0..1024usize).sum::<usize>() as f64);
42    });
43
44    // Nested borrows: outer scratch holds the signal; inner holds a window.
45    f64_pool.with_scratch(256, |signal| {
46        for (i, s) in signal.iter_mut().enumerate() {
47            *s = (i as f64).sin();
48        }
49        f64_pool.with_scratch(16, |window| {
50            window.copy_from_slice(&signal[..16]);
51            let dp = dot_product(window, &signal[..16]);
52            println!("windowed dot-product: {dp:.6}");
53            assert!(dp.is_finite());
54        });
55    });
56
57    // F32 pool: same API, different element type.
58    f32_pool.with_scratch(64, |buf| {
59        buf.fill(1.0_f32);
60        let total: f32 = buf.iter().sum();
61        println!("f32 scratch sum: {total}");
62        assert_eq!(total, 64.0_f32);
63    });
64
65    // Banked scratch keeps multiple related roles in one const-generic group,
66    // which matches transform pipelines that need independent temporary views
67    // without falling back to the system allocator.
68    let bank: ScratchBank<f64, 2> = ScratchBank::new();
69    bank.with_scratch::<1, _>(32, |scratch| {
70        scratch.fill(3.5);
71        assert!(scratch.iter().all(|v| *v == 3.5));
72    });
73
74    // Bounded provisioning retains the working set while making geometric
75    // growth headroom reclaimable at a consumer-selected quiescent point.
76    let bounded_pool: ScratchPool<u32> = ScratchPool::new();
77    bounded_pool.with_scratch_bounded(1024, |_| {});
78    bounded_pool.with_scratch_bounded(1025, |_| {});
79    assert_eq!(bounded_pool.capacity(), 2048);
80    let retained = bounded_pool.release();
81    println!(
82        "bounded release: capacity={} -> retained={}",
83        2048, retained[0]
84    );
85    assert_eq!(retained[0], 1025);
86    bounded_pool.reset();
87    assert_eq!(bounded_pool.release()[0], 0);
88
89    println!("MAX_POOL_SLOTS = {MAX_POOL_SLOTS} (max concurrent nested borrows)");
90    println!("all scratch-pool assertions passed");
91}
book_scratch_pool/book_scratch_pool.rs

book_scratch_pool/
book_scratch_pool.rs
