remove global_allocator in zk-alloc

Cargo.toml

@@ -83,8 +83,6 @@ include_dir = "0.7"
 
 [features]
 prox-gaps-conjecture = ["rec_aggregation/prox-gaps-conjecture"]
-# Build with the plain system allocator instead of zk-alloc (for comparison/debugging).
-standard-alloc = ["rec_aggregation/standard-alloc"]
 
 [dependencies]
 clap.workspace = true
@@ -102,3 +100,4 @@ system-info.workspace = true
 
 [profile.release]
 lto = "thin"
+codegen-units = 1

crates/backend/field/Cargo.toml

@@ -5,6 +5,7 @@ edition.workspace = true
 
 [dependencies]
 utils = { path = "../utils", package = "utils" }
+zk-alloc.workspace = true
 
 itertools.workspace = true
 num-bigint = "*"

crates/backend/field/src/field.rs

@@ -526,6 +526,31 @@ pub trait BasedVectorSpace<F: PrimeCharacteristicRing>: Sized {
     /// different basis might have been used.
     #[must_use]
     fn reconstitute_from_base(vec: Vec<F>) -> Vec<Self>;
+
+    /// [`flatten_to_base`](Self::flatten_to_base) for `ArenaVec` proof buffers. Defaults to a layout
+    /// reinterpret; override alongside `flatten_to_base` if `Self` isn't `[F; DIMENSION]`.
+    ///
+    /// # Safety
+    /// Same basis-portability caveat as [`flatten_to_base`](Self::flatten_to_base).
+    #[must_use]
+    fn flatten_to_base_in(vec: zk_alloc::ArenaVec<Self>) -> zk_alloc::ArenaVec<F> {
+        // SAFETY: `Self` is `[F; DIMENSION]` (the `flatten_to_base` contract); the const asserts reject mismatches.
+        unsafe { utils::flatten_to_base_in::<F, Self>(vec) }
+    }
+
+    /// [`reconstitute_from_base`](Self::reconstitute_from_base) for `ArenaVec` proof buffers (see
+    /// [`flatten_to_base_in`](Self::flatten_to_base_in)).
+    ///
+    /// # Safety
+    /// Same basis-portability caveat as [`reconstitute_from_base`](Self::reconstitute_from_base).
+    #[must_use]
+    fn reconstitute_from_base_in(vec: zk_alloc::ArenaVec<F>) -> zk_alloc::ArenaVec<Self>
+    where
+        Self: Clone,
+    {
+        // SAFETY: as above.
+        unsafe { utils::reconstitute_from_base_in::<F, Self>(vec) }
+    }
 }
 
 impl<F: PrimeCharacteristicRing> BasedVectorSpace<F> for F {

crates/backend/koala-bear/src/quintic_extension/packed_extension.rs

@@ -48,7 +48,7 @@ impl<F: Field, PF: PackedField<Scalar = F>> From<QuinticExtensionField<F>> for P
     #[inline]
     fn from(x: QuinticExtensionField<F>) -> Self {
         Self {
-            value: x.value.map(Into::into),
+            value: array::from_fn(|i| x.value[i].into()),
         }
     }
 }
@@ -117,10 +117,11 @@ macro_rules! impl_packed_ext_scalar_ops {
         impl Mul<KoalaBear> for PackedQuinticExtensionField<KoalaBear, $pf> {
             type Output = Self;
             #[inline]
-            fn mul(self, rhs: KoalaBear) -> Self {
-                Self {
-                    value: self.value.map(|x| x * rhs),
+            fn mul(mut self, rhs: KoalaBear) -> Self {
+                for v in &mut self.value {
+                    *v *= rhs;
                 }
+                self
             }
         }
 
@@ -281,10 +282,12 @@ where
     type Output = Self;
 
     #[inline]
-    fn neg(self) -> Self {
-        Self {
-            value: self.value.map(PF::neg),
+    fn neg(mut self) -> Self {
+        // Loop, not `self.value.map(..)`: avoids a thin-LTO de-inlined `Wrapped` closure.
+        for v in &mut self.value {
+            *v = -*v;
         }
+        self
     }
 }
 
@@ -478,7 +481,7 @@ where
 
     #[inline(always)]
     fn mul(self, rhs: QuinticExtensionField<F>) -> Self {
-        let b: [PF; 5] = rhs.value.map(|x| x.into());
+        let b: [PF; 5] = array::from_fn(|i| rhs.value[i].into());
         Self {
             value: super::extension::quintic_mul(&self.value, &b, PF::dot_product::<5>),
         }
@@ -493,10 +496,11 @@ where
     type Output = Self;
 
     #[inline]
-    fn mul(self, rhs: PF) -> Self {
-        Self {
-            value: self.value.map(|x| x * rhs),
+    fn mul(mut self, rhs: PF) -> Self {
+        for v in &mut self.value {
+            *v *= rhs;
         }
+        self
     }
 }
 

crates/backend/poly/Cargo.toml

@@ -8,6 +8,8 @@ field = { path = "../field", package = "mt-field" }
 utils = { path = "../utils", package = "utils" }
 system-info.workspace = true
 parallel.workspace = true
+zk-alloc.workspace = true
+
 tracing.workspace = true
 itertools.workspace = true
 rand.workspace = true

crates/backend/poly/src/alloc.rs

@@ -0,0 +1,63 @@
+//! Explicit arena allocation for proof buffers.
+//!
+//! [`ArenaVec`] is [`zk_alloc::ArenaVec`]: an owning vector that bumps from the proving arena inside a
+//! phase and falls back to the system allocator outside one. `Deref<Target = [T]>` lets it drop
+//! into slice-based APIs unchanged.
+
+pub use zk_alloc::{ArenaVec, PhaseGuard, enter_phase};
+
+// Construct empty / pre-sized arena vectors with the inherent `ArenaVec::new()` and
+// `ArenaVec::with_capacity()`. The helpers below cover the cases with no inherent equivalent
+// (`vec![v; n]`, collect, `to_vec`, uninitialized, parallel fill).
+
+/// Arena-backed `vec![value; n]`.
+#[inline]
+#[must_use]
+pub fn arena_filled<T: Clone>(value: T, n: usize) -> ArenaVec<T> {
+    let mut v = ArenaVec::with_capacity(n);
+    v.resize(n, value);
+    v
+}
+
+/// Collect an iterator into an `ArenaVec`.
+#[inline]
+#[must_use]
+pub fn arena_collect<T, I: IntoIterator<Item = T>>(iter: I) -> ArenaVec<T> {
+    let iter = iter.into_iter();
+    let mut v = ArenaVec::with_capacity(iter.size_hint().0);
+    v.extend(iter);
+    v
+}
+
+/// Arena-backed `slice.to_vec()`.
+#[inline]
+#[must_use]
+pub fn arena_from_slice<T: Clone>(slice: &[T]) -> ArenaVec<T> {
+    let mut v = ArenaVec::with_capacity(slice.len());
+    v.extend_from_slice(slice);
+    v
+}
+
+/// Arena-backed [`uninitialized_vec`](crate::uninitialized_vec): `len` uninitialized slots.
+///
+/// # Safety
+/// Every element must be overwritten before it is read.
+#[inline]
+#[must_use]
+pub unsafe fn uninitialized_arena_vec<T>(len: usize) -> ArenaVec<T> {
+    let mut v = ArenaVec::with_capacity(len);
+    // SAFETY: caller guarantees all `len` slots are written before being read.
+    unsafe { v.set_len(len) };
+    v
+}
+
+/// Arena-backed parallel `(0..n).map(f).collect()`: fill an `ArenaVec` of length `n` in parallel.
+/// The single allocation happens on the calling thread; workers write disjoint slots.
+#[inline]
+#[must_use]
+pub fn arena_par_collect<T: Send, F: Fn(usize) -> T + Sync>(n: usize, f: F) -> ArenaVec<T> {
+    // SAFETY: `par_fill` writes every slot in `0..n` exactly once before any is read.
+    let mut v = unsafe { uninitialized_arena_vec(n) };
+    parallel::par_fill(&mut v, f);
+    v
+}

crates/backend/poly/src/evals.rs

@@ -3,8 +3,6 @@ use crate::{EFPacking, PF};
 use ::utils::log2_ceil_usize;
 use field::{ExtensionField, Field, PrimeCharacteristicRing};
 use itertools::Itertools;
-use std::borrow::Borrow;
-
 pub trait EvaluationsList<F: Field> {
     fn num_variables(&self) -> usize;
     fn num_evals(&self) -> usize;
@@ -14,30 +12,30 @@ pub trait EvaluationsList<F: Field> {
     fn evaluate_sparse<EF: ExtensionField<F>>(&self, selector: usize, point: &MultilinearPoint<EF>) -> EF;
 }
 
-impl<F: Field, EL: Borrow<[F]>> EvaluationsList<F> for EL {
+impl<F: Field, EL: AsRef<[F]>> EvaluationsList<F> for EL {
     fn num_variables(&self) -> usize {
-        self.borrow().len().ilog2() as usize
+        self.as_ref().len().ilog2() as usize
     }
 
     fn num_evals(&self) -> usize {
-        self.borrow().len()
+        self.as_ref().len()
     }
 
     fn evaluate<EF: ExtensionField<F>>(&self, point: &MultilinearPoint<EF>) -> EF {
-        eval_multilinear::<_, _, true>(self.borrow(), point)
+        eval_multilinear::<_, _, true>(self.as_ref(), point)
     }
 
     fn evaluate_sequential<EF: ExtensionField<F>>(&self, point: &MultilinearPoint<EF>) -> EF {
-        eval_multilinear::<_, _, false>(self.borrow(), point)
+        eval_multilinear::<_, _, false>(self.as_ref(), point)
     }
 
     fn as_constant(&self) -> F {
-        assert_eq!(self.borrow().len(), 1);
-        self.borrow()[0]
+        assert_eq!(self.as_ref().len(), 1);
+        self.as_ref()[0]
     }
 
     fn evaluate_sparse<EF: ExtensionField<F>>(&self, selector: usize, point: &MultilinearPoint<EF>) -> EF {
-        (&self.borrow()[selector << point.len()..][..(1 << point.len())]).evaluate(point)
+        (&self.as_ref()[selector << point.len()..][..(1 << point.len())]).evaluate(point)
     }
 }
 

crates/backend/poly/src/lib.rs

@@ -1,5 +1,8 @@
 #![cfg_attr(not(test), warn(unused_crate_dependencies))]
 
+mod alloc;
+pub use alloc::*;
+
 mod mle;
 pub use mle::*;
 

crates/backend/poly/src/mle/mle_group_owned.rs

@@ -4,28 +4,28 @@ use field::ExtensionField;
 
 #[derive(Debug)]
 pub enum MleGroupOwned<EF: ExtensionField<PF<EF>>> {
-    Base(Vec<Vec<PF<EF>>>),
-    Extension(Vec<Vec<EF>>),
-    BasePacked(Vec<Vec<PFPacking<EF>>>),
-    ExtensionPacked(Vec<Vec<EFPacking<EF>>>),
+    Base(Vec<ArenaVec<PF<EF>>>),
+    Extension(Vec<ArenaVec<EF>>),
+    BasePacked(Vec<ArenaVec<PFPacking<EF>>>),
+    ExtensionPacked(Vec<ArenaVec<EFPacking<EF>>>),
 }
 
 impl<EF: ExtensionField<PF<EF>>> MleGroupOwned<EF> {
-    pub fn as_extension_mut(&mut self) -> Option<&mut Vec<Vec<EF>>> {
+    pub fn as_extension_mut(&mut self) -> Option<&mut Vec<ArenaVec<EF>>> {
         match self {
             Self::Extension(e) => Some(e),
             _ => None,
         }
     }
 
-    pub fn as_extension_packed_mut(&mut self) -> Option<&mut Vec<Vec<EFPacking<EF>>>> {
+    pub fn as_extension_packed_mut(&mut self) -> Option<&mut Vec<ArenaVec<EFPacking<EF>>>> {
         match self {
             Self::ExtensionPacked(e) => Some(e),
             _ => None,
         }
     }
 
-    pub fn as_extension(self) -> Option<Vec<Vec<EF>>> {
+    pub fn as_extension(self) -> Option<Vec<ArenaVec<EF>>> {
         match self {
             Self::Extension(e) => Some(e),
             _ => None,

crates/backend/poly/src/mle/mle_group_ref.rs

@@ -85,7 +85,7 @@ impl<'a, EF: ExtensionField<PF<EF>>> MleGroupRef<'a, EF> {
             }
             Self::Extension(ext) => {
                 // the only case where there is real work
-                MleGroupOwned::ExtensionPacked(ext.iter().map(|v| pack_extension(v)).collect()).into()
+                MleGroupOwned::ExtensionPacked(ext.iter().map(|v| pack_extension_in(v)).collect()).into()
             }
             Self::BasePacked(_) | Self::ExtensionPacked(_) => self.soft_clone().into(),
         }
@@ -99,7 +99,7 @@ impl<'a, EF: ExtensionField<PF<EF>>> MleGroupRef<'a, EF> {
                 MleGroupRef::Base(pols.iter().map(|v| PFPacking::<EF>::unpack_slice(v)).collect()).into()
             }
             Self::ExtensionPacked(pols) => {
-                MleGroupOwned::Extension(pols.iter().map(|v| unpack_extension(v)).collect()).into()
+                MleGroupOwned::Extension(pols.iter().map(|v| unpack_extension_in(v)).collect()).into()
             }
         }
     }
@@ -158,10 +158,12 @@ impl<'a, EF: ExtensionField<PF<EF>>> MleGroupRef<'a, EF> {
 
     pub fn clone_to_owned(&self) -> MleGroupOwned<EF> {
         match self {
-            Self::Base(pols) => MleGroupOwned::Base(pols.iter().map(|v| v.to_vec()).collect()),
-            Self::Extension(pols) => MleGroupOwned::Extension(pols.iter().map(|v| v.to_vec()).collect()),
-            Self::BasePacked(pols) => MleGroupOwned::BasePacked(pols.iter().map(|v| v.to_vec()).collect()),
-            Self::ExtensionPacked(pols) => MleGroupOwned::ExtensionPacked(pols.iter().map(|v| v.to_vec()).collect()),
+            Self::Base(pols) => MleGroupOwned::Base(pols.iter().map(|v| arena_from_slice(v)).collect()),
+            Self::Extension(pols) => MleGroupOwned::Extension(pols.iter().map(|v| arena_from_slice(v)).collect()),
+            Self::BasePacked(pols) => MleGroupOwned::BasePacked(pols.iter().map(|v| arena_from_slice(v)).collect()),
+            Self::ExtensionPacked(pols) => {
+                MleGroupOwned::ExtensionPacked(pols.iter().map(|v| arena_from_slice(v)).collect())
+            }
         }
     }
 }