Make ArenaVec a Vec over a std Allocator

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<T>` is `Vec<T, zk_alloc::Arena>`, `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.

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`.

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};
 

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)]

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 -----------------------------------------------------------
 

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<F192>, mu: usize) -> Vec<ArenaVec<F192>> {
         } 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<F192>; 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<F192> = (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<F192> = (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<F192> = (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<F192> = (0..width.ilog2() - 1)

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

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

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();

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<F64> {

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,

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))
     }
 }
 

crates/primitives/src/lib.rs

@@ -111,7 +111,7 @@ pub fn log2_ceil_usize(n: usize) -> usize {
 pub fn par_collect_arena<T: Send>(n: usize, build: impl Fn(usize) -> T + Sync) -> zk_alloc::ArenaVec<T> {
     // 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 { <zk_alloc::ArenaVec<T> as zk_alloc::ArenaVecExt<T>>::uninitialized(n) };
     parallel::fill(&mut out, build);
     out
 }

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<T>: DerefMut<Target = [T]> {
-    /// Keep the first `len` elements, dropping the rest.
-    fn shrink_to(&mut self, len: usize);
-}
-
-impl<T> Shrink<T> for Vec<T> {
-    #[inline]
-    fn shrink_to(&mut self, len: usize) {
-        self.truncate(len);
-    }
-}
-
-impl<T> Shrink<T> for ArenaVec<T> {
-    #[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<F192> {
 /// 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<B: Shrink<F192>>(table: &mut B, chi: F192) {
+pub fn fold_high_inplace<A: Allocator>(table: &mut Vec<F192, A>, chi: F192) {
     debug_assert_eq!(table.len() % 2, 0);
     let half = table.len() / 2;
     {
@@ -191,7 +168,7 @@ pub fn fold_high_inplace<B: Shrink<F192>>(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<B: Shrink<F192>>(table: &mut B) {
+pub fn shrink_eq_low<A: Allocator>(table: &mut Vec<F192, A>) {
     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<B: Shrink<F192>>(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<B: Shrink<F192>>(table: &mut B) {
+pub fn shrink_eq_high<A: Allocator>(table: &mut Vec<F192, A>) {
     let half = table.len() / 2;
     {
         // Sliced, as in `fold_high_inplace`.
@@ -237,7 +214,7 @@ pub fn shrink_eq_high<B: Shrink<F192>>(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)`,