From 468b61353d3a32288781d91d6a695bd3e0ffa8b4 Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Wed, 30 Sep 2026 19:20:05 +0100 Subject: [PATCH] Make ArenaVec a Vec over a std Allocator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `Allocator` trait, `Vec` and `Box` are stable from Rust 1.100 (rust-lang/rust#156882). `zk_alloc::Arena` is a zero-sized `Allocator` over `raw_alloc`/`raw_dealloc`, which already tell arena from system memory by address, and `ArenaVec` is `Vec`. The hand-written vector is gone and every `Vec` method works on an `ArenaVec`. - `Vec`'s constructors belong to the global allocator, so `new`, `with_capacity`, `from_slice`, `filled`, `zeroed`, `uninitialized` and `from_iter` move to an `ArenaVecExt` trait. `collect()`, `Default` and `mem::take` are unavailable for the same reason; four call sites use `from_iter` and `mem::replace` instead. - `primitives::multilinear::Shrink` only let the folds take either vector type. It is deleted and the folds are generic over `A: Allocator`, calling `truncate` (not `Vec::shrink_to`, which shrinks capacity). - An allocator's methods must never unwind. The address space was reserved lazily inside the first allocation behind an `assert!`; `enable_arena` now reserves it. - `clippy::uninit_vec` now sees two deliberate uninitialized vectors it could not see through the old type; each gets an `#[expect]` with its reason. Requires Rust 1.100 (nightly until it ships); stable 1.97 rejects `allocator_api`. Measured on x86-64 (32 threads), same nightly for both builds, alternating runs, proving time: fibonacci --n 2000000: base 1.864 s, 1.883 s; new 1.889 s, 1.923 s leanxmss --n 400: base 3.247, 3.275, 3.152 (±7.2%) s; new 3.284, 3.244, 3.248 s No difference beyond run-to-run noise; peak memory within 0.3%. --- AGENTS.md | 2 + crates/flock/src/witness.rs | 2 +- crates/flock/src/zerocheck.rs | 2 +- crates/flock/src/zerocheck/multilinear.rs | 2 +- crates/leanvm_core/src/cpu/layout.rs | 1 + crates/leanvm_core/src/gkr.rs | 21 +- crates/leanvm_core/src/leaf.rs | 3 +- crates/leanvm_core/src/witness.rs | 2 +- crates/pcs/benches/throughput.rs | 1 + crates/pcs/src/ring_switch.rs | 1 + crates/pcs/src/whir.rs | 2 +- crates/pcs/src/whir_induce.rs | 4 +- crates/primitives/src/lib.rs | 2 +- crates/primitives/src/multilinear.rs | 37 +- crates/zk_alloc/src/arena_vec.rs | 474 +++++----------------- crates/zk_alloc/src/lib.rs | 24 +- crates/zk_alloc/tests/arena.rs | 2 +- 17 files changed, 151 insertions(+), 431 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index eb7c2505d..1a7479de8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -124,6 +124,8 @@ CI tracks benchmarks on Bencher in three workflows, Bencher's pattern for PRs fr One proof is one **phase**, opened by `cpu::prove`. `ArenaVec` bumps a per-thread slab, a small block's release is at most a cursor pop while a large one is recycled (below), and the next `begin_phase()` reclaims everything. Not a `#[global_allocator]`: `raw_dealloc` picks arena-vs-system by address range, so with no phase open `ArenaVec` is an ordinary system vector (used in particular by the verifier, where correctness and simplicity matters much more than performance). +`ArenaVec` is `Vec`, `Arena` being a zero-sized `std::alloc::Allocator` (stable from Rust 1.100), so every `Vec` method works on it. Its constructors (`new`, `with_capacity`, `from_slice`, `from_iter`, the unsafe `uninitialized` and `zeroed`) come from the `ArenaVecExt` trait, `Vec`'s own being the global allocator's; for the same reason it has no `collect()`, `Default` or `mem::take` (use `ArenaVec::from_iter` and `mem::replace(.., ArenaVec::new())`). A function taking either vector is generic over `A: Allocator`. The allocator's methods must never unwind, which is why `enable_arena` reserves the address space rather than the first allocation. + **The rule:** an `ArenaVec` allocated in a phase dies at the next `begin_phase()`. A reset neither clears nor unmaps, so a buffer that outlives its phase reads the previous proof's plausible bytes, so the symptom is a proof that stops verifying, never a crash. Anything outliving a phase (a `Proof`, a cache, a table) must be a plain `Vec`. And **`drop` means something**: a large released block is handed back out within the phase (a per-thread free list, see the crate docs), so dropping a big buffer where it dies is worth doing, and a use-after-free the bump arena used to mask now reads another buffer's live data. Run `ZK_ALLOC_POISON=1 cargo testall` after changing buffer lifetimes; it fills released blocks and fills what a phase used when it ends, turning a silent wrong answer into a loud failure. That covers both shapes: a buffer read after being dropped, and a buffer that outlives its phase. `setup_prover_without_arena` (or `leanvm_core::init_prover_pool` alone) leaves the arena disengaged, sending every `ArenaVec` to the system allocator. It is the escape hatch for a host where even the recycled peak does not fit; on one that it does fit, the arena is faster, since its pages stay faulted in across proofs. diff --git a/crates/flock/src/witness.rs b/crates/flock/src/witness.rs index 4e3f6a964..db5563391 100644 --- a/crates/flock/src/witness.rs +++ b/crates/flock/src/witness.rs @@ -4,7 +4,7 @@ use primitives::bits::bit_transpose_64bytes; use primitives::stream::Stream; -use zk_alloc::ArenaVec; +use zk_alloc::{ArenaVec, ArenaVecExt}; /// OR the low 32 bits of `val` into `buf` starting at bit-offset `bit_off`. /// Handles u64 straddling when `bit_off % 64 > 32`. diff --git a/crates/flock/src/zerocheck.rs b/crates/flock/src/zerocheck.rs index 0b537c75e..028f8456d 100644 --- a/crates/flock/src/zerocheck.rs +++ b/crates/flock/src/zerocheck.rs @@ -24,7 +24,7 @@ use fiat_shamir::transcript::{Challenger, ProverState, Receiver, Transmitter, VerifierState}; use primitives::field::{F8, F192}; -use zk_alloc::ArenaVec; +use zk_alloc::{ArenaVec, ArenaVecExt}; use pcs::ntt::{AdditiveNttGf8, InvNttTableByteSingleGf8}; diff --git a/crates/flock/src/zerocheck/multilinear.rs b/crates/flock/src/zerocheck/multilinear.rs index abf9cc769..fca951cb7 100644 --- a/crates/flock/src/zerocheck/multilinear.rs +++ b/crates/flock/src/zerocheck/multilinear.rs @@ -41,7 +41,7 @@ use crate::zerocheck::univariate_skip::pack_bits; use crate::zerocheck::univariate_skip::{SplitEq, build_eq}; use primitives::field::{F192, F192Unreduced, PHI_8_TABLE_192 as PHI_8_TABLE}; use primitives::stream::Stream; -use zk_alloc::ArenaVec; +use zk_alloc::{ArenaVec, ArenaVecExt}; /// Four independent products. Tuples keep the scalar and NEON paths in registers, while AVX-512 uses the batched helper. #[inline(always)] diff --git a/crates/leanvm_core/src/cpu/layout.rs b/crates/leanvm_core/src/cpu/layout.rs index 2360b2018..1f75ca083 100644 --- a/crates/leanvm_core/src/cpu/layout.rs +++ b/crates/leanvm_core/src/cpu/layout.rs @@ -5,6 +5,7 @@ use super::*; use crate::leaf::SparseColumn; use crate::rv::{ADVICE_BASE, LOG_REGS, RAM_BASE, TEXT_BASE}; +use zk_alloc::ArenaVecExt; // ---- column schema ----------------------------------------------------------- diff --git a/crates/leanvm_core/src/gkr.rs b/crates/leanvm_core/src/gkr.rs index 8191d86e2..17ded68d5 100644 --- a/crates/leanvm_core/src/gkr.rs +++ b/crates/leanvm_core/src/gkr.rs @@ -11,7 +11,7 @@ use crate::transcript::{Challenger, ProverState, Receiver, Transmitter, Verifier use primitives::field::{F192, F192Unreduced, mul_unreduced4, mul2, mul4}; use primitives::multilinear::{eq_table, interp, poly_eval, shrink_eq_low}; use primitives::stream::Stream; -use zk_alloc::ArenaVec; +use zk_alloc::{ArenaVec, ArenaVecExt}; #[derive(Clone, Debug, PartialEq, Eq)] pub enum GkrError { @@ -53,7 +53,7 @@ fn build_layers(leaves: ArenaVec, mu: usize) -> Vec> { } else if full_rows >= PAR_THRESHOLD { primitives::par_collect_arena(full_rows, product) } else { - (0..full_rows).map(product).collect() + ArenaVec::from_iter((0..full_rows).map(product)) }; if !current.len().is_multiple_of(4) && current.len() > 2 { let row = full_rows; @@ -367,8 +367,9 @@ pub fn prove_product_triple(leaves: [ArenaVec; 3], ps: &mut ProverState, s } let width = 1usize << round_count; - let mut trees = - [0, 1, 2].map(|tree| QuaternaryLayerState::new(std::mem::take(&mut layers[tree][layer - 2]), width)); + let mut trees = [0, 1, 2].map(|tree| { + QuaternaryLayerState::new(std::mem::replace(&mut layers[tree][layer - 2], ArenaVec::new()), width) + }); let mut equality = if round_count > 0 { eq_table(&point[1..]) } else { @@ -521,9 +522,9 @@ mod tests { #[test] fn quartic_round_message_matches_direct_evaluation() { for width in [2, 4, 8, 16] { - let below: ArenaVec = (0..4 * width) - .map(|i| F192::new((17 * i + width + 1) as u64, (i * i + 3) as u64, (5 * i + 7) as u64)) - .collect(); + let below = ArenaVec::from_iter( + (0..4 * width).map(|i| F192::new((17 * i + width + 1) as u64, (i * i + 3) as u64, (5 * i + 7) as u64)), + ); let state = QuaternaryLayerState::new(below, width); let equality: Vec = (0..width / 2) .map(|i| F192::new((31 * i + 5) as u64, (7 * i + 1) as u64, (11 * i + 9) as u64)) @@ -554,9 +555,9 @@ mod tests { if len > 4 * width { continue; } - let values: ArenaVec = (0..len) - .map(|i| F192::new((17 * i + 1) as u64, (i * i + 3) as u64, (5 * i + 7) as u64)) - .collect(); + let values = ArenaVec::from_iter( + (0..len).map(|i| F192::new((17 * i + 1) as u64, (i * i + 3) as u64, (5 * i + 7) as u64)), + ); let mut reference = QuaternaryLayerState::new(ArenaVec::from_slice(&values), width); let mut fused = QuaternaryLayerState::new(values, width); let point: Vec = (0..width.ilog2() - 1) diff --git a/crates/leanvm_core/src/leaf.rs b/crates/leanvm_core/src/leaf.rs index 712373933..4e4308f75 100644 --- a/crates/leanvm_core/src/leaf.rs +++ b/crates/leanvm_core/src/leaf.rs @@ -15,7 +15,7 @@ use primitives::field::{F64, F192, F192Unreduced, g_pow, index_mle, int_index_ml use primitives::multilinear::{eq_eval, eq_table_arena, mle_eval}; use std::collections::HashMap; use std::sync::Arc; -use zk_alloc::ArenaVec; +use zk_alloc::{ArenaVec, ArenaVecExt}; /// One tuple coordinate as a function of the block's row `z`. #[derive(Clone, Debug)] @@ -356,6 +356,7 @@ pub fn build_leaves( // Capacity is rounded to whole four-tuples because `gkr::QuaternaryLayerState` // pads this level to that before reading it, and growing it here would copy it. let covered: usize = blocks.iter().map(|blk| 1usize << blk.kappa).sum(); + #[expect(clippy::uninit_vec, reason = "the per-block fills below cover `0..explicit`")] let mut leaves = if covered == explicit { let mut values = ArenaVec::with_capacity(explicit.next_multiple_of(4)); // SAFETY: the per-block fills below cover `0..explicit` exactly, and each diff --git a/crates/leanvm_core/src/witness.rs b/crates/leanvm_core/src/witness.rs index 8444c2237..9bfeeab4b 100644 --- a/crates/leanvm_core/src/witness.rs +++ b/crates/leanvm_core/src/witness.rs @@ -5,7 +5,7 @@ //! the column's offset. use primitives::field::F64; -use zk_alloc::ArenaVec; +use zk_alloc::{ArenaVec, ArenaVecExt}; /// Where a column sits in the stacked witness. A [`Placement::VIRTUAL`] column is /// NOT committed: it carries data for the bus, but its evaluation claims settle diff --git a/crates/pcs/benches/throughput.rs b/crates/pcs/benches/throughput.rs index 055e3c96f..75236f7db 100644 --- a/crates/pcs/benches/throughput.rs +++ b/crates/pcs/benches/throughput.rs @@ -28,6 +28,7 @@ use primitives::{ pretty_integer, test_rng::Rng, }; +use zk_alloc::ArenaVecExt; fn main() { bench::init_tracing_from_env(); diff --git a/crates/pcs/src/ring_switch.rs b/crates/pcs/src/ring_switch.rs index 78a648fe9..d7e54aebf 100644 --- a/crates/pcs/src/ring_switch.rs +++ b/crates/pcs/src/ring_switch.rs @@ -497,6 +497,7 @@ mod tests { }; use crate::whir_config::test_config_for; use primitives::test_rng::Rng; + use zk_alloc::ArenaVecExt; /// Pack bit `64 * y + i` of `bits` into bit `i` of word `y`. fn pack_witness(bits: &[bool]) -> Vec { diff --git a/crates/pcs/src/whir.rs b/crates/pcs/src/whir.rs index 3b0f46539..8790bc2e0 100644 --- a/crates/pcs/src/whir.rs +++ b/crates/pcs/src/whir.rs @@ -30,7 +30,7 @@ use primitives::{ multilinear::eq_eval, stream::Stream, }; -use zk_alloc::ArenaVec; +use zk_alloc::{ArenaVec, ArenaVecExt}; pub use super::whir_config::{ FinalBlockConfig, INITIAL_FOLDING_FACTOR, LOG_INV_RATE_0, LevelShapes, MAX_LOG_INV_RATE, MIN_LOG_INV_RATE, diff --git a/crates/pcs/src/whir_induce.rs b/crates/pcs/src/whir_induce.rs index c34b18982..bec2fe1aa 100644 --- a/crates/pcs/src/whir_induce.rs +++ b/crates/pcs/src/whir_induce.rs @@ -10,7 +10,7 @@ use crate::ntt::AdditiveNttF64; use crate::whir::build_eq_table_ext; use primitives::field::{F64, F192}; -use zk_alloc::ArenaVec; +use zk_alloc::{ArenaVec, ArenaVecExt}; // =================================================================== // LCH novel-basis evaluations over K (mirror of whir's extension-field block) @@ -272,7 +272,7 @@ pub(crate) fn induce_sumcheck_evaluate_at_residual( if yr_len > PAR_FLOOR { primitives::par_collect_arena(yr_len, compute_y) } else { - (0..yr_len).map(compute_y).collect() + ArenaVec::from_iter((0..yr_len).map(compute_y)) } } diff --git a/crates/primitives/src/lib.rs b/crates/primitives/src/lib.rs index 762bfee2a..846669e8c 100644 --- a/crates/primitives/src/lib.rs +++ b/crates/primitives/src/lib.rs @@ -111,7 +111,7 @@ pub fn log2_ceil_usize(n: usize) -> usize { pub fn par_collect_arena(n: usize, build: impl Fn(usize) -> T + Sync) -> zk_alloc::ArenaVec { // SAFETY: the fill below writes every slot in `0..n` exactly once, and // the dispatch joins before the buffer is observable. - let mut out = unsafe { zk_alloc::ArenaVec::uninitialized(n) }; + let mut out = unsafe { as zk_alloc::ArenaVecExt>::uninitialized(n) }; parallel::fill(&mut out, build); out } diff --git a/crates/primitives/src/multilinear.rs b/crates/primitives/src/multilinear.rs index ce2741495..e37d4823d 100644 --- a/crates/primitives/src/multilinear.rs +++ b/crates/primitives/src/multilinear.rs @@ -6,35 +6,12 @@ //! `K`-valued (`F64`) while randomness is `E`-valued (`F192`), so the first //! fold of a committed table also lifts it into `E`. -use std::ops::DerefMut; - +use std::alloc::Allocator; use std::mem::MaybeUninit; use crate::field::{F64, F192, PHI_8_TABLE_192 as PHI_8_TABLE, Weights8, dot_base, mul_base8, mul4}; use zk_alloc::ArenaVec; -/// The one thing the in-place folds need beyond a mutable slice: the ability to -/// drop a suffix. Implemented for `Vec` and `ArenaVec`, so a fold works on either -/// without duplicating the kernel or naming a container in its signature. -pub trait Shrink: DerefMut { - /// Keep the first `len` elements, dropping the rest. - fn shrink_to(&mut self, len: usize); -} - -impl Shrink for Vec { - #[inline] - fn shrink_to(&mut self, len: usize) { - self.truncate(len); - } -} - -impl Shrink for ArenaVec { - #[inline] - fn shrink_to(&mut self, len: usize) { - self.truncate(len); - } -} - /// Multilinear interpolation in one variable over `E`: `lo + t·(lo+hi)`, the /// char-2 form of `(1−t)·lo + t·hi`. #[inline] @@ -181,7 +158,7 @@ pub fn fold_high_k(table: &[F64], chi: F192) -> ArenaVec { /// Bind the highest free variable of `table` to `chi` in place: `table[i] = /// interp(table[i], table[i + half], chi)`. Binding from the top down leaves the /// low variables, the ones every table of a batch shares, for last. -pub fn fold_high_inplace>(table: &mut B, chi: F192) { +pub fn fold_high_inplace(table: &mut Vec, chi: F192) { debug_assert_eq!(table.len() % 2, 0); let half = table.len() / 2; { @@ -191,7 +168,7 @@ pub fn fold_high_inplace>(table: &mut B, chi: F192) { let (lo, hi) = (**table).split_at_mut(half); interp_into(lo, hi, chi); } - table.shrink_to(half); + table.truncate(half); } /// `lo[i] = interp(lo[i], hi[i], chi)`, four products per batch. @@ -211,7 +188,7 @@ fn interp_into(lo: &mut [F192], hi: &[F192], chi: F192) { /// Marginalize the lowest variable out of an `eq` table (in place). `eq(r_0, 0) + /// eq(r_0, 1) = 1`, so summing adjacent entries drops `r_0` with no multiplies, /// versus `2^{n-1}` to rebuild the table. -pub fn shrink_eq_low>(table: &mut B) { +pub fn shrink_eq_low(table: &mut Vec) { let half = table.len() / 2; { // Sliced, as in `fold_high_inplace`: reading the pair and writing the @@ -223,12 +200,12 @@ pub fn shrink_eq_low>(table: &mut B) { t[i] = a + b; } } - table.shrink_to(half); + table.truncate(half); } /// Marginalize the highest variable out of an `eq` table (in place), the /// [`shrink_eq_low`] counterpart for a top-down sumcheck. -pub fn shrink_eq_high>(table: &mut B) { +pub fn shrink_eq_high(table: &mut Vec) { let half = table.len() / 2; { // Sliced, as in `fold_high_inplace`. @@ -237,7 +214,7 @@ pub fn shrink_eq_high>(table: &mut B) { *l += *h; } } - table.shrink_to(half); + table.truncate(half); } /// The one barycentric denominator an aligned `size`-node window of the φ₈ table has: `∏_{k≠0} φ₈(k)`, diff --git a/crates/zk_alloc/src/arena_vec.rs b/crates/zk_alloc/src/arena_vec.rs index 878a489fe..21cfcb55a 100644 --- a/crates/zk_alloc/src/arena_vec.rs +++ b/crates/zk_alloc/src/arena_vec.rs @@ -1,297 +1,156 @@ -//! Growable buffer backed by the proving arena during a phase and the system allocator otherwise. Size arena-backed buffers up front anyway: growing copies, and the old allocation is only recycled if it clears the arena's reuse floor. +//! The proving arena as a [`std::alloc::Allocator`], and [`ArenaVec`], the `Vec` over it. Size arena-backed buffers up front anyway: growing copies, and the old allocation is only recycled if it clears the arena's reuse floor. -use std::alloc::{Layout, handle_alloc_error}; -use std::cmp; -use std::fmt; -use std::marker::PhantomData; -use std::mem::{ManuallyDrop, MaybeUninit, align_of, needs_drop, size_of}; -use std::ops::{Deref, DerefMut}; +use std::alloc::{AllocError, Allocator, Layout}; +use std::mem::MaybeUninit; use std::ptr::{self, NonNull}; -use std::slice; use crate::{raw_alloc, raw_dealloc}; -/// An owning, growable buffer allocated from the proving arena. -pub struct ArenaVec { - ptr: NonNull, - len: usize, - cap: usize, - _marker: PhantomData, +/// The proving arena: this thread's slab inside a phase, the system allocator otherwise. +/// +/// Zero-sized, so an [`ArenaVec`] is the size of a `Vec`. Which side a block came from is decided on release by +/// address range, so one handle serves both. +#[derive(Clone, Copy, Debug, Default)] +pub struct Arena; + +// SAFETY: `raw_alloc` returns a block valid for `layout.size()` bytes at `layout.align()` (it only ever raises the +// alignment), or null, reported as `AllocError`; `raw_dealloc` takes back exactly what `raw_alloc` handed out. The +// block stays valid until it is released or the next phase opens, which the crate's phase rule forbids any live +// `ArenaVec` to see. Zero-sized requests never reach the arena. +unsafe impl Allocator for Arena { + #[inline(always)] + fn allocate(&self, layout: Layout) -> Result, AllocError> { + if layout.size() == 0 { + // SAFETY: an alignment is nonzero. + let dangling = unsafe { NonNull::new_unchecked(ptr::without_provenance_mut(layout.align())) }; + return Ok(NonNull::slice_from_raw_parts(dangling, 0)); + } + // SAFETY: `layout` has a power-of-two alignment and a nonzero size. + let p = unsafe { raw_alloc(layout.size(), layout.align()) }; + NonNull::new(p) + .map(|p| NonNull::slice_from_raw_parts(p, layout.size())) + .ok_or(AllocError) + } + + #[inline(always)] + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + if layout.size() != 0 { + // SAFETY: the caller guarantees `ptr` came from `allocate` with this layout. + unsafe { raw_dealloc(ptr.as_ptr(), layout.size(), layout.align()) }; + } + } } -// SAFETY: `ArenaVec` owns its elements exactly as `Vec` does; the arena imposes -// no additional thread affinity (a slab is bumped only by its owning thread, and -// the resulting pointer is plain memory). -unsafe impl Send for ArenaVec {} -unsafe impl Sync for ArenaVec {} - -impl ArenaVec { - /// `usize::MAX` stands in for "unbounded" for zero-sized elements, which - /// never allocate; `0` otherwise. - const EMPTY_CAP: usize = if size_of::() == 0 { usize::MAX } else { 0 }; +/// An owning, growable buffer allocated from the proving arena. +pub type ArenaVec = Vec; +/// The constructors `Vec` defines for the global allocator only, for [`ArenaVec`]. +pub trait ArenaVecExt: Sized { /// An empty buffer. Allocates nothing. - #[inline] - #[must_use] - pub const fn new() -> Self { - Self { - ptr: NonNull::dangling(), - len: 0, - cap: Self::EMPTY_CAP, - _marker: PhantomData, - } - } + fn new() -> Self; - /// An empty buffer with room for exactly `cap` elements (no over-allocation: - /// the caller knows the size, and a slab is not a heap to grow into). - #[inline] - #[must_use] - pub fn with_capacity(cap: usize) -> Self { - let mut v = Self::new(); - if size_of::() != 0 && cap != 0 { - v.grow_to(cap); - } - v - } + /// An empty buffer with room for `cap` elements. + fn with_capacity(cap: usize) -> Self; - /// `n` elements, each a clone of `value`, the arena's `vec![value; n]`. - /// Prefer [`zeroed`](Self::zeroed) when the value is all-zero bytes. - #[inline] - #[must_use] - pub fn filled(value: T, n: usize) -> Self + /// `n` elements, each a clone of `value`, the arena's `vec![value; n]`. Prefer [`zeroed`](Self::zeroed) when + /// the value is all-zero bytes. + fn filled(value: T, n: usize) -> Self where - T: Clone, - { - let mut v = Self::with_capacity(n); - v.resize(n, value); - v - } + T: Clone; /// `n` zero-filled elements. /// /// # Safety - /// The all-zero bit pattern must be a valid, fully initialized `T`, true of - /// the field types here and their SIMD packings, whose zero is all-zero bytes. - #[inline] - #[must_use] - pub unsafe fn zeroed(n: usize) -> Self { - // SAFETY: every slot is initialized by the write_bytes below before any - // read is possible. - let mut v = unsafe { Self::uninitialized(n) }; - // SAFETY: `v` owns `n` slots; the caller guarantees all-zero is a valid `T`. - unsafe { ptr::write_bytes(v.as_mut_ptr(), 0u8, n) }; - v - } + /// The all-zero bit pattern must be a valid, fully initialized `T`, true of the field types here and their SIMD + /// packings, whose zero is all-zero bytes. + unsafe fn zeroed(n: usize) -> Self; /// The arena's `slice.to_vec()`. - #[inline] - #[must_use] - pub fn from_slice(slice: &[T]) -> Self + fn from_slice(slice: &[T]) -> Self where - T: Clone, - { - let mut v = Self::with_capacity(slice.len()); - v.extend_from_slice(slice); - v - } + T: Clone; - /// `len` slots of uninitialized memory, ready to be filled in place (in - /// parallel, typically: `ArenaVec` derefs to `&mut [T]`). + /// `len` slots of uninitialized memory, ready to be filled in place (in parallel, typically). /// /// # Safety /// Every one of the `len` elements must be written before it is read. - #[inline] - #[must_use] - pub unsafe fn uninitialized(len: usize) -> Self { - let mut v = Self::with_capacity(len); - // SAFETY: the caller guarantees all `len` slots are written before read. - unsafe { v.set_len(len) }; - v - } + unsafe fn uninitialized(len: usize) -> Self; - #[inline] - #[must_use] - pub const fn len(&self) -> usize { - self.len - } - - #[inline] - #[must_use] - pub const fn capacity(&self) -> usize { - self.cap - } - - #[inline] - #[must_use] - pub const fn is_empty(&self) -> bool { - self.len == 0 - } - - #[inline] - #[must_use] - pub const fn as_ptr(&self) -> *const T { - self.ptr.as_ptr() - } + /// The arena's `collect()`. + fn from_iter>(iter: I) -> Self; +} +impl ArenaVecExt for ArenaVec { #[inline] - pub const fn as_mut_ptr(&mut self) -> *mut T { - self.ptr.as_ptr() + fn new() -> Self { + Vec::new_in(Arena) } #[inline] - #[must_use] - pub fn as_slice(&self) -> &[T] { - self + fn with_capacity(cap: usize) -> Self { + Vec::with_capacity_in(cap, Arena) } - /// Set the length without touching the buffer. - /// - /// # Safety - /// `new_len <= capacity()`, and every element below `new_len` must be - /// initialized. #[inline] - pub unsafe fn set_len(&mut self, new_len: usize) { - debug_assert!(new_len <= self.cap); - self.len = new_len; - } - - /// Make room for `additional` more elements, doubling to keep repeated - /// [`push`](Self::push) amortized. - #[inline] - pub fn reserve(&mut self, additional: usize) { - if size_of::() == 0 { - return; // capacity is conceptually unbounded for zero-sized types - } - let required = self.len.checked_add(additional).expect("ArenaVec capacity overflow"); - if required > self.cap { - self.grow_to(cmp::max(required, self.cap.saturating_mul(2))); - } + fn filled(value: T, n: usize) -> Self + where + T: Clone, + { + let mut v = Self::with_capacity(n); + v.resize(n, value); + v } #[inline] - pub fn push(&mut self, value: T) { - if self.len == self.cap { - // Zero-sized types never reach here: their capacity is `usize::MAX`. - self.grow_to(cmp::max(self.cap.saturating_mul(2), 4)); + unsafe fn zeroed(n: usize) -> Self { + let mut v = Self::with_capacity(n); + // SAFETY: `v` has room for `n`, and the caller guarantees all-zero is a valid `T`. + unsafe { + ptr::write_bytes(v.as_mut_ptr(), 0u8, n); + v.set_len(n); } - // SAFETY: `len < cap` now, so slot `len` is allocated and uninitialized. - unsafe { self.ptr.as_ptr().add(self.len).write(value) }; - self.len += 1; + v } - /// Append a clone of every element of `other`. #[inline] - pub fn extend_from_slice(&mut self, other: &[T]) + fn from_slice(slice: &[T]) -> Self where T: Clone, { - self.reserve(other.len()); - // Bump `len` per element, so a panic mid-clone leaves a consistent buffer - // whose written clones still drop. - for x in other { - // SAFETY: `reserve` guaranteed room for `other.len()` more elements. - unsafe { self.ptr.as_ptr().add(self.len).write(x.clone()) }; - self.len += 1; - } - } - - /// Grow or shrink to `new_len`, cloning `value` into any new slots. - pub fn resize(&mut self, new_len: usize, value: T) - where - T: Clone, - { - if new_len > self.len { - self.reserve(new_len - self.len); - while self.len < new_len { - // SAFETY: room reserved above, so `len < new_len <= cap`. - unsafe { self.ptr.as_ptr().add(self.len).write(value.clone()) }; - self.len += 1; - } - } else { - self.truncate(new_len); - } - } - - /// Drop everything past `len`, keeping capacity. - pub fn truncate(&mut self, len: usize) { - if len < self.len { - let dropped = self.len - len; - // Shorten first, so a panicking `Drop` cannot observe or re-drop the tail. - self.len = len; - // SAFETY: `[len, old_len)` were initialized and are now logically gone. - unsafe { - ptr::drop_in_place(ptr::slice_from_raw_parts_mut(self.ptr.as_ptr().add(len), dropped)); - } - } + let mut v = Self::with_capacity(slice.len()); + v.extend_from_slice(slice); + v } - /// Decompose into raw parts, leaking the buffer. Inverse of - /// [`from_raw_parts`](Self::from_raw_parts). #[inline] - #[must_use] - pub(crate) fn into_raw_parts(self) -> (*mut T, usize, usize) { - let me = ManuallyDrop::new(self); - (me.ptr.as_ptr(), me.len, me.cap) + #[expect( + clippy::uninit_vec, + reason = "the caller's contract: every slot is written before it is read" + )] + unsafe fn uninitialized(len: usize) -> Self { + let mut v = Self::with_capacity(len); + // SAFETY: the caller guarantees all `len` slots are written before read. + unsafe { v.set_len(len) }; + v } - /// Reassemble from parts produced by [`into_raw_parts`](Self::into_raw_parts) - /// (or a layout-compatible reinterpretation of them). - /// - /// # Safety - /// `ptr` must be non-null and aligned for `T`, `len <= cap`, and either `ptr` - /// came from `raw_alloc` for `cap * size_of::()` bytes - /// at `align_of::()`, or `cap == 0` and `ptr` is dangling-but-aligned. - /// Exactly one `ArenaVec` may own a given pointer. #[inline] - #[must_use] - pub(crate) unsafe fn from_raw_parts(ptr: *mut T, len: usize, cap: usize) -> Self { - Self { - // SAFETY: the caller guarantees `ptr` is non-null. - ptr: unsafe { NonNull::new_unchecked(ptr) }, - len, - cap, - _marker: PhantomData, - } - } - - /// Allocate a fresh `new_cap`-element buffer, move the live elements across, - /// and release the old one. Only reached for sized `T` with - /// `new_cap >= len` and `new_cap > 0`. - fn grow_to(&mut self, new_cap: usize) { - debug_assert!(size_of::() != 0 && new_cap >= self.len && new_cap > 0); - let align = align_of::(); - let bytes = new_cap.checked_mul(size_of::()).expect("ArenaVec capacity overflow"); - assert!(bytes <= isize::MAX as usize, "ArenaVec capacity overflow"); - - // SAFETY: `align` is a valid power of two and `bytes > 0`. - let raw = unsafe { raw_alloc(bytes, align) }.cast::(); - let Some(new_ptr) = NonNull::new(raw) else { - // Matching `Vec`: allocation failure aborts rather than unwinds. - // SAFETY: `align` is a power of two and `bytes <= isize::MAX`. - handle_alloc_error(unsafe { Layout::from_size_align_unchecked(bytes, align) }); - }; - - if self.cap != 0 { - // SAFETY: the buffers are distinct and `len` initialized elements move. - unsafe { ptr::copy_nonoverlapping(self.ptr.as_ptr(), new_ptr.as_ptr(), self.len) }; - // SAFETY: the old buffer came from `raw_alloc` with this size/align. - unsafe { raw_dealloc(self.ptr.as_ptr().cast::(), self.cap * size_of::(), align) }; - } - self.ptr = new_ptr; - self.cap = new_cap; + fn from_iter>(iter: I) -> Self { + let mut v = Self::new(); + v.extend(iter); + v } } /// `n` arena-backed slots to be initialized in place, one element at a time. /// -/// The `MaybeUninit` element type is what makes a partial fill expressible: write -/// through the slice, then [`assume_init`] once every slot is set. Prefer -/// [`ArenaVec::uninitialized`] when the fill is a bulk write over `&mut [T]`. +/// The `MaybeUninit` element type is what makes a partial fill expressible: write through the slice, then +/// [`assume_init`] once every slot is set. Prefer [`ArenaVecExt::uninitialized`] when the fill is a bulk write over +/// `&mut [T]`. #[inline] #[must_use] pub fn alloc_uninit(n: usize) -> ArenaVec> { - // SAFETY: `MaybeUninit` is valid uninitialized, so every slot already - // holds a valid value of the element type. + // SAFETY: `MaybeUninit` is valid uninitialized. unsafe { ArenaVec::uninitialized(n) } } @@ -302,139 +161,8 @@ pub fn alloc_uninit(n: usize) -> ArenaVec> { #[inline] #[must_use] pub unsafe fn assume_init(v: ArenaVec>) -> ArenaVec { - let (ptr, len, cap) = v.into_raw_parts(); - // SAFETY: `MaybeUninit` has the same size and alignment as `T`, so the - // allocation matches `T`'s layout; the caller guarantees every slot is - // initialized, and `into_raw_parts` transferred sole ownership. - unsafe { ArenaVec::from_raw_parts(ptr.cast::(), len, cap) } -} - -impl Drop for ArenaVec { - fn drop(&mut self) { - // Drop the live elements first. Elided entirely for plain-data `T`. - if needs_drop::() { - // SAFETY: `0..len` are initialized. - unsafe { ptr::drop_in_place(ptr::slice_from_raw_parts_mut(self.ptr.as_ptr(), self.len)) }; - } - // Release the buffer. Zero-sized types and never-grown vectors own nothing. - if size_of::() != 0 && self.cap != 0 { - // SAFETY: the buffer came from `raw_alloc(cap * size, align)`; - // `raw_dealloc` range-checks arena-vs-system, and an arena pointer is - // recycled or left to the next phase reset rather than unmapped. - unsafe { - raw_dealloc( - self.ptr.as_ptr().cast::(), - self.cap * size_of::(), - align_of::(), - ); - } - } - } -} - -impl Deref for ArenaVec { - type Target = [T]; - #[inline] - fn deref(&self) -> &[T] { - // SAFETY: `ptr` is aligned and `0..len` are initialized. Valid for - // zero-sized `T` too: a dangling aligned pointer is a sound base for a - // zero-stride slice. - unsafe { slice::from_raw_parts(self.ptr.as_ptr(), self.len) } - } -} - -impl DerefMut for ArenaVec { - #[inline] - fn deref_mut(&mut self) -> &mut [T] { - // SAFETY: as `deref`, with unique access. - unsafe { slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len) } - } -} - -impl Default for ArenaVec { - #[inline] - fn default() -> Self { - Self::new() - } -} - -impl Clone for ArenaVec { - fn clone(&self) -> Self { - Self::from_slice(self) - } -} - -impl fmt::Debug for ArenaVec { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt::Debug::fmt(&**self, f) - } -} - -impl PartialEq for ArenaVec { - #[inline] - fn eq(&self, other: &Self) -> bool { - **self == **other - } -} - -impl PartialEq<[T]> for ArenaVec { - #[inline] - fn eq(&self, other: &[T]) -> bool { - **self == *other - } -} - -impl PartialEq> for ArenaVec { - #[inline] - fn eq(&self, other: &Vec) -> bool { - **self == **other - } -} - -impl PartialEq> for Vec { - #[inline] - fn eq(&self, other: &ArenaVec) -> bool { - **self == **other - } -} - -impl Eq for ArenaVec {} - -impl Extend for ArenaVec { - #[inline] - fn extend>(&mut self, iter: I) { - let iter = iter.into_iter(); - self.reserve(iter.size_hint().0); - for x in iter { - self.push(x); - } - } -} - -impl FromIterator for ArenaVec { - #[inline] - fn from_iter>(iter: I) -> Self { - let iter = iter.into_iter(); - let mut v = Self::with_capacity(iter.size_hint().0); - v.extend(iter); - v - } -} - -impl<'a, T> IntoIterator for &'a ArenaVec { - type Item = &'a T; - type IntoIter = slice::Iter<'a, T>; - #[inline] - fn into_iter(self) -> Self::IntoIter { - self.iter() - } -} - -impl<'a, T> IntoIterator for &'a mut ArenaVec { - type Item = &'a mut T; - type IntoIter = slice::IterMut<'a, T>; - #[inline] - fn into_iter(self) -> Self::IntoIter { - self.iter_mut() - } + let (ptr, len, cap, alloc) = v.into_parts_with_allocator(); + // SAFETY: `MaybeUninit` has the layout of `T`, so the allocation matches; the caller guarantees every slot is + // initialized, and `into_parts_with_allocator` transferred sole ownership. + unsafe { Vec::from_parts_in(ptr.cast::(), len, cap, alloc) } } diff --git a/crates/zk_alloc/src/lib.rs b/crates/zk_alloc/src/lib.rs index 5880e32a5..c36b3f4ba 100644 --- a/crates/zk_alloc/src/lib.rs +++ b/crates/zk_alloc/src/lib.rs @@ -16,9 +16,10 @@ //! # Usage //! //! ```no_run +//! use zk_alloc::{ArenaVec, ArenaVecExt}; //! zk_alloc::enable_arena(); // once, at startup //! let _phase = zk_alloc::enter_phase(); // bind before the phase's buffers -//! let buf: zk_alloc::ArenaVec = zk_alloc::ArenaVec::with_capacity(1 << 20); +//! let buf: ArenaVec = ArenaVec::with_capacity(1 << 20); //! ``` use std::alloc::{GlobalAlloc, Layout, System}; @@ -29,7 +30,7 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; mod arena_vec; mod syscall; -pub use arena_vec::{ArenaVec, alloc_uninit, assume_init}; +pub use arena_vec::{Arena, ArenaVec, ArenaVecExt, alloc_uninit, assume_init}; /// Address space reserved per thread. Overflow falls back to the system allocator. const SLAB_SIZE: usize = 64 << 30; @@ -182,8 +183,9 @@ thread_local! { static NO_SLAB: Cell = const { Cell::new(false) }; } -fn region() -> usize { - *REGION.get_or_init(|| { +/// Map the reservation. Done in [`enable_arena`], never on the allocation path, which must not unwind. +fn reserve_region() { + REGION.get_or_init(|| { let size = region_size(); // SAFETY: `reserve` returns a page-aligned pointer or null. let ptr = unsafe { syscall::reserve(size) }; @@ -195,7 +197,7 @@ fn region() -> usize { // SAFETY: the mapping we just made is live and exactly `size` bytes. unsafe { syscall::disable_huge_pages(ptr, size) }; ptr as usize - }) + }); } /// Opt into the arena. Call once at startup, before any proving. @@ -205,6 +207,7 @@ fn region() -> usize { /// memory-constrained host, at the cost of the page-fault churn described in the /// module docs. pub fn enable_arena() { + reserve_region(); syscall::retain_system_heap(); POISON.store(std::env::var_os("ZK_ALLOC_POISON").is_some(), Ordering::Relaxed); ARENA_ENGAGED.store(true, Ordering::Release); @@ -340,7 +343,12 @@ unsafe fn alloc_slow(size: usize, align: usize) -> *mut u8 { NO_SLAB.set(true); return unsafe { system_alloc(size, align) }; } - base = region() + idx * SLAB_SIZE; + // Set by `enable_arena`, which a phase requires. + let Some(®ion) = REGION.get() else { + NO_SLAB.set(true); + return unsafe { system_alloc(size, align) }; + }; + base = region + idx * SLAB_SIZE; BASE.set(base); END.set(base + SLAB_SIZE); } else { @@ -458,8 +466,8 @@ fn dealloc_large(addr: usize, size: usize) { /// leave it to the next [`begin_phase`]; for a system pointer, a system free. /// /// Which one applies is decided by address range, not by a flag stored beside -/// the buffer, which is what lets [`ArenaVec`] carry no allocator parameter and -/// stay a single type whether or not a phase was open when it was built. +/// the buffer, which is what lets [`Arena`] be zero-sized and one allocator +/// whether or not a phase was open when a buffer was built. /// /// Recycling is safe NOT because of a cursor property (reuse puts live /// allocations below the cursor), but because a block reaches the list only from diff --git a/crates/zk_alloc/tests/arena.rs b/crates/zk_alloc/tests/arena.rs index f501252ec..5db1be4e0 100644 --- a/crates/zk_alloc/tests/arena.rs +++ b/crates/zk_alloc/tests/arena.rs @@ -5,7 +5,7 @@ use std::sync::{Mutex, MutexGuard}; -use zk_alloc::{ArenaVec, enable_arena, enter_phase}; +use zk_alloc::{ArenaVec, ArenaVecExt, enable_arena, enter_phase}; const N: usize = 4096;