From c83923c3ce86df67b277afbc8b919c9d8e170d9a Mon Sep 17 00:00:00 2001 From: Tom Wambsgans Date: Sun, 26 Jul 2026 21:14:36 +0200 Subject: [PATCH 1/5] field: pin the prover field's prime subfield to KoalaBear Replace the self-referential bound `EF: ExtensionField>` (where `PF = ::PrimeSubfield`) with a named bound that states what the codebase already relies on: pub trait KoalaBearExtension: Field + ExtensionField + PrimeCharacteristicRing {} Both accept exactly the same types -- there is one prime field and one extension field in the tree -- so this is a no-op for type checking. What changes is that `PF` now *normalizes* to the concrete `KoalaBear` instead of remaining an opaque projection, which pays off three ways: 1. Aeneas extraction. The old form made Charon diverge: translating the projection `PF` needs a proof of `EF: PrimeCharacteristicRing`, and with only that bound in scope the proof was rooted at a predicate that itself mentioned `PF`. Naming the base concretely removes the cycle at the root. `sumcheck::verify::sumcheck_verify` goes from stack overflow to a 1.3 MB LLBC; `whir::verify` and `lean_prover::verify_execution` also extract. 2. `EF: Algebra` now follows from `ExtensionField`'s supertrait, so code can operate on prime-field constants directly. 3. `PF == KoalaBear` becomes a fact the compiler knows. Clippy immediately flagged the two transmutes in `restore_merkle_paths` as transmutes-to-self, so that `TypeId` assert and both `unsafe` blocks are deleted. 16 `TypeId` sites remain elsewhere and are now provably unnecessary in the same way. Also drops 12 `PF: TwoAdicField` / `PrimeField64` bounds that the pinning makes redundant. One `WhirConfig` impl keeps its `PF: TwoAdicField` because it only requires `EF: Field`, so the projection stays opaque there; and `ConstraintFolderPacked`'s `AirBuilder` impl keeps the old bound because its `low_degree_block` override mentions `&mut Self` in a higher-ranked closure bound, where the solver will not look through the blanket-implemented alias. Verified: 66 tests pass, clippy -Dwarnings and rustfmt clean, and the workspace builds under no-SIMD, AVX2 and AVX-512. Co-Authored-By: Claude Opus 5 (1M context) --- .../air/src/constraint_folder/normal.rs | 6 ++-- .../air/src/constraint_folder/packed.rs | 4 +++ crates/backend/fiat-shamir/src/prover.rs | 19 ++++------- crates/backend/fiat-shamir/src/traits.rs | 6 ++-- crates/backend/fiat-shamir/src/verifier.rs | 32 +++++------------- .../backend/koala-bear/src/extension_bound.rs | 13 ++++++++ crates/backend/koala-bear/src/lib.rs | 2 ++ crates/backend/poly/Cargo.toml | 2 +- crates/backend/poly/src/eq_mle.rs | 18 +++++----- crates/backend/poly/src/evals.rs | 4 +-- crates/backend/poly/src/mle/mle_group.rs | 9 +++-- .../backend/poly/src/mle/mle_group_owned.rs | 5 ++- crates/backend/poly/src/mle/mle_group_ref.rs | 5 ++- crates/backend/poly/src/mle/mle_single.rs | 11 +++---- .../backend/poly/src/mle/mle_single_owned.rs | 13 +++++--- crates/backend/poly/src/mle/mle_single_ref.rs | 5 ++- crates/backend/poly/src/multilinear_utils.rs | 4 +-- crates/backend/poly/src/next_mle.rs | 7 ++-- crates/backend/poly/src/utils.rs | 10 +++--- crates/backend/poly/src/wrappers.rs | 2 ++ .../sumcheck/src/product_computation.rs | 8 ++--- crates/backend/sumcheck/src/prove.rs | 11 +++---- crates/backend/sumcheck/src/sc_computation.rs | 33 +++++++++---------- crates/backend/sumcheck/src/split_eq.rs | 6 ++-- crates/backend/sumcheck/src/verify.rs | 3 +- crates/lean_prover/src/verify_execution.rs | 2 +- crates/lean_vm/src/tables/table_trait.rs | 6 ++-- crates/lean_vm/src/tables/utils.rs | 2 +- crates/sub_protocols/src/air_sumcheck.rs | 22 ++++++------- .../sub_protocols/src/quotient_gkr/layers.rs | 12 +++---- crates/sub_protocols/src/quotient_gkr/mod.rs | 17 ++++------ .../src/quotient_gkr/sumcheck_utils.rs | 16 ++++----- crates/whir/src/commit.rs | 10 +++--- crates/whir/src/config.rs | 2 ++ crates/whir/src/open.rs | 30 ++++++++--------- crates/whir/src/utils.rs | 11 +++---- crates/whir/src/verify.rs | 16 ++++----- 37 files changed, 185 insertions(+), 199 deletions(-) create mode 100644 crates/backend/koala-bear/src/extension_bound.rs diff --git a/crates/backend/air/src/constraint_folder/normal.rs b/crates/backend/air/src/constraint_folder/normal.rs index 68c627fc2..8bb685f70 100644 --- a/crates/backend/air/src/constraint_folder/normal.rs +++ b/crates/backend/air/src/constraint_folder/normal.rs @@ -3,7 +3,7 @@ use field::*; use poly::*; #[derive(Debug)] -pub struct ConstraintFolder<'a, IF, EF: ExtensionField>, ExtraData: AlphaPowers> { +pub struct ConstraintFolder<'a, IF, EF: KoalaBearExtension, ExtraData: AlphaPowers> { pub flat: &'a [IF], pub shift: &'a [IF], pub extra_data: &'a ExtraData, @@ -13,7 +13,7 @@ pub struct ConstraintFolder<'a, IF, EF: ExtensionField>, ExtraData: Alpha impl<'a, IF, EF, ExtraData> ConstraintFolder<'a, IF, EF, ExtraData> where - EF: ExtensionField>, + EF: KoalaBearExtension, ExtraData: AlphaPowers, { pub fn new(flat: &'a [IF], shift: &'a [IF], extra_data: &'a ExtraData) -> Self { @@ -30,7 +30,7 @@ where impl<'a, IF, EF, ExtraData> AirBuilder for ConstraintFolder<'a, IF, EF, ExtraData> where IF: Algebra> + 'static, - EF: Field + ExtensionField> + Mul + Add, + EF: Field + KoalaBearExtension + Mul + Add, ExtraData: AlphaPowers, { type F = PF; diff --git a/crates/backend/air/src/constraint_folder/packed.rs b/crates/backend/air/src/constraint_folder/packed.rs index 2aa9ed734..0e6729675 100644 --- a/crates/backend/air/src/constraint_folder/packed.rs +++ b/crates/backend/air/src/constraint_folder/packed.rs @@ -2,6 +2,10 @@ use crate::*; use field::*; use poly::*; +// Keeps `ExtensionField>` where the rest of the stack uses `KoalaBearExtension`: the +// `low_degree_block` override below names `&mut Self` in a higher-ranked closure bound, and there +// the solver won't see through a blanket-implemented alias, so every bound on the impl stops +// applying. Callers using `KoalaBearExtension` satisfy this weaker bound anyway. #[derive(Debug)] pub struct ConstraintFolderPacked<'a, IF, EF: ExtensionField>, ExtraData: AlphaPowers> { pub flat: &'a [IF], diff --git a/crates/backend/fiat-shamir/src/prover.rs b/crates/backend/fiat-shamir/src/prover.rs index 94d9fd714..748beb087 100644 --- a/crates/backend/fiat-shamir/src/prover.rs +++ b/crates/backend/fiat-shamir/src/prover.rs @@ -3,8 +3,9 @@ use crate::{MerklePaths, PrunedMerklePaths, *}; use field::Field; use field::PackedValue; use field::PrimeCharacteristicRing; +use field::PrimeField64; use field::integers::QuotientMap; -use field::{ExtensionField, PrimeField64}; +use koala_bear::KoalaBearExtension; use koala_bear::symmetric::Permutation; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::time::Duration; @@ -24,16 +25,13 @@ pub fn reset_pow_grinding_time() { } #[derive(Debug)] -pub struct ProverState>, P> { +pub struct ProverState { challenger: Challenger, P>, transcript: Vec>, merkle_paths: Vec, PF>>, } -impl>, P: Permutation<[PF; WIDTH]>> ProverState -where - PF: PrimeField64, -{ +impl; WIDTH]>> ProverState { #[must_use] pub fn new(permutation: P, capacity: [PF; CAPACITY]) -> Self { assert!(EF::DIMENSION <= RATE); @@ -52,10 +50,7 @@ where } } -impl>, P: Permutation<[PF; WIDTH]>> ChallengeSampler for ProverState -where - PF: PrimeField64, -{ +impl; WIDTH]>> ChallengeSampler for ProverState { fn sample_vec(&mut self, len: usize) -> Vec { sample_vec(&mut self.challenger, len) } @@ -65,10 +60,8 @@ where } } -impl>, P: Permutation<[PF; WIDTH]> + Permutation<[ as Field>::Packing; WIDTH]>> +impl; WIDTH]> + Permutation<[ as Field>::Packing; WIDTH]>> FSProver for ProverState -where - PF: PrimeField64, { fn add_base_scalars(&mut self, scalars: &[PF]) { self.challenger.observe_many(scalars); diff --git a/crates/backend/fiat-shamir/src/traits.rs b/crates/backend/fiat-shamir/src/traits.rs index 722ea761e..1da4f7e29 100644 --- a/crates/backend/fiat-shamir/src/traits.rs +++ b/crates/backend/fiat-shamir/src/traits.rs @@ -1,4 +1,4 @@ -use field::ExtensionField; +use koala_bear::KoalaBearExtension; use crate::{ MerkleOpening, MerklePath, PF, ProofError, ProofResult, flatten_scalars_to_base, pack_scalars_to_extension, @@ -12,7 +12,7 @@ pub trait ChallengeSampler { fn sample_in_range(&mut self, bits: usize, n_samples: usize) -> Vec; } -pub trait FSProver>>: ChallengeSampler { +pub trait FSProver: ChallengeSampler { fn state(&self) -> String; fn add_base_scalars(&mut self, scalars: &[PF]); fn observe_scalars(&mut self, scalars: &[PF]); @@ -43,7 +43,7 @@ pub trait FSProver>>: ChallengeSampler { } } -pub trait FSVerifier>>: ChallengeSampler { +pub trait FSVerifier: ChallengeSampler { fn state(&self) -> String; fn next_base_scalars_vec(&mut self, n: usize) -> Result>, ProofError>; fn observe_scalars(&mut self, scalars: &[PF]); diff --git a/crates/backend/fiat-shamir/src/verifier.rs b/crates/backend/fiat-shamir/src/verifier.rs index 1b5a81234..ed4829a6b 100644 --- a/crates/backend/fiat-shamir/src/verifier.rs +++ b/crates/backend/fiat-shamir/src/verifier.rs @@ -5,17 +5,17 @@ use crate::{ *, }; use field::PrimeCharacteristicRing; -use field::{ExtensionField, PrimeField64}; +use field::PrimeField64; +use koala_bear::KoalaBearExtension; use koala_bear::symmetric::Permutation; use koala_bear::{KoalaBear, default_koalabear_poseidon1_16}; -use std::any::TypeId; use std::collections::VecDeque; use std::iter::repeat_n; use symetric::CAPACITY; use symetric::RATE; use symetric::WIDTH; -pub struct VerifierState>, P> { +pub struct VerifierState { challenger: Challenger, P>, transcript: Vec>, transcript_offset: usize, @@ -25,10 +25,7 @@ pub struct VerifierState>, P> { raw_transcript: Vec>, // reconstructed during the proof verification, it's the format that the zkVM recursion program expects (no Merkle pruning, no sumcheck optimization to send less data, etc) } -impl>, P: Permutation<[PF; WIDTH]>> VerifierState -where - PF: PrimeField64, -{ +impl; WIDTH]>> VerifierState { pub fn new(proof: Proof>, permutation: P, capacity: [PF; CAPACITY]) -> Result { Ok(Self { challenger: Challenger::new(permutation, capacity), @@ -75,16 +72,12 @@ where Ok(scalars) } - #[allow(clippy::missing_transmute_annotations)] fn restore_merkle_paths( - paths: PrunedMerklePaths, PF>, + paths: PrunedMerklePaths, indices: &[usize], merkle_height: usize, leaf_len: usize, - ) -> Option>>> { - assert_eq!(TypeId::of::>(), TypeId::of::()); - // SAFETY: We've confirmed PF == KoalaBear - let paths: PrunedMerklePaths = unsafe { std::mem::transmute(paths) }; + ) -> Option>> { let perm = default_koalabear_poseidon1_16(); let hash_fn = |data: &[KoalaBear]| symetric::hash_slice_rtl::<_, _, 16, 8, DIGEST_LEN_FE>(&perm, data); let combine_fn = |left: &[KoalaBear; DIGEST_LEN_FE], right: &[KoalaBear; DIGEST_LEN_FE]| { @@ -100,15 +93,11 @@ where path: path.sibling_hashes, }) .collect(); - // SAFETY: PF == KoalaBear - Some(unsafe { std::mem::transmute(openings) }) + Some(openings) } } -impl>, P: Permutation<[PF; WIDTH]>> ChallengeSampler for VerifierState -where - PF: PrimeField64, -{ +impl; WIDTH]>> ChallengeSampler for VerifierState { fn sample_vec(&mut self, len: usize) -> Vec { sample_vec(&mut self.challenger, len) } @@ -117,10 +106,7 @@ where } } -impl>, P: Permutation<[PF; WIDTH]>> FSVerifier for VerifierState -where - PF: PrimeField64, -{ +impl; WIDTH]>> FSVerifier for VerifierState { fn state(&self) -> String { format!( "state {} (offset: {}, merkle_idx: {})", diff --git a/crates/backend/koala-bear/src/extension_bound.rs b/crates/backend/koala-bear/src/extension_bound.rs new file mode 100644 index 000000000..b47847671 --- /dev/null +++ b/crates/backend/koala-bear/src/extension_bound.rs @@ -0,0 +1,13 @@ +use field::{ExtensionField, Field, PrimeCharacteristicRing}; + +use crate::KoalaBear; + +pub trait KoalaBearExtension: + Field + ExtensionField + PrimeCharacteristicRing +{ +} + +impl + PrimeCharacteristicRing> KoalaBearExtension + for T +{ +} diff --git a/crates/backend/koala-bear/src/lib.rs b/crates/backend/koala-bear/src/lib.rs index 959ed3ada..44932248b 100644 --- a/crates/backend/koala-bear/src/lib.rs +++ b/crates/backend/koala-bear/src/lib.rs @@ -4,6 +4,7 @@ extern crate alloc; +mod extension_bound; mod koala_bear; pub mod monty_31; mod poseidon1_koalabear_16; @@ -22,6 +23,7 @@ mod x86_64_avx2; #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] mod x86_64_avx512; +pub use extension_bound::*; pub use koala_bear::*; pub use monty_31::*; pub use poseidon1_koalabear_16::*; diff --git a/crates/backend/poly/Cargo.toml b/crates/backend/poly/Cargo.toml index eb0f917f0..65b610438 100644 --- a/crates/backend/poly/Cargo.toml +++ b/crates/backend/poly/Cargo.toml @@ -5,6 +5,7 @@ edition.workspace = true [dependencies] field = { path = "../field", package = "field" } +koala-bear = { path = "../koala-bear", package = "koala-bear" } utils = { path = "../utils", package = "utils" } system-info.workspace = true parallel.workspace = true @@ -14,4 +15,3 @@ rand.workspace = true serde.workspace = true [dev-dependencies] -koala-bear = { path = "../koala-bear", package = "koala-bear" } diff --git a/crates/backend/poly/src/eq_mle.rs b/crates/backend/poly/src/eq_mle.rs index 3b1e12960..288a8fcee 100644 --- a/crates/backend/poly/src/eq_mle.rs +++ b/crates/backend/poly/src/eq_mle.rs @@ -59,11 +59,11 @@ fn par_eval_eq( /// defined on the boolean hypercube by: ∀ (x_1, ..., x_n) ∈ {0, 1}^n, /// P(x_1, ..., x_n) = Π_{i=1}^{n} (x_i.α_i + (1 - x_i).(1 - α_i)) /// (often denoted as P(x) = eq(x, evals)) -pub fn eval_eq>>(eval: &[F]) -> ArenaVec { +pub fn eval_eq(eval: &[F]) -> ArenaVec { eval_eq_scaled(eval, F::ONE) } -pub fn eval_eq_scaled>>(eval: &[F], scalar: F) -> ArenaVec { +pub fn eval_eq_scaled(eval: &[F], scalar: F) -> ArenaVec { // SAFETY: fully written by `compute_eval_eq`. let mut out = unsafe { ArenaVec::uninitialized(1 << eval.len()) }; compute_eval_eq::, F, false>(eval, &mut out, scalar); @@ -71,17 +71,17 @@ pub fn eval_eq_scaled>>(eval: &[F], scalar: F) -> ArenaV } /// Single-threaded, arena-free [`eval_eq`] (verifier side). -pub fn eval_eq_sequential>>(eval: &[F]) -> Vec { +pub fn eval_eq_sequential(eval: &[F]) -> Vec { let mut out = vec![F::ZERO; 1 << eval.len()]; eval_eq_basic::, F, F, false>(eval, &mut out, F::ONE); out } -pub fn eval_eq_packed>>(eval: &[F]) -> ArenaVec> { +pub fn eval_eq_packed(eval: &[F]) -> ArenaVec> { eval_eq_packed_scaled(eval, F::ONE) } -pub fn eval_eq_packed_scaled>>(eval: &[F], scalar: F) -> ArenaVec> { +pub fn eval_eq_packed_scaled(eval: &[F], scalar: F) -> ArenaVec> { // Alloc memory without initializing it to zero. // This is safe because we overwrite it inside `compute_eval_eq_packed`. let mut out = unsafe { ArenaVec::uninitialized(1 << (eval.len() - packing_log_width::())) }; @@ -89,7 +89,7 @@ pub fn eval_eq_packed_scaled>>(eval: &[F], scalar: F) -> out } -pub fn compute_sparse_eval_eq>>(selector: usize, eval: &[F], out: &mut [F], scalar: F) { +pub fn compute_sparse_eval_eq(selector: usize, eval: &[F], out: &mut [F], scalar: F) { if eval.is_empty() { out[selector] += scalar; return; @@ -103,7 +103,7 @@ pub fn compute_sparse_eval_eq>>(selector: usize, eval: & pub fn compute_sparse_eval_eq_packed(selector: usize, eval: &[EF], out: &mut [EFPacking], scalar: EF) where - EF: ExtensionField>, + EF: KoalaBearExtension, { let log_packing = packing_log_width::(); if eval.len() < log_packing { @@ -171,7 +171,7 @@ where #[inline] pub fn compute_eval_eq_packed(eval: &[EF], out: &mut [EF::ExtensionPacking], scalar: EF) where - EF: ExtensionField>, + EF: KoalaBearExtension, { // `packing_width` may be 1 (e.g. Goldilocks on Neon, or without `target-cpu=native`), // so nothing here may assume it is > 1. @@ -879,7 +879,7 @@ pub fn compute_eval_eq_packed_dual( scalar_a: EF, scalar_b: EF, ) where - EF: ExtensionField>, + EF: KoalaBearExtension, { let packing_width = packing_width::(); let log_packing_width = log2_strict_usize(packing_width); diff --git a/crates/backend/poly/src/evals.rs b/crates/backend/poly/src/evals.rs index 6a917ee17..d054a710f 100644 --- a/crates/backend/poly/src/evals.rs +++ b/crates/backend/poly/src/evals.rs @@ -167,7 +167,7 @@ const EVAL_BASE_PACKED_MIN_VARS: usize = 6; /// subtrees, which is several times more base multiplications in total. pub fn eval_base_packed(evals: &[PF], point: &[EF]) -> EF where - EF: ExtensionField>, + EF: KoalaBearExtension, { debug_assert_eq!(evals.len(), 1 << point.len()); if point.len() < EVAL_BASE_PACKED_MIN_VARS { @@ -200,7 +200,7 @@ where pub fn eval_packed(evals: &[EFPacking], point: &[EF]) -> EF where - EF: ExtensionField>, + EF: KoalaBearExtension, { let log_width = packing_log_width::(); let res_packed: EFPacking = eval_multilinear_generic::<_, _, _, _, _, _, PARALLEL>( diff --git a/crates/backend/poly/src/mle/mle_group.rs b/crates/backend/poly/src/mle/mle_group.rs index 912ec271c..5ac631b63 100644 --- a/crates/backend/poly/src/mle/mle_group.rs +++ b/crates/backend/poly/src/mle/mle_group.rs @@ -1,25 +1,24 @@ use crate::*; -use field::ExtensionField; #[derive(Debug)] -pub enum MleGroup<'a, EF: ExtensionField>> { +pub enum MleGroup<'a, EF: KoalaBearExtension> { Owned(MleGroupOwned), Ref(MleGroupRef<'a, EF>), } -impl<'a, EF: ExtensionField>> From> for MleGroup<'a, EF> { +impl<'a, EF: KoalaBearExtension> From> for MleGroup<'a, EF> { fn from(owned: MleGroupOwned) -> Self { MleGroup::Owned(owned) } } -impl<'a, EF: ExtensionField>> From> for MleGroup<'a, EF> { +impl<'a, EF: KoalaBearExtension> From> for MleGroup<'a, EF> { fn from(r: MleGroupRef<'a, EF>) -> Self { MleGroup::Ref(r) } } -impl<'a, EF: ExtensionField>> MleGroup<'a, EF> { +impl<'a, EF: KoalaBearExtension> MleGroup<'a, EF> { pub fn by_ref(&'a self) -> MleGroupRef<'a, EF> { match self { Self::Owned(owned) => owned.by_ref(), diff --git a/crates/backend/poly/src/mle/mle_group_owned.rs b/crates/backend/poly/src/mle/mle_group_owned.rs index d4e57ed1d..7487f189d 100644 --- a/crates/backend/poly/src/mle/mle_group_owned.rs +++ b/crates/backend/poly/src/mle/mle_group_owned.rs @@ -1,17 +1,16 @@ use crate::*; use ::utils::log2_strict_usize; -use field::ExtensionField; use zk_alloc::ArenaVec; #[derive(Debug)] -pub enum MleGroupOwned>> { +pub enum MleGroupOwned { Base(Vec>>), Extension(Vec>), BasePacked(Vec>>), ExtensionPacked(Vec>>), } -impl>> MleGroupOwned { +impl MleGroupOwned { pub fn as_extension_packed_mut(&mut self) -> Option<&mut Vec>>> { match self { Self::ExtensionPacked(e) => Some(e), diff --git a/crates/backend/poly/src/mle/mle_group_ref.rs b/crates/backend/poly/src/mle/mle_group_ref.rs index 335399168..f5c38a6e9 100644 --- a/crates/backend/poly/src/mle/mle_group_ref.rs +++ b/crates/backend/poly/src/mle/mle_group_ref.rs @@ -1,18 +1,17 @@ use crate::*; use ::utils::log2_strict_usize; -use field::ExtensionField; use field::PackedValue; use zk_alloc::ArenaVec; #[derive(Debug)] -pub enum MleGroupRef<'a, EF: ExtensionField>> { +pub enum MleGroupRef<'a, EF: KoalaBearExtension> { Base(Vec<&'a [PF]>), Extension(Vec<&'a [EF]>), BasePacked(Vec<&'a [PFPacking]>), ExtensionPacked(Vec<&'a [EFPacking]>), } -impl<'a, EF: ExtensionField>> MleGroupRef<'a, EF> { +impl<'a, EF: KoalaBearExtension> MleGroupRef<'a, EF> { pub const fn group_size(&self) -> usize { match self { Self::Base(v) => v.len(), diff --git a/crates/backend/poly/src/mle/mle_single.rs b/crates/backend/poly/src/mle/mle_single.rs index 6afa328a8..580381b7b 100644 --- a/crates/backend/poly/src/mle/mle_single.rs +++ b/crates/backend/poly/src/mle/mle_single.rs @@ -1,13 +1,12 @@ use crate::*; -use field::ExtensionField; #[derive(Debug)] -pub enum Mle<'a, EF: ExtensionField>> { +pub enum Mle<'a, EF: KoalaBearExtension> { Owned(MleOwned), Ref(MleRef<'a, EF>), } -impl>> Clone for Mle<'_, EF> { +impl Clone for Mle<'_, EF> { fn clone(&self) -> Self { match self { Self::Owned(o) => Self::Owned(o.clone()), @@ -15,19 +14,19 @@ impl>> Clone for Mle<'_, EF> { } } } -impl>> From> for Mle<'_, EF> { +impl From> for Mle<'_, EF> { fn from(value: MleOwned) -> Self { Self::Owned(value) } } -impl<'a, EF: ExtensionField>> From> for Mle<'a, EF> { +impl<'a, EF: KoalaBearExtension> From> for Mle<'a, EF> { fn from(value: MleRef<'a, EF>) -> Self { Self::Ref(value) } } -impl<'a, EF: ExtensionField>> Mle<'a, EF> { +impl<'a, EF: KoalaBearExtension> Mle<'a, EF> { pub fn by_ref(&'a self) -> MleRef<'a, EF> { match self { Self::Owned(owned) => owned.by_ref(), diff --git a/crates/backend/poly/src/mle/mle_single_owned.rs b/crates/backend/poly/src/mle/mle_single_owned.rs index 538060296..55c1f9234 100644 --- a/crates/backend/poly/src/mle/mle_single_owned.rs +++ b/crates/backend/poly/src/mle/mle_single_owned.rs @@ -1,23 +1,26 @@ -use crate::{EFPacking, Mle, MleRef, MultilinearPoint, PF, PFPacking, pack_extension, packing_width, unpack_extension}; +use crate::{ + EFPacking, KoalaBearExtension, Mle, MleRef, MultilinearPoint, PF, PFPacking, pack_extension, packing_width, + unpack_extension, +}; +use field::PackedFieldExtension; use field::PackedValue; -use field::{ExtensionField, PackedFieldExtension}; use zk_alloc::ArenaVec; #[derive(Debug, Clone)] -pub enum MleOwned>> { +pub enum MleOwned { Base(ArenaVec>), Extension(ArenaVec), BasePacked(ArenaVec>), ExtensionPacked(ArenaVec>), } -impl>> Default for MleOwned { +impl Default for MleOwned { fn default() -> Self { Self::Base(ArenaVec::new()) } } -impl>> MleOwned { +impl MleOwned { pub fn by_ref<'a>(&'a self) -> MleRef<'a, EF> { match self { Self::Base(v) => MleRef::Base(v), diff --git a/crates/backend/poly/src/mle/mle_single_ref.rs b/crates/backend/poly/src/mle/mle_single_ref.rs index 4c630b826..844855226 100644 --- a/crates/backend/poly/src/mle/mle_single_ref.rs +++ b/crates/backend/poly/src/mle/mle_single_ref.rs @@ -1,18 +1,17 @@ use crate::*; use ::utils::log2_strict_usize; -use field::ExtensionField; use field::PackedValue; use zk_alloc::ArenaVec; #[derive(Debug)] -pub enum MleRef<'a, EF: ExtensionField>> { +pub enum MleRef<'a, EF: KoalaBearExtension> { Base(&'a [PF]), Extension(&'a [EF]), BasePacked(&'a [PFPacking]), ExtensionPacked(&'a [EFPacking]), } -impl<'a, EF: ExtensionField>> MleRef<'a, EF> { +impl<'a, EF: KoalaBearExtension> MleRef<'a, EF> { pub const fn packed_len(&self) -> usize { match self { Self::Base(v) => v.len(), diff --git a/crates/backend/poly/src/multilinear_utils.rs b/crates/backend/poly/src/multilinear_utils.rs index d96cbc57e..7f02e8522 100644 --- a/crates/backend/poly/src/multilinear_utils.rs +++ b/crates/backend/poly/src/multilinear_utils.rs @@ -1,7 +1,7 @@ use field::{ExtensionField, Field, dot_product}; use utils::*; -use crate::{EFPacking, EvaluationsList as _, MultilinearPoint, PF, PFPacking}; +use crate::{EFPacking, EvaluationsList as _, KoalaBearExtension, MultilinearPoint, PFPacking}; pub fn multilinear_eval_constants_at_right(limit: usize, point: &[F]) -> F { let n_vars = point.len(); @@ -68,7 +68,7 @@ pub fn finger_print(domainsep: EF, data: &[EF], alphas_eq_poly: &[EF] /// Packed variant of [`finger_print`]. #[inline(always)] -pub fn finger_print_packed>>( +pub fn finger_print_packed( domainsep: PFPacking, data: &[PFPacking], alphas_packed: &[EFPacking], diff --git a/crates/backend/poly/src/next_mle.rs b/crates/backend/poly/src/next_mle.rs index af8c08e96..213ce6bba 100644 --- a/crates/backend/poly/src/next_mle.rs +++ b/crates/backend/poly/src/next_mle.rs @@ -1,7 +1,7 @@ -use field::{ExtensionField, Field, PrimeCharacteristicRing}; +use field::Field; use zk_alloc::ArenaVec; -use crate::{PF, eval_eq_scaled}; +use crate::{KoalaBearExtension, eval_eq_scaled}; /// Evaluates the "next" multilinear polynomial at two n-variable points (x, y). /// @@ -33,9 +33,8 @@ pub fn next_mle(x: &[F], y: &[F]) -> F { /// /// This is the "folded" version: the first argument (outer_challenges) is fixed, /// and the result is a vector indexed by the second argument. -pub fn matrix_next_mle_folded>>(outer_challenges: &[F]) -> ArenaVec +pub fn matrix_next_mle_folded(outer_challenges: &[F]) -> ArenaVec where - PF: PrimeCharacteristicRing, { let n = outer_challenges.len(); let mut res = unsafe { ArenaVec::::zeroed(1 << n) }; diff --git a/crates/backend/poly/src/utils.rs b/crates/backend/poly/src/utils.rs index 4f94b3f59..3fcd2667c 100644 --- a/crates/backend/poly/src/utils.rs +++ b/crates/backend/poly/src/utils.rs @@ -3,13 +3,13 @@ use std::ops::{Add, Sub}; use field::*; use zk_alloc::{ArenaVec, OwnedBuffer}; -use crate::{EFPacking, PF, PFPacking}; +use crate::{EFPacking, KoalaBearExtension, PFPacking}; pub const PARALLEL_THRESHOLD: usize = 1 << 9; /// AoS->SoA transpose of `slice` into the already-sized packed buffer `out` (`out.len()` /// packed elements, each consuming `packing_width` scalars). -fn fill_packed_extension>>(slice: &[EF], out: &mut [EFPacking]) { +fn fill_packed_extension(slice: &[EF], out: &mut [EFPacking]) { let width = packing_width::(); let write = |slot: &mut EFPacking, chunk: &[EF]| { *slot = EFPacking::::from_ext_slice(chunk); @@ -25,13 +25,13 @@ fn fill_packed_extension>>(slice: &[EF], out: &mut [EF } } -pub fn pack_extension>, B: OwnedBuffer>>(slice: &[EF]) -> B { +pub fn pack_extension>>(slice: &[EF]) -> B { B::build(slice.len() / packing_width::(), |out| { fill_packed_extension(slice, out) }) } -fn fill_unpacked_extension>>(vec: &[EFPacking], out: &mut [EF]) { +fn fill_unpacked_extension(vec: &[EFPacking], out: &mut [EF]) { let width = packing_width::(); let total = out.len(); let write = |out_chunk: &mut [EF], x: &EFPacking| { @@ -56,7 +56,7 @@ fn fill_unpacked_extension>>(vec: &[EFPacking], ou } } -pub fn unpack_extension>, B: OwnedBuffer>(vec: &[EFPacking]) -> B { +pub fn unpack_extension>(vec: &[EFPacking]) -> B { B::build(vec.len() * packing_width::(), |out| { fill_unpacked_extension(vec, out) }) diff --git a/crates/backend/poly/src/wrappers.rs b/crates/backend/poly/src/wrappers.rs index 8ec72d4d9..167ab40a4 100644 --- a/crates/backend/poly/src/wrappers.rs +++ b/crates/backend/poly/src/wrappers.rs @@ -4,3 +4,5 @@ pub type PF = ::PrimeSubfield; pub type FPacking = ::Packing; pub type PFPacking = as Field>::Packing; pub type EFPacking = >>::ExtensionPacking; + +pub use koala_bear::KoalaBearExtension; diff --git a/crates/backend/sumcheck/src/product_computation.rs b/crates/backend/sumcheck/src/product_computation.rs index b054e09a1..167fe734e 100644 --- a/crates/backend/sumcheck/src/product_computation.rs +++ b/crates/backend/sumcheck/src/product_computation.rs @@ -9,7 +9,7 @@ use crate::{SumcheckComputation, packing_unpack_sum, sumcheck_prove_many_rounds} #[derive(Debug)] pub struct ProductComputation; -impl>> SumcheckComputation for ProductComputation { +impl SumcheckComputation for ProductComputation { type ExtraData = Vec; fn degree(&self) -> usize { @@ -34,7 +34,7 @@ impl>> SumcheckComputation for ProductComputation } #[instrument(skip_all)] -pub fn run_product_sumcheck>>( +pub fn run_product_sumcheck( pol_a: &MleRef<'_, EF>, // evals pol_b: &MleRef<'_, EF>, // weights prover_state: &mut impl FSProver, @@ -69,7 +69,7 @@ pub fn run_product_sumcheck>>( /// Rounds 1+ of the product sumcheck, for callers that computed round 0 themselves. /// `sum` is the running sum after binding `r1`. -pub fn run_product_sumcheck_from_round1>>( +pub fn run_product_sumcheck_from_round1( pol_a: &MleRef<'_, EF>, // evals pol_b: &MleRef<'_, EF>, // weights prover_state: &mut impl FSProver, @@ -273,7 +273,7 @@ where /// Algo 3 of https://eprint.iacr.org/2024/1046.pdf. Requires n_rounds >= 3. #[allow(clippy::too_many_arguments)] -pub fn run_product_sumcheck_from_round1_delayed>>( +pub fn run_product_sumcheck_from_round1_delayed( evals: &[PFPacking], weights: &[EFPacking], prover_state: &mut impl FSProver, diff --git a/crates/backend/sumcheck/src/prove.rs b/crates/backend/sumcheck/src/prove.rs index 5e0bdd5c1..b3a64320e 100644 --- a/crates/backend/sumcheck/src/prove.rs +++ b/crates/backend/sumcheck/src/prove.rs @@ -1,6 +1,5 @@ use air::AlphaPowers; use fiat_shamir::*; -use field::ExtensionField; use field::PrimeCharacteristicRing; use poly::*; @@ -18,7 +17,7 @@ pub fn sumcheck_fold_and_prove<'a, EF, SC, M: Into>>( store_intermediate_foldings: bool, ) -> (MultilinearPoint, Vec, EF) where - EF: ExtensionField>, + EF: KoalaBearExtension, SC: SumcheckComputation + 'static, SC::ExtraData: AlphaPowers, { @@ -70,7 +69,7 @@ pub fn sumcheck_prove_many_rounds<'a, EF, SC, M: Into>>( pow_bits: usize, ) -> (MultilinearPoint, MleGroupOwned, EF) where - EF: ExtensionField>, + EF: KoalaBearExtension, SC: SumcheckComputation + 'static, SC::ExtraData: AlphaPowers, { @@ -137,7 +136,7 @@ pub fn compute_round_polynomial<'a, EF, SC>( missing_mul_factor: Option, ) -> DensePolynomial where - EF: ExtensionField>, + EF: KoalaBearExtension, SC: SumcheckComputation + 'static, SC::ExtraData: AlphaPowers, { @@ -197,7 +196,7 @@ fn compute_and_send_polynomial<'a, EF, SC>( missing_mul_factor: Option, ) -> DensePolynomial where - EF: ExtensionField>, + EF: KoalaBearExtension, SC: SumcheckComputation + 'static, SC::ExtraData: AlphaPowers, { @@ -216,7 +215,7 @@ where } #[allow(clippy::too_many_arguments)] -pub fn on_challenge_received<'a, EF: ExtensionField>>( +pub fn on_challenge_received<'a, EF: KoalaBearExtension>( multilinears: &mut MleGroup<'a, EF>, n_vars: &mut usize, eq_factor: &mut Option<(Vec, SplitEq)>, diff --git a/crates/backend/sumcheck/src/sc_computation.rs b/crates/backend/sumcheck/src/sc_computation.rs index 3879e18be..0e2f85732 100644 --- a/crates/backend/sumcheck/src/sc_computation.rs +++ b/crates/backend/sumcheck/src/sc_computation.rs @@ -13,7 +13,7 @@ fn add_assign_vec(mut a: Vec, b: Vec) -> Vec { a } -pub trait SumcheckComputation>>: Sync { +pub trait SumcheckComputation: Sync { type ExtraData: Send + Sync + 'static; fn degree(&self) -> usize; @@ -34,7 +34,7 @@ macro_rules! impl_air_eval { impl SumcheckComputation for A where - EF: ExtensionField>, + EF: KoalaBearExtension, A: Send + Sync + Air, A::ExtraData: AlphaPowers, { @@ -65,17 +65,14 @@ where } } -fn build_evals>>( - sums: impl IntoIterator, - missing_mul_factor: Option, -) -> Vec { +fn build_evals(sums: impl IntoIterator, missing_mul_factor: Option) -> Vec { sums.into_iter() .map(|sum| missing_mul_factor.map_or(sum, |f| sum * f)) .collect() } #[inline(always)] -fn poly_to_evals>>(poly: &DensePolynomial) -> Vec { +fn poly_to_evals(poly: &DensePolynomial) -> Vec { vec![poly.coeffs[0], poly.evaluate(EF::TWO)] } @@ -85,16 +82,16 @@ pub(crate) fn identity_decompose(e: EF) -> Vec { } #[inline(always)] -pub(crate) fn packing_decompose>>(e: EFPacking) -> Vec { +pub(crate) fn packing_decompose(e: EFPacking) -> Vec { EFPacking::::to_ext_iter([e]).collect() } #[inline(always)] -pub fn packing_unpack_sum>>(s: EFPacking) -> EF { +pub fn packing_unpack_sum(s: EFPacking) -> EF { EFPacking::::to_ext_iter([s]).sum::() } -fn handle_product_computation<'a, EF: ExtensionField>>(group: &MleGroupRef<'a, EF>, sum: EF) -> Vec { +fn handle_product_computation<'a, EF: KoalaBearExtension>(group: &MleGroupRef<'a, EF>, sum: EF) -> Vec { let poly = match group { MleGroupRef::Extension(multilinears) => { compute_product_sumcheck_polynomial(multilinears[0], multilinears[1], sum, identity_decompose) @@ -108,7 +105,7 @@ fn handle_product_computation<'a, EF: ExtensionField>>(group: &MleGroupRe } #[allow(clippy::type_complexity)] -fn handle_product_computation_with_fold<'a, EF: ExtensionField>>( +fn handle_product_computation_with_fold<'a, EF: KoalaBearExtension>( group: &MleGroupRef<'a, EF>, prev_folding_factor: EF, sum: EF, @@ -139,7 +136,7 @@ fn handle_product_computation_with_fold<'a, EF: ExtensionField>>( (poly_to_evals(&poly), folded_f) } -pub struct SumcheckComputeParams<'a, EF: ExtensionField>, SC: SumcheckComputation> { +pub struct SumcheckComputeParams<'a, EF: KoalaBearExtension, SC: SumcheckComputation> { pub split_eq: Option<&'a SplitEq>, pub computation: &'a SC, pub extra_data: &'a SC::ExtraData, @@ -147,7 +144,7 @@ pub struct SumcheckComputeParams<'a, EF: ExtensionField>, SC: SumcheckCom pub sum: EF, } -pub fn sumcheck_compute<'a, EF: ExtensionField>, SC>( +pub fn sumcheck_compute<'a, EF: KoalaBearExtension, SC>( group: &MleGroupRef<'a, EF>, params: SumcheckComputeParams<'a, EF, SC>, degree: usize, @@ -250,7 +247,7 @@ where } #[allow(clippy::type_complexity)] -pub fn fold_and_sumcheck_compute<'a, EF: ExtensionField>, SC>( +pub fn fold_and_sumcheck_compute<'a, EF: KoalaBearExtension, SC>( prev_folding_factor: EF, group: &MleGroupRef<'a, EF>, params: SumcheckComputeParams<'a, EF, SC>, @@ -388,7 +385,7 @@ fn sumcheck_compute_core( unpack_sum: impl Fn(EFT) -> EF, ) -> Vec where - EF: ExtensionField>, + EF: KoalaBearExtension, IF: Copy + Sub + Add + AddAssign + Send + Sync, EFT: PrimeCharacteristicRing + Copy @@ -460,7 +457,7 @@ fn sumcheck_fold_and_compute_core( wrap_f: impl FnOnce(Vec>) -> MleGroupOwned, ) -> (Vec, MleGroupOwned) where - EF: ExtensionField>, + EF: KoalaBearExtension, IF: Copy + Send + Sync, FT: PrimeCharacteristicRing + Copy + Sub + Add + Send + Sync, SC: SumcheckComputation, @@ -533,7 +530,7 @@ fn sumcheck_compute_with_split_eq( unpack_sum: impl Fn(EFPacking) -> EF, ) -> Vec where - EF: ExtensionField>, + EF: KoalaBearExtension, SC: SumcheckComputation, { let n_lo = split_eq.n_lo(); @@ -614,7 +611,7 @@ fn sumcheck_fold_and_compute_with_split_eq( wrap_f: impl FnOnce(Vec>>) -> MleGroupOwned, ) -> (Vec, MleGroupOwned) where - EF: ExtensionField>, + EF: KoalaBearExtension, IF: Copy + Send + Sync, SC: SumcheckComputation, { diff --git a/crates/backend/sumcheck/src/split_eq.rs b/crates/backend/sumcheck/src/split_eq.rs index 16a8908b8..b784849d3 100644 --- a/crates/backend/sumcheck/src/split_eq.rs +++ b/crates/backend/sumcheck/src/split_eq.rs @@ -1,9 +1,9 @@ -use field::{ExtensionField, PackedFieldExtension}; +use field::PackedFieldExtension; use poly::*; use zk_alloc::ArenaVec; #[derive(Debug)] -pub struct SplitEq>> { +pub struct SplitEq { pub eq_lo: ArenaVec, pub eq_hi_packed: ArenaVec>, pub log_packed_hi: u32, // = log2(eq_hi_packed.len()), cached for bit-shift in get_packed @@ -11,7 +11,7 @@ pub struct SplitEq>> { pub remainder: ArenaVec, } -impl>> SplitEq { +impl SplitEq { pub fn new(eq_point: &[EF]) -> Self { let n = eq_point.len(); diff --git a/crates/backend/sumcheck/src/verify.rs b/crates/backend/sumcheck/src/verify.rs index 59c0ff4f1..2e8db3bf2 100644 --- a/crates/backend/sumcheck/src/verify.rs +++ b/crates/backend/sumcheck/src/verify.rs @@ -1,8 +1,7 @@ use fiat_shamir::*; -use field::*; use poly::*; -pub fn sumcheck_verify>>( +pub fn sumcheck_verify( verifier_state: &mut impl FSVerifier, n_vars: usize, degree: usize, diff --git a/crates/lean_prover/src/verify_execution.rs b/crates/lean_prover/src/verify_execution.rs index 7b66edac3..912ac807a 100644 --- a/crates/lean_prover/src/verify_execution.rs +++ b/crates/lean_prover/src/verify_execution.rs @@ -231,7 +231,7 @@ pub fn verify_execution( )) } -fn back_loaded_table_contribution>>( +fn back_loaded_table_contribution( bus_point: &[EF], sumcheck_air_point: &[EF], natural_ordering_point: &[EF], diff --git a/crates/lean_vm/src/tables/table_trait.rs b/crates/lean_vm/src/tables/table_trait.rs index 54d931db4..4cbc24951 100644 --- a/crates/lean_vm/src/tables/table_trait.rs +++ b/crates/lean_vm/src/tables/table_trait.rs @@ -159,13 +159,13 @@ pub fn sort_tables_by_height(tables_log_heights: &BTreeMap) -> Vec } #[derive(Debug, Default)] -pub struct ExtraDataForBuses>> { +pub struct ExtraDataForBuses { // GKR quotient challenges pub logup_alphas_eq_poly: Vec, pub logup_alphas_eq_poly_packed: Vec>, pub alpha_powers: Vec, } -impl>> ExtraDataForBuses { +impl ExtraDataForBuses { pub fn new(logup_alphas_eq_poly: &[EF], alpha_powers: Vec) -> Self { let logup_alphas_eq_poly_packed = logup_alphas_eq_poly.iter().map(|a| EFPacking::::from(*a)).collect(); Self { @@ -188,7 +188,7 @@ impl AlphaPowers for ExtraDataForBuses { } } -impl>> ExtraDataForBuses { +impl ExtraDataForBuses { pub fn transmute_bus_data(&self) -> &Vec { if TypeId::of::() == TypeId::of::() { unsafe { transmute::<&Vec, &Vec>(&self.logup_alphas_eq_poly) } diff --git a/crates/lean_vm/src/tables/utils.rs b/crates/lean_vm/src/tables/utils.rs index ca1594190..6d6b06336 100644 --- a/crates/lean_vm/src/tables/utils.rs +++ b/crates/lean_vm/src/tables/utils.rs @@ -2,7 +2,7 @@ use backend::*; use crate::ExtraDataForBuses; -pub(crate) fn eval_bus_virtual>>( +pub(crate) fn eval_bus_virtual( builder: &mut AB, extra_data: &ExtraDataForBuses, multiplicity: AB::IF, diff --git a/crates/sub_protocols/src/air_sumcheck.rs b/crates/sub_protocols/src/air_sumcheck.rs index eb9ca0a02..f0bb38cdb 100644 --- a/crates/sub_protocols/src/air_sumcheck.rs +++ b/crates/sub_protocols/src/air_sumcheck.rs @@ -31,7 +31,7 @@ use tracing::info_span; const ENDIANNESS_PIVOT_AIR: usize = 12; -pub trait OuterSumcheckSession>>: Debug { +pub trait OuterSumcheckSession: Debug { fn initial_n_vars(&self) -> usize; fn sum(&self) -> EF; fn bare_degree(&self) -> usize; @@ -42,7 +42,7 @@ pub trait OuterSumcheckSession>>: Debug { } #[derive(Debug)] -pub struct AirSumcheckSession<'a, EF: ExtensionField>, A: Air> +pub struct AirSumcheckSession<'a, EF: KoalaBearExtension, A: Air> where A::ExtraData: AlphaPowers, { @@ -60,7 +60,7 @@ where rounds_done: usize, } -impl<'a, EF: ExtensionField>, A: Air> AirSumcheckSession<'a, EF, A> +impl<'a, EF: KoalaBearExtension, A: Air> AirSumcheckSession<'a, EF, A> where A::ExtraData: AlphaPowers + AlphaPowersMut, { @@ -125,7 +125,7 @@ where impl<'a, EF, A> AirSumcheckSession<'a, EF, A> where - EF: ExtensionField>, + EF: KoalaBearExtension, A: Air + 'static, A::ExtraData: AlphaPowers, { @@ -200,7 +200,7 @@ where impl<'a, EF, A> OuterSumcheckSession for AirSumcheckSession<'a, EF, A> where - EF: ExtensionField>, + EF: KoalaBearExtension, A: Air + Debug + 'static, A::ExtraData: AlphaPowers + AlphaPowersMut + Debug, { @@ -289,7 +289,7 @@ where } } -fn column_evals>>(multilinears: &MleGroupRef<'_, EF>, i: usize) -> Vec { +fn column_evals(multilinears: &MleGroupRef<'_, EF>, i: usize) -> Vec { match multilinears { MleGroupRef::Base(cols) => cols.iter().map(|c| EF::from(c[i])).collect(), MleGroupRef::Extension(cols) => cols.iter().map(|c| c[i]).collect(), @@ -315,7 +315,7 @@ fn compute_raw_poly<'a, EF, A>( active_count_pairs: usize, ) -> Vec where - EF: ExtensionField>, + EF: KoalaBearExtension, A: Air + 'static, A::ExtraData: AlphaPowers, { @@ -410,7 +410,7 @@ fn compute_raw_poly_degree_split( unpack_sum: UnpackSum, ) -> Vec where - EF: ExtensionField>, + EF: KoalaBearExtension, A: Air + 'static, A::ExtraData: AlphaPowers, IF: Algebra> + Copy + Send + Sync + Sub + AddAssign + PrimeCharacteristicRing + 'static, @@ -558,7 +558,7 @@ fn compute_raw_poly_impl( unpack_sum: UnpackSum, ) -> Vec where - EF: ExtensionField>, + EF: KoalaBearExtension, A: Air + 'static, A::ExtraData: AlphaPowers, IF: Copy + Send + Sync + Sub + AddAssign + PrimeCharacteristicRing, @@ -612,7 +612,7 @@ where acc.into_iter().map(unpack_sum).collect() } -pub fn prove_batched_air_sumcheck<'a, EF: ExtensionField>>( +pub fn prove_batched_air_sumcheck<'a, EF: KoalaBearExtension>( prover_state: &mut impl FSProver, sessions: &mut [Box + 'a>], ) -> MultilinearPoint { @@ -678,7 +678,7 @@ pub fn natural_ordering_point_for_session(sumcheck_air_point: &[EF], l .collect() } -pub fn columns_evals_flat_and_shift>, A: Air>( +pub fn columns_evals_flat_and_shift( air: &A, col_evals: &[EF], natural_ordering_point: &[EF], diff --git a/crates/sub_protocols/src/quotient_gkr/layers.rs b/crates/sub_protocols/src/quotient_gkr/layers.rs index 7a19da588..0aae17a55 100644 --- a/crates/sub_protocols/src/quotient_gkr/layers.rs +++ b/crates/sub_protocols/src/quotient_gkr/layers.rs @@ -2,7 +2,7 @@ use backend::PackedValue; use backend::*; -pub(super) enum LayerStorage<'a, EF: ExtensionField>> { +pub(super) enum LayerStorage<'a, EF: KoalaBearExtension> { Initial { nums: ArenaCow<'a, PFPacking>, dens: ArenaCow<'a, EFPacking>, @@ -19,7 +19,7 @@ pub(super) enum LayerStorage<'a, EF: ExtensionField>> { }, } -impl<'a, EF: ExtensionField>> LayerStorage<'a, EF> { +impl<'a, EF: KoalaBearExtension> LayerStorage<'a, EF> { pub(super) fn convert_to_natural(&self) -> Self { match self { Self::Initial { nums, dens, chunk_log } => { @@ -119,7 +119,7 @@ pub(super) fn bit_reverse_chunks(v: &[T], chunk_log: usiz out } -fn sum_quotients_2_by_2>>(nums: &[EF], dens: &[EF]) -> (ArenaVec, ArenaVec) { +fn sum_quotients_2_by_2(nums: &[EF], dens: &[EF]) -> (ArenaVec, ArenaVec) { assert_eq!(nums.len(), dens.len()); let active_len = nums.len(); let new_active = active_len.div_ceil(2); @@ -150,7 +150,7 @@ fn sum_quotients_2_by_2>>(nums: &[EF], dens: &[EF]) -> (new_nums, new_dens) } -fn sum_quotients_2_by_2_packed_br>, N>( +fn sum_quotients_2_by_2_packed_br( nums: &[N], dens: &[EFPacking], chunk_log: usize, @@ -182,14 +182,14 @@ where (new_nums, new_dens) } -pub(super) fn unpack_and_unreverse_active>>( +pub(super) fn unpack_and_unreverse_active( v: &[EFPacking], chunk_log: usize, ) -> ArenaVec { bit_reverse_chunks(&unpack_extension::>(v), chunk_log) } -fn unpack_base_and_unreverse_active>>(v: &[PFPacking], chunk_log: usize) -> ArenaVec { +fn unpack_base_and_unreverse_active(v: &[PFPacking], chunk_log: usize) -> ArenaVec { let active_unpacked: ArenaVec = PFPacking::::unpack_slice(v).iter().map(|x| EF::from(*x)).collect(); bit_reverse_chunks(&active_unpacked, chunk_log) } diff --git a/crates/sub_protocols/src/quotient_gkr/mod.rs b/crates/sub_protocols/src/quotient_gkr/mod.rs index 0d2153536..7b4494b8a 100644 --- a/crates/sub_protocols/src/quotient_gkr/mod.rs +++ b/crates/sub_protocols/src/quotient_gkr/mod.rs @@ -26,7 +26,7 @@ mod sumcheck_utils; pub const ENDIANNESS_PIVOT_GKR: usize = 12; #[instrument(skip_all, name = "prove GKR")] -pub fn prove_gkr_quotient<'a, EF: ExtensionField>>( +pub fn prove_gkr_quotient<'a, EF: KoalaBearExtension>( prover_state: &mut impl FSProver, nums_br: &'a [PFPacking], // already bit-reversed at `pivot`, ACTIVE prefix only (i.e. length may not be a power of 2) dens_br: &'a [EFPacking], // same as above @@ -75,7 +75,7 @@ pub fn prove_gkr_quotient<'a, EF: ExtensionField>>( (quotient, point) } -fn prove_gkr_layer>>( +fn prove_gkr_layer( prover_state: &mut impl FSProver, layer: &LayerStorage<'_, EF>, claim_point: &MultilinearPoint, // K coords, natural order @@ -138,7 +138,7 @@ fn prove_gkr_layer>>( (MultilinearPoint(q_natural), next_num, next_den) } -fn compute_quotient>>(numerators: &[EF], denominators: &[EF]) -> Option { +fn compute_quotient(numerators: &[EF], denominators: &[EF]) -> Option { let mut acc = EF::ZERO; for (&n, &d) in numerators.iter().zip(denominators) { acc += n * d.try_inverse()?; @@ -146,7 +146,7 @@ fn compute_quotient>>(numerators: &[EF], denominators: Some(acc) } -pub fn verify_gkr_quotient>>( +pub fn verify_gkr_quotient( verifier_state: &mut impl FSVerifier, n_vars: usize, ) -> Result<(EF, MultilinearPoint, EF, EF), ProofError> { @@ -164,7 +164,7 @@ pub fn verify_gkr_quotient>>( Ok((quotient, point, claims_num, claims_den)) } -fn verify_gkr_quotient_step>>( +fn verify_gkr_quotient_step( verifier_state: &mut impl FSVerifier, n_vars: usize, point: &MultilinearPoint, @@ -207,14 +207,11 @@ mod tests { nums.iter().zip(den).map(|(&n, &d)| EF::from(n) / d).sum() } - fn bit_reverse_chunks_and_pack_ext>>(v: &[EF], chunk_log: usize) -> Vec> { + fn bit_reverse_chunks_and_pack_ext(v: &[EF], chunk_log: usize) -> Vec> { pack_extension(&bit_reverse_chunks(v, chunk_log)) } - fn bit_reverse_chunks_and_pack_base>>( - v: &[PF], - chunk_log: usize, - ) -> Vec> { + fn bit_reverse_chunks_and_pack_base(v: &[PF], chunk_log: usize) -> Vec> { let width: usize = packing_width::(); let mut res = unsafe { uninitialized_vec::>(v.len() / width) }; let unpacked = PFPacking::::unpack_slice_mut(&mut res); diff --git a/crates/sub_protocols/src/quotient_gkr/sumcheck_utils.rs b/crates/sub_protocols/src/quotient_gkr/sumcheck_utils.rs index c771d6e04..30adff7d8 100644 --- a/crates/sub_protocols/src/quotient_gkr/sumcheck_utils.rs +++ b/crates/sub_protocols/src/quotient_gkr/sumcheck_utils.rs @@ -85,7 +85,7 @@ fn within_pt(remaining_eq: &[EF], head_len: usize) -> Vec { .collect() } -fn finalize_round>>( +fn finalize_round( prover_state: &mut impl FSProver, coeffs: RoundCoeffs>, alpha: EF, @@ -107,7 +107,7 @@ fn finalize_round>>( } #[allow(clippy::too_many_arguments)] -pub(super) fn quotient_sumcheck_prove_packed_br_base>>( +pub(super) fn quotient_sumcheck_prove_packed_br_base( prover_state: &mut impl FSProver, packed_nums: &[PFPacking], packed_dens: &[EFPacking], @@ -181,7 +181,7 @@ pub(super) fn quotient_sumcheck_prove_packed_br_base>> /// bit-reversed by chunk + Packed #[allow(clippy::too_many_arguments)] -pub(super) fn run_phase1_sumcheck<'a, EF: ExtensionField>>( +pub(super) fn run_phase1_sumcheck<'a, EF: KoalaBearExtension>( prover_state: &mut impl FSProver, mut nums: ArenaCow<'a, EFPacking>, mut dens: ArenaCow<'a, EFPacking>, @@ -277,7 +277,7 @@ pub(super) fn run_phase1_sumcheck<'a, EF: ExtensionField>>( // Normal ordering (not bit-reversed) + not packed #[allow(clippy::too_many_arguments)] -pub(super) fn run_phase2_sumcheck>>( +pub(super) fn run_phase2_sumcheck( prover_state: &mut impl FSProver, mut num_l: ArenaVec, mut num_r: ArenaVec, @@ -372,7 +372,7 @@ pub(super) fn run_phase2_sumcheck>>( (q_natural, evals) } -fn fold_normal_with_padding>>(m: &[EF], r: EF, pad_value: EF) -> ArenaVec { +fn fold_normal_with_padding(m: &[EF], r: EF, pad_value: EF) -> ArenaVec { let active = m.len(); let new_active = active.div_ceil(2); assert!(new_active != 0); @@ -392,7 +392,7 @@ fn fold_normal_with_padding>>(m: &[EF], r: EF, pad_val out } -fn compute_round_packed>, N>( +fn compute_round_packed( nums: &[N], dens: &[EFPacking], layer_chunk_log: usize, @@ -437,7 +437,7 @@ where } #[allow(clippy::type_complexity)] -fn fold_and_compute_round_packed>, N>( +fn fold_and_compute_round_packed( nums: &[N], dens: &[EFPacking], layer_chunk_log_old: usize, @@ -510,7 +510,7 @@ where (new_nums, new_dens, coeffs) } -fn build_bare_from_coeffs>>( +fn build_bare_from_coeffs( c0_raw: EF, c2_raw: EF, eq_alpha: EF, diff --git a/crates/whir/src/commit.rs b/crates/whir/src/commit.rs index f74bcb7fe..da1f17649 100644 --- a/crates/whir/src/commit.rs +++ b/crates/whir/src/commit.rs @@ -1,7 +1,6 @@ // Credits: whir-p3 (https://github.com/tcoratger/whir-p3) (MIT and Apache-2.0 licenses). use fiat_shamir::FSProver; -use field::{ExtensionField, TwoAdicField}; use poly::*; use tracing::{info_span, instrument}; use zk_alloc::ArenaVec; @@ -9,12 +8,12 @@ use zk_alloc::ArenaVec; use crate::*; #[derive(Debug, Clone)] -pub enum MerkleData>> { +pub enum MerkleData { Base(RoundMerkleTree>), Extension(RoundMerkleTree>), } -impl>> MerkleData { +impl MerkleData { pub(crate) fn build( matrix: DftOutput, full_n_cols: usize, @@ -49,7 +48,7 @@ impl>> MerkleData { #[derive(Debug, Clone)] pub struct Witness where - EF: ExtensionField>, + EF: KoalaBearExtension, { pub prover_data: MerkleData, pub ood_points: Vec, @@ -58,8 +57,7 @@ where impl WhirConfig where - EF: ExtensionField>, - PF: TwoAdicField, + EF: KoalaBearExtension, { #[instrument(skip_all)] pub fn commit( diff --git a/crates/whir/src/config.rs b/crates/whir/src/config.rs index a2de41af2..418ce9344 100644 --- a/crates/whir/src/config.rs +++ b/crates/whir/src/config.rs @@ -134,6 +134,8 @@ pub struct WhirConfig { impl WhirConfig where EF: Field, + // Not implied here: this impl only requires `EF: Field`, so `PF` stays an opaque + // projection rather than normalizing to `KoalaBear`. PF: TwoAdicField, { /// `log_c` controls the proximity parameter `η` (η = √ρ/c for JB, η = ρ/c for CB). diff --git a/crates/whir/src/open.rs b/crates/whir/src/open.rs index cfafa784a..a5b4604ba 100644 --- a/crates/whir/src/open.rs +++ b/crates/whir/src/open.rs @@ -3,7 +3,7 @@ use ::utils::log2_strict_usize; use fiat_shamir::{FSProver, MerklePath, ProofResult}; use field::PrimeCharacteristicRing; -use field::{ExtensionField, Field, TwoAdicField}; +use field::{Field, TwoAdicField}; use sumcheck::{ ProductComputation, packing_unpack_sum, run_product_sumcheck_from_round1, run_product_sumcheck_from_round1_delayed, sumcheck_prove_many_rounds, @@ -15,8 +15,7 @@ use crate::{config::WhirConfig, *}; impl WhirConfig where - EF: ExtensionField>, - PF: TwoAdicField, + EF: KoalaBearExtension, { fn validate_parameters(&self) -> bool { self.num_variables == self.folding_factor.total_number(self.n_rounds()) + self.final_sumcheck_rounds @@ -279,7 +278,7 @@ where } } -fn open_merkle_tree_at_challenges>>( +fn open_merkle_tree_at_challenges( merkle_tree: &MerkleData, prover_state: &mut impl FSProver, stir_challenges_indexes: &[usize], @@ -322,7 +321,7 @@ fn open_merkle_tree_at_challenges>>( } #[derive(Debug, Clone)] -pub struct SumcheckSingle>> { +pub struct SumcheckSingle { /// Evaluations of the polynomial `p(X)`. pub(crate) evals: MleOwned, /// Evaluations of the equality polynomial used for enforcing constraints. @@ -333,7 +332,7 @@ pub struct SumcheckSingle>> { impl SumcheckSingle where - EF: ExtensionField>, + EF: KoalaBearExtension, { #[instrument(skip_all)] pub(crate) fn add_new_equality( @@ -463,7 +462,7 @@ where #[derive(Debug)] pub(crate) struct RoundState where - EF: ExtensionField>, + EF: KoalaBearExtension, { domain_size: usize, next_domain_gen: PF, @@ -476,8 +475,7 @@ where #[allow(clippy::mismatching_type_param_order)] impl RoundState where - EF: ExtensionField>, - PF: TwoAdicField, + EF: KoalaBearExtension, { pub(crate) fn initialize_first_round_state( prover: &WhirConfig, @@ -534,21 +532,21 @@ where const LAZY_OVERLAY_SPAN_MAX: usize = 8; // packed words; small blocks are pre-expanded -struct LazyFullTerm>> { +struct LazyFullTerm { left: ArenaVec, // prefix eq-table, scalar folded in right: ArenaVec>, // packed suffix eq-table rshift: usize, // hi = j >> rshift, lo = j & ((1 << rshift) - 1) } /// scalar·eq(point,·) on the packed range [start, start + 2^ishift). -struct LazyBlock>> { +struct LazyBlock { start: usize, ishift: usize, inner_id: u32, scalar: EF, } -pub(crate) struct LazyCombineTerms>> { +pub(crate) struct LazyCombineTerms { full: Vec>, inners: Vec>>, grid_blocks: Vec>, @@ -558,7 +556,7 @@ pub(crate) struct LazyCombineTerms>> { pub(crate) combined_sum: EF, } -impl>> LazyCombineTerms { +impl LazyCombineTerms { #[inline(always)] fn value_at(&self, j: usize) -> EFPacking { let mut acc = EFPacking::::ZERO; @@ -579,7 +577,7 @@ impl>> LazyCombineTerms { /// Builds the lazy term tables; `combined_sum` is the exact combined value. pub(crate) fn build_lazy_combine_terms(statements: &[SparseStatement], gamma: EF) -> LazyCombineTerms where - EF: ExtensionField>, + EF: KoalaBearExtension, { let num_variables = statements[0].total_num_variables; assert!(statements.iter().all(|e| e.total_num_variables == num_variables)); @@ -709,7 +707,7 @@ const MAX_RUN_TERMS: usize = 12; /// (packed slice, scalar) contributing `slice[t] * scalar` to the weight at offset t. /// Returns the term count, or None if more than `MAX_RUN_TERMS` cover the run. #[inline] -fn gather_run_terms<'a, EF: ExtensionField>>( +fn gather_run_terms<'a, EF: KoalaBearExtension>( terms: &'a LazyCombineTerms, base: usize, run: usize, @@ -744,7 +742,7 @@ fn combine_and_compute_first_round( terms: &LazyCombineTerms, ) -> (DensePolynomial, ArenaVec>) where - EF: ExtensionField>, + EF: KoalaBearExtension, EFPacking: std::ops::Mul, Output = EFPacking> + std::ops::Mul>, { let n = evals.len(); diff --git a/crates/whir/src/utils.rs b/crates/whir/src/utils.rs index 9ef86a8b8..8a196c1bb 100644 --- a/crates/whir/src/utils.rs +++ b/crates/whir/src/utils.rs @@ -5,7 +5,7 @@ use field::BasedVectorSpace; use field::Field; use field::PackedValue; use field::PrimeCharacteristicRing; -use field::{ExtensionField, TwoAdicField}; +use field::TwoAdicField; use poly::*; use std::any::{Any, TypeId}; use std::collections::HashMap; @@ -69,7 +69,7 @@ pub(crate) fn get_challenge_stir_queries>( /// A utility function to sample Out-of-Domain (OOD) points and evaluate them. /// /// This should be used on the prover side. -pub(crate) fn sample_ood_points>, E>( +pub(crate) fn sample_ood_points( prover_state: &mut impl FSProver, num_samples: usize, num_variables: usize, @@ -108,14 +108,13 @@ pub(crate) enum DftOutput { Extension(Matrix>), } -pub(crate) fn reorder_and_dft>>( +pub(crate) fn reorder_and_dft( evals: &MleRef<'_, EF>, folding_factor: usize, log_inv_rate: usize, dft_n_cols: usize, ) -> DftOutput where - PF: TwoAdicField, { let prepared_evals = prepare_evals_for_fft(evals, folding_factor, log_inv_rate, dft_n_cols); let dft = global_dft::>(); @@ -131,7 +130,7 @@ where } } -fn prepare_evals_for_fft>>( +fn prepare_evals_for_fft( evals: &MleRef<'_, EF>, folding_factor: usize, log_inv_rate: usize, @@ -202,7 +201,7 @@ fn prepare_evals_for_fft_unpacked( out } -fn prepare_evals_for_fft_packed_extension>>( +fn prepare_evals_for_fft_packed_extension( evals: &[EFPacking], folding_factor: usize, log_inv_rate: usize, diff --git a/crates/whir/src/verify.rs b/crates/whir/src/verify.rs index 6d6e99318..d25891a80 100644 --- a/crates/whir/src/verify.rs +++ b/crates/whir/src/verify.rs @@ -5,6 +5,7 @@ use std::{fmt::Debug, marker::PhantomData}; use ::utils::log2_strict_usize; use fiat_shamir::{FSVerifier, ProofError, ProofResult, try_pack_scalars_to_extension}; use field::{ExtensionField, Field, PrimeCharacteristicRing, TwoAdicField}; +use koala_bear::KoalaBear; use crate::*; @@ -26,7 +27,7 @@ impl> ParsedCommitment { where F: TwoAdicField, EF: ExtensionField + TwoAdicField, - EF: ExtensionField>, + EF: KoalaBearExtension, { let root = verifier_state.next_base_scalars_vec(DIGEST_ELEMS)?.try_into().unwrap(); let mut ood_points = vec![]; @@ -61,7 +62,7 @@ impl> ParsedCommitment { impl WhirConfig where - EF: TwoAdicField + ExtensionField>, + EF: TwoAdicField + KoalaBearExtension, { pub fn parse_commitment( &self, @@ -76,8 +77,7 @@ where impl WhirConfig where - EF: TwoAdicField + ExtensionField>, - PF: TwoAdicField, + EF: TwoAdicField + KoalaBearExtension, { #[allow(clippy::too_many_lines)] pub fn verify( @@ -87,7 +87,7 @@ where statement: Vec>, ) -> ProofResult> where - F: TwoAdicField + ExtensionField>, + F: TwoAdicField + ExtensionField, EF: ExtensionField, { statement @@ -235,7 +235,7 @@ where round_index: usize, ) -> ProofResult>> where - F: Field + ExtensionField>, + F: Field + ExtensionField, EF: ExtensionField, { let leafs_base_field = round_index == 0; @@ -298,7 +298,7 @@ where _var_shift: usize, ) -> ProofResult>> where - F: Field + ExtensionField>, + F: Field + ExtensionField, EF: ExtensionField, { let merkle_height = log2_strict_usize(dimensions[0].height); @@ -422,7 +422,7 @@ pub(crate) fn verify_sumcheck_rounds( ) -> ProofResult> where F: TwoAdicField, - EF: ExtensionField + TwoAdicField + ExtensionField>, + EF: ExtensionField + TwoAdicField + KoalaBearExtension, { // Preallocate vector to hold the randomness values let mut randomness = Vec::with_capacity(rounds); From f848d14538d85d36b2fda31f27a46398c37e82c8 Mon Sep 17 00:00:00 2001 From: Tom Wambsgans Date: Sun, 26 Jul 2026 21:54:02 +0200 Subject: [PATCH 2/5] poseidon: drop the add_kb/mul_kb type-punning dispatch `AirBuilder::IF` now requires `Algebra` alongside `Algebra`: round constants are plain `KoalaBear` scalars while `Self::F` may be a SIMD packing, so `Algebra` gave `+= PackedKoalaBear` but not `+= KoalaBear`. That gap is the entire reason the Poseidon AIR reached for `TypeId` + pointer casts. With the bound in place `add_kb(s, c)` becomes `*s += c` and `mul_kb(x, c)` becomes `x * c`, and both helpers are deleted along with `mds_air_16`'s five-way numeric dispatch -- 55 lines and 11 unsafe pointer casts, replaced by ordinary arithmetic. `mds_air_16` keeps its `SymbolicExpression` check, which is a real semantic fork (the dense matrix makes the zkDSL emit `dot_product_be` instead of Karatsuba), not type recovery. This is what `trace_gen.rs` already did with `F: Algebra`; the AIR side can now do the same. Pinning the prime subfield in the previous commit is what made this affordable. `PF` normalizes to `KoalaBear`, so the normal constraint folder satisfies `Algebra` for free and only one obligation propagates -- `EFPacking: Algebra`, since `PackedFieldExtension` supplies `Algebra` but not the unpacked scalar. Adding that to `PackedFieldExtension` directly is not possible: the quintic impl would then need a blanket `Algebra` that overlaps its existing `Algebra` when the packing degenerates. So the two `Algebra`/`From` impls for `PackedQuinticExtensionField` are declared per concrete packing, and the bound is carried explicitly through the packed sumcheck path. One call site needed disambiguating: with the new bound in scope the solver has a second `PackedFieldExtension` candidate and can no longer infer the trait's parameters for `to_ext_iter`. Verified: 66 tests pass (including `test_prove_poseidon` and `display_poseidon_air_in_zk_dsl`, which exercise both the numeric and symbolic paths), clippy -Dwarnings and rustfmt clean, no-SIMD/AVX2/AVX-512 all build, and Charon still extracts `sumcheck_verify` and `verify_execution`. Co-Authored-By: Claude Opus 5 (1M context) --- Cargo.lock | 1 + crates/backend/air/Cargo.toml | 1 + .../air/src/constraint_folder/packed.rs | 7 +- crates/backend/air/src/lib.rs | 3 +- crates/backend/air/src/symbolic.rs | 7 +- .../src/quintic_extension/packed_extension.rs | 9 ++ crates/backend/sumcheck/Cargo.toml | 1 + crates/backend/sumcheck/src/sc_computation.rs | 2 + crates/lean_vm/src/tables/poseidon/mod.rs | 83 ++++--------------- crates/sub_protocols/src/air_sumcheck.rs | 22 ++++- crates/whir/src/config.rs | 2 - 11 files changed, 57 insertions(+), 81 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9f6f8c5f6..fc55cca11 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -16,6 +16,7 @@ name = "air" version = "0.1.0" dependencies = [ "field", + "koala-bear", "poly", ] diff --git a/crates/backend/air/Cargo.toml b/crates/backend/air/Cargo.toml index 931fb611e..9387236b4 100644 --- a/crates/backend/air/Cargo.toml +++ b/crates/backend/air/Cargo.toml @@ -5,4 +5,5 @@ edition.workspace = true [dependencies] field = { path = "../field", package = "field" } +koala-bear = { path = "../koala-bear", package = "koala-bear" } poly = { path = "../poly", package = "poly" } diff --git a/crates/backend/air/src/constraint_folder/packed.rs b/crates/backend/air/src/constraint_folder/packed.rs index 0e6729675..76824662f 100644 --- a/crates/backend/air/src/constraint_folder/packed.rs +++ b/crates/backend/air/src/constraint_folder/packed.rs @@ -1,11 +1,8 @@ use crate::*; use field::*; +use koala_bear::KoalaBear; use poly::*; -// Keeps `ExtensionField>` where the rest of the stack uses `KoalaBearExtension`: the -// `low_degree_block` override below names `&mut Self` in a higher-ranked closure bound, and there -// the solver won't see through a blanket-implemented alias, so every bound on the impl stops -// applying. Callers using `KoalaBearExtension` satisfy this weaker bound anyway. #[derive(Debug)] pub struct ConstraintFolderPacked<'a, IF, EF: ExtensionField>, ExtraData: AlphaPowers> { pub flat: &'a [IF], @@ -42,7 +39,7 @@ where impl<'a, IF, EF, ExtraData> AirBuilder for ConstraintFolderPacked<'a, IF, EF, ExtraData> where - IF: Algebra> + 'static, + IF: Algebra> + Algebra + 'static, EF: Field + ExtensionField>, EFPacking: PrimeCharacteristicRing + Mul> + Add>, ExtraData: AlphaPowers, diff --git a/crates/backend/air/src/lib.rs b/crates/backend/air/src/lib.rs index deee0bf0f..d48413fca 100644 --- a/crates/backend/air/src/lib.rs +++ b/crates/backend/air/src/lib.rs @@ -2,6 +2,7 @@ use core::ops::{Add, Mul, Sub}; use field::{Algebra, PrimeCharacteristicRing}; +use koala_bear::KoalaBear; mod symbolic; pub use symbolic::*; @@ -36,7 +37,7 @@ pub trait AirBuilder: Sized { type F: PrimeCharacteristicRing + 'static; /// Intermediate field: equals F in base-field rounds, EF in extension rounds /// (or their respective SIMD packings). - type IF: Algebra + 'static; + type IF: Algebra + Algebra + 'static; /// Always the extension field (or its SIMD packing). type EF: PrimeCharacteristicRing + 'static diff --git a/crates/backend/air/src/symbolic.rs b/crates/backend/air/src/symbolic.rs index cbc5283a3..bff773476 100644 --- a/crates/backend/air/src/symbolic.rs +++ b/crates/backend/air/src/symbolic.rs @@ -7,6 +7,7 @@ use core::marker::PhantomData; use core::ops::{Add, AddAssign, Deref, Mul, MulAssign, Neg, Sub, SubAssign}; use field::{Algebra, Field, InjectiveMonomial, PrimeCharacteristicRing}; +use koala_bear::KoalaBear; use crate::{Air, AirBuilder}; @@ -284,7 +285,10 @@ impl SymbolicAirBuilder { } } -impl AirBuilder for SymbolicAirBuilder { +impl AirBuilder for SymbolicAirBuilder +where + SymbolicExpression: Algebra, +{ type F = F; type IF = SymbolicExpression; type EF = SymbolicExpression; @@ -325,6 +329,7 @@ pub type SymbolicAirData = ( pub fn get_symbolic_constraints_and_bus_data_values(air: &A) -> SymbolicAirData where A::ExtraData: Default, + SymbolicExpression: Algebra, { let mut builder = SymbolicAirBuilder::::new(air.n_columns(), air.n_shift_columns()); air.eval(&mut builder, &Default::default()); diff --git a/crates/backend/koala-bear/src/quintic_extension/packed_extension.rs b/crates/backend/koala-bear/src/quintic_extension/packed_extension.rs index 34e374a68..352144c19 100644 --- a/crates/backend/koala-bear/src/quintic_extension/packed_extension.rs +++ b/crates/backend/koala-bear/src/quintic_extension/packed_extension.rs @@ -132,6 +132,15 @@ macro_rules! impl_packed_ext_scalar_ops { } } } + + impl From for PackedQuinticExtensionField { + #[inline] + fn from(x: KoalaBear) -> Self { + Self::from(<$pf>::from(x)) + } + } + + impl Algebra for PackedQuinticExtensionField {} }; } diff --git a/crates/backend/sumcheck/Cargo.toml b/crates/backend/sumcheck/Cargo.toml index 47e7a2c7e..ae9af9cd9 100644 --- a/crates/backend/sumcheck/Cargo.toml +++ b/crates/backend/sumcheck/Cargo.toml @@ -6,6 +6,7 @@ edition.workspace = true [dependencies] field = { path = "../field", package = "field" } air = { path = "../air", package = "air" } +koala-bear = { path = "../koala-bear", package = "koala-bear" } poly = { path = "../poly", package = "poly" } zk-alloc.workspace = true fiat-shamir = { path = "../fiat-shamir", package = "fiat-shamir" } diff --git a/crates/backend/sumcheck/src/sc_computation.rs b/crates/backend/sumcheck/src/sc_computation.rs index 0e2f85732..a72a19284 100644 --- a/crates/backend/sumcheck/src/sc_computation.rs +++ b/crates/backend/sumcheck/src/sc_computation.rs @@ -1,6 +1,7 @@ use crate::*; use air::*; use field::*; +use koala_bear::KoalaBear; use poly::*; use std::any::TypeId; use std::ops::{Add, AddAssign, Mul, MulAssign, Sub}; @@ -35,6 +36,7 @@ macro_rules! impl_air_eval { impl SumcheckComputation for A where EF: KoalaBearExtension, + EFPacking: Algebra, A: Send + Sync + Air, A::ExtraData: AlphaPowers, { diff --git a/crates/lean_vm/src/tables/poseidon/mod.rs b/crates/lean_vm/src/tables/poseidon/mod.rs index c0bf33075..de82860d9 100644 --- a/crates/lean_vm/src/tables/poseidon/mod.rs +++ b/crates/lean_vm/src/tables/poseidon/mod.rs @@ -4,28 +4,15 @@ use crate::*; use crate::{execution::memory::MemoryAccess, tables::poseidon::trace_gen::generate_trace_rows_for_perm}; use backend::*; -/// Dispatch `mds_fft_16` through concrete types. /// For `SymbolicExpression` we use the dense form so the zkDSL generator can /// emit `dot_product_be` precompile calls instead of Karatsuba arithmetic. #[inline(always)] -fn mds_air_16(state: &mut [A; WIDTH]) { +fn mds_air_16 + 'static>(state: &mut [A; WIDTH]) { if TypeId::of::() == TypeId::of::>() { dense_mat_vec_air_16(mds_dense_16(), state); - return; + } else { + mds_fft_16(state); } - macro_rules! dispatch { - ($t:ty) => { - if TypeId::of::() == TypeId::of::<$t>() { - mds_fft_16::<$t>(unsafe { &mut *(state as *mut [A; WIDTH] as *mut [$t; WIDTH]) }); - return; - } - }; - } - dispatch!(F); - dispatch!(EF); - dispatch!(FPacking); - dispatch!(EFPacking); - unreachable!() } fn mds_dense_16() -> &'static [[F; 16]; 16] { @@ -42,44 +29,6 @@ fn mds_dense_16() -> &'static [[F; 16]; 16] { }) } -/// Add a `KoalaBear` constant to any AIR type. -#[inline(always)] -fn add_kb(a: &mut A, value: F) { - macro_rules! dispatch { - ($t:ty) => { - if TypeId::of::() == TypeId::of::<$t>() { - *unsafe { &mut *(a as *mut A as *mut $t) } += value; - return; - } - }; - } - dispatch!(F); - dispatch!(EF); - dispatch!(FPacking); - dispatch!(EFPacking); - dispatch!(SymbolicExpression); - unreachable!() -} - -/// Multiply any AIR type by a `KoalaBear` constant. -#[inline(always)] -fn mul_kb(a: A, value: F) -> A { - macro_rules! dispatch { - ($t:ty) => { - if TypeId::of::() == TypeId::of::<$t>() { - let r = unsafe { std::ptr::read(&a as *const A as *const $t) } * value; - return unsafe { std::ptr::read(&r as *const $t as *const A) }; - } - }; - } - dispatch!(F); - dispatch!(EF); - dispatch!(FPacking); - dispatch!(EFPacking); - dispatch!(SymbolicExpression); - unreachable!() -} - mod trace_gen; pub use trace_gen::fill_trace_poseidon_16; @@ -405,7 +354,7 @@ fn eval_poseidon1_16(builder: &mut AB, local: &Poseidon1Cols16(builder: &mut AB, local: &Poseidon1Cols16( builder: &mut AB, ) { for (s, r) in state.iter_mut().zip(round_constants_1.iter()) { - add_kb(s, *r); + *s += *r; *s = s.cube(); } mds_air_16(state); for (s, r) in state.iter_mut().zip(round_constants_2.iter()) { - add_kb(s, *r); + *s += *r; *s = s.cube(); } mds_air_16(state); @@ -499,12 +448,12 @@ fn eval_last_2_full_rounds_16( builder: &mut AB, ) { for (s, r) in state.iter_mut().zip(round_constants_1.iter()) { - add_kb(s, *r); + *s += *r; *s = s.cube(); } mds_air_16(state); for (s, r) in state.iter_mut().zip(round_constants_2.iter()) { - add_kb(s, *r); + *s += *r; *s = s.cube(); } mds_air_16(state); @@ -523,30 +472,26 @@ fn eval_last_2_full_rounds_16( } #[inline] -fn dense_mat_vec_air_16(mat: &[[F; 16]; 16], state: &mut [A; WIDTH]) { +fn dense_mat_vec_air_16>(mat: &[[F; 16]; 16], state: &mut [A; WIDTH]) { let input = *state; for i in 0..WIDTH { let mut acc = A::ZERO; for j in 0..WIDTH { - acc += mul_kb(input[j], mat[i][j]); + acc += input[j] * mat[i][j]; } state[i] = acc; } } #[inline] -fn sparse_mat_air_16( - state: &mut [A; WIDTH], - first_row: &[F; WIDTH], - v: &[F; WIDTH], -) { +fn sparse_mat_air_16>(state: &mut [A; WIDTH], first_row: &[F; WIDTH], v: &[F; WIDTH]) { let old_s0 = state[0]; let mut new_s0 = A::ZERO; for j in 0..WIDTH { - new_s0 += mul_kb(state[j], first_row[j]); + new_s0 += state[j] * first_row[j]; } state[0] = new_s0; for i in 1..WIDTH { - state[i] += mul_kb(old_s0, v[i - 1]); + state[i] += old_s0 * v[i - 1]; } } diff --git a/crates/sub_protocols/src/air_sumcheck.rs b/crates/sub_protocols/src/air_sumcheck.rs index f0bb38cdb..3263cf4a1 100644 --- a/crates/sub_protocols/src/air_sumcheck.rs +++ b/crates/sub_protocols/src/air_sumcheck.rs @@ -45,6 +45,7 @@ pub trait OuterSumcheckSession: Debug { pub struct AirSumcheckSession<'a, EF: KoalaBearExtension, A: Air> where A::ExtraData: AlphaPowers, + EFPacking: Algebra, { multilinears: MleGroup<'a, EF>, eq_factor: Vec, // The last element is removed at each round @@ -63,6 +64,7 @@ where impl<'a, EF: KoalaBearExtension, A: Air> AirSumcheckSession<'a, EF, A> where A::ExtraData: AlphaPowers + AlphaPowersMut, + EFPacking: Algebra, { pub fn new( packed_multilinears: MleGroup<'a, EF>, @@ -126,6 +128,7 @@ where impl<'a, EF, A> AirSumcheckSession<'a, EF, A> where EF: KoalaBearExtension, + EFPacking: Algebra, A: Air + 'static, A::ExtraData: AlphaPowers, { @@ -201,6 +204,7 @@ where impl<'a, EF, A> OuterSumcheckSession for AirSumcheckSession<'a, EF, A> where EF: KoalaBearExtension, + EFPacking: Algebra, A: Air + Debug + 'static, A::ExtraData: AlphaPowers + AlphaPowersMut + Debug, { @@ -316,10 +320,12 @@ fn compute_raw_poly<'a, EF, A>( ) -> Vec where EF: KoalaBearExtension, + EFPacking: Algebra, A: Air + 'static, A::ExtraData: AlphaPowers, { - let unpack_sum_packed = |s: EFPacking| -> EF { EFPacking::::to_ext_iter([s]).sum::() }; + let unpack_sum_packed = + |s: EFPacking| -> EF { as PackedFieldExtension, EF>>::to_ext_iter([s]).sum::() }; if let Some((low_degree, low_n_constraints)) = computation.low_degree_air() { match multilinears { @@ -411,9 +417,18 @@ fn compute_raw_poly_degree_split( ) -> Vec where EF: KoalaBearExtension, + EFPacking: Algebra, A: Air + 'static, A::ExtraData: AlphaPowers, - IF: Algebra> + Copy + Send + Sync + Sub + AddAssign + PrimeCharacteristicRing + 'static, + IF: Algebra> + + Algebra + + Copy + + Send + + Sync + + Sub + + AddAssign + + PrimeCharacteristicRing + + 'static, EFPacking: PrimeCharacteristicRing + Mul> + Add> @@ -559,9 +574,10 @@ fn compute_raw_poly_impl( ) -> Vec where EF: KoalaBearExtension, + EFPacking: Algebra, A: Air + 'static, A::ExtraData: AlphaPowers, - IF: Copy + Send + Sync + Sub + AddAssign + PrimeCharacteristicRing, + IF: Copy + Send + Sync + Sub + AddAssign + PrimeCharacteristicRing + Algebra, EFT: Copy + Send + Sync + Add + AddAssign + Mul + PrimeCharacteristicRing, GetEq: Fn(usize) -> EFT + Sync + Send, UnpackSum: Fn(EFT) -> EF + Sync + Send, diff --git a/crates/whir/src/config.rs b/crates/whir/src/config.rs index 418ce9344..a2de41af2 100644 --- a/crates/whir/src/config.rs +++ b/crates/whir/src/config.rs @@ -134,8 +134,6 @@ pub struct WhirConfig { impl WhirConfig where EF: Field, - // Not implied here: this impl only requires `EF: Field`, so `PF` stays an opaque - // projection rather than normalizing to `KoalaBear`. PF: TwoAdicField, { /// `log_c` controls the proximity parameter `η` (η = √ρ/c for JB, η = ρ/c for CB). From f81b0d600f85a832e9c385ae2c5316ca5e0fab69 Mon Sep 17 00:00:00 2001 From: Tom Wambsgans Date: Sun, 26 Jul 2026 22:14:20 +0200 Subject: [PATCH 3/5] whir: merge the merkle TypeId arms into one generic body `merkle_commit`, `merkle_open` and `merkle_verify` were generic over `>` but each dispatched at runtime on `TypeId::of::<(F, EF)>()` against `(KoalaBear, QuinticExtensionFieldKB)` and `(KoalaBear, KoalaBear)`, transmuting into the concrete types and falling into `unimplemented!()` otherwise. `F` was only ever `PF`, which now normalizes to `KoalaBear`, so it stops being a parameter. That removes the reason for the dispatch: the two arms were the same code at `BasedVectorSpace::DIMENSION` 5 and 1, and the flattening helpers (`flatten_to_base_arena`, `reconstitute_from_base`, `flatten_to_base`) are already generic over the degree. Each function collapses to a single body generic in `EF`, with the degree-1 case falling out of the reflexive `ExtensionField for F` impl. merkle.rs goes from 224 to 161 lines and from 12 `unsafe` blocks to one (`ArenaVec::uninitialized`, unrelated). Workspace `TypeId::of` sites: 13 -> 7. Verified both degrees are actually exercised, not just compiled: instrumenting `test_run_whir` shows commit/open/verify each running at dim=1 (round 0, base field leaves) and dim=5 (later rounds), and the prove/verify roundtrip passes -- a width or flattening mistake in either path would break the Merkle root. 66 tests pass, clippy -Dwarnings and rustfmt clean, no-SIMD/AVX2/AVX-512 build, and `whir::verify` still extracts through Charon. Co-Authored-By: Claude Opus 5 (1M context) --- crates/whir/src/commit.rs | 8 +-- crates/whir/src/merkle.rs | 114 +++++++++----------------------------- crates/whir/src/verify.rs | 4 +- 3 files changed, 32 insertions(+), 94 deletions(-) diff --git a/crates/whir/src/commit.rs b/crates/whir/src/commit.rs index da1f17649..34b9283b1 100644 --- a/crates/whir/src/commit.rs +++ b/crates/whir/src/commit.rs @@ -21,11 +21,11 @@ impl MerkleData { ) -> (Self, [PF; DIGEST_ELEMS]) { match matrix { DftOutput::Base(m) => { - let (root, prover_data) = merkle_commit::, PF>(m, full_n_cols, effective_n_cols); + let (root, prover_data) = merkle_commit::>(m, full_n_cols, effective_n_cols); (MerkleData::Base(prover_data), root) } DftOutput::Extension(m) => { - let (root, prover_data) = merkle_commit::, EF>(m, full_n_cols, effective_n_cols); + let (root, prover_data) = merkle_commit::(m, full_n_cols, effective_n_cols); (MerkleData::Extension(prover_data), root) } } @@ -34,11 +34,11 @@ impl MerkleData { pub(crate) fn open(&self, index: usize) -> (MleOwned, Vec<[PF; DIGEST_ELEMS]>) { match self { MerkleData::Base(prover_data) => { - let (leaf, proof) = merkle_open::, PF>(prover_data, index); + let (leaf, proof) = merkle_open::>(prover_data, index); (MleOwned::Base(ArenaVec::from_slice(&leaf)), proof) } MerkleData::Extension(prover_data) => { - let (leaf, proof) = merkle_open::, EF>(prover_data, index); + let (leaf, proof) = merkle_open::(prover_data, index); (MleOwned::Extension(ArenaVec::from_slice(&leaf)), proof) } } diff --git a/crates/whir/src/merkle.rs b/crates/whir/src/merkle.rs index 72010f7e8..8eda5c5cf 100644 --- a/crates/whir/src/merkle.rs +++ b/crates/whir/src/merkle.rs @@ -2,14 +2,11 @@ // - whir-p3 (https://github.com/tcoratger/whir-p3) (MIT and Apache-2.0 licenses). // - Plonky3 (https://github.com/Plonky3/Plonky3) (MIT and Apache-2.0 licenses). -use std::any::TypeId; - use field::BasedVectorSpace; use field::ExtensionField; -use field::Field; use field::PackedValue; use field::PrimeCharacteristicRing; -use koala_bear::{KoalaBear, QuinticExtensionFieldKB, default_koalabear_poseidon1_16}; +use koala_bear::{KoalaBear, default_koalabear_poseidon1_16}; use poly::*; use symetric::merkle::unpack_array; @@ -24,37 +21,16 @@ pub use symetric::DIGEST_ELEMS; pub(crate) type RoundMerkleTree = WhirMerkleTree; -#[allow(clippy::missing_transmute_annotations)] -pub(crate) fn merkle_commit>( +pub(crate) fn merkle_commit>( matrix: Matrix>, full_n_cols: usize, effective_n_cols: usize, -) -> ([F; DIGEST_ELEMS], RoundMerkleTree) { - if TypeId::of::<(F, EF)>() == TypeId::of::<(KoalaBear, QuinticExtensionFieldKB)>() { - let matrix = unsafe { - std::mem::transmute::<_, Matrix>>(matrix) - }; - let dim = >::DIMENSION; - let dft_base_width = matrix.width * dim; - let full_base_width = full_n_cols * dim; - let effective_base_width = effective_n_cols * dim; - let base_values = flatten_to_base_arena::(matrix.values); - let base_matrix = Matrix::new(base_values, dft_base_width); - let tree = build_merkle_tree_koalabear(base_matrix, full_base_width, effective_base_width); - let root: [_; DIGEST_ELEMS] = tree.root(); - let root = unsafe { std::mem::transmute_copy::<_, [F; DIGEST_ELEMS]>(&root) }; - let tree = unsafe { std::mem::transmute::<_, RoundMerkleTree>(tree) }; - (root, tree) - } else if TypeId::of::<(F, EF)>() == TypeId::of::<(KoalaBear, KoalaBear)>() { - let matrix = unsafe { std::mem::transmute::<_, Matrix>>(matrix) }; - let tree = build_merkle_tree_koalabear(matrix, full_n_cols, effective_n_cols); - let root: [_; DIGEST_ELEMS] = tree.root(); - let root = unsafe { std::mem::transmute_copy::<_, [F; DIGEST_ELEMS]>(&root) }; - let tree = unsafe { std::mem::transmute::<_, RoundMerkleTree>(tree) }; - (root, tree) - } else { - unimplemented!() - } +) -> ([KoalaBear; DIGEST_ELEMS], RoundMerkleTree) { + let dim = >::DIMENSION; + let base_values = flatten_to_base_arena::(matrix.values); + let base_matrix = Matrix::new(base_values, matrix.width * dim); + let tree = build_merkle_tree_koalabear(base_matrix, full_n_cols * dim, effective_n_cols * dim); + (tree.root(), tree) } #[instrument(name = "build merkle tree", skip_all)] @@ -89,71 +65,33 @@ fn build_merkle_tree_koalabear( } } -#[allow(clippy::missing_transmute_annotations)] -pub(crate) fn merkle_open>( - merkle_tree: &RoundMerkleTree, +pub(crate) fn merkle_open>( + merkle_tree: &RoundMerkleTree, index: usize, -) -> (Vec, Vec<[F; DIGEST_ELEMS]>) { - if TypeId::of::<(F, EF)>() == TypeId::of::<(KoalaBear, QuinticExtensionFieldKB)>() { - let merkle_tree = unsafe { std::mem::transmute::<_, &RoundMerkleTree>(merkle_tree) }; - let (inner_leaf, proof) = merkle_tree.open(index); - let leaf = QuinticExtensionFieldKB::reconstitute_from_base(inner_leaf); - let leaf = unsafe { std::mem::transmute::<_, Vec>(leaf) }; - let proof = unsafe { std::mem::transmute::<_, Vec<[F; DIGEST_ELEMS]>>(proof) }; - (leaf, proof) - } else if TypeId::of::<(F, EF)>() == TypeId::of::<(KoalaBear, KoalaBear)>() { - let merkle_tree = unsafe { std::mem::transmute::<_, &RoundMerkleTree>(merkle_tree) }; - let (inner_leaf, proof) = merkle_tree.open(index); - let leaf = KoalaBear::reconstitute_from_base(inner_leaf); - let leaf = unsafe { std::mem::transmute::<_, Vec>(leaf) }; - let proof = unsafe { std::mem::transmute::<_, Vec<[F; DIGEST_ELEMS]>>(proof) }; - (leaf, proof) - } else { - unimplemented!() - } +) -> (Vec, Vec<[KoalaBear; DIGEST_ELEMS]>) { + let (inner_leaf, proof) = merkle_tree.open(index); + (EF::reconstitute_from_base(inner_leaf), proof) } -#[allow(clippy::missing_transmute_annotations)] -pub(crate) fn merkle_verify>( - merkle_root: [F; DIGEST_ELEMS], +pub(crate) fn merkle_verify>( + merkle_root: [KoalaBear; DIGEST_ELEMS], index: usize, dimension: Dimensions, data: Vec, - proof: &Vec<[F; DIGEST_ELEMS]>, + proof: &Vec<[KoalaBear; DIGEST_ELEMS]>, ) -> bool { let perm = default_koalabear_poseidon1_16(); let log_max_height = utils::log2_strict_usize(dimension.height); - if TypeId::of::<(F, EF)>() == TypeId::of::<(KoalaBear, QuinticExtensionFieldKB)>() { - let merkle_root = unsafe { std::mem::transmute_copy::<_, [KoalaBear; DIGEST_ELEMS]>(&merkle_root) }; - let data = unsafe { std::mem::transmute::<_, Vec>(data) }; - let proof = unsafe { std::mem::transmute::<_, &Vec<[KoalaBear; DIGEST_ELEMS]>>(proof) }; - let base_data = QuinticExtensionFieldKB::flatten_to_base(data); - symetric::merkle::merkle_verify::<_, _, _, DIGEST_ELEMS, 16, 8>( - &perm, - &perm, - &merkle_root, - log_max_height, - index, - &base_data, - proof, - ) - } else if TypeId::of::<(F, EF)>() == TypeId::of::<(KoalaBear, KoalaBear)>() { - let merkle_root = unsafe { std::mem::transmute_copy::<_, [KoalaBear; DIGEST_ELEMS]>(&merkle_root) }; - let data = unsafe { std::mem::transmute::<_, Vec>(data) }; - let proof = unsafe { std::mem::transmute::<_, &Vec<[KoalaBear; DIGEST_ELEMS]>>(proof) }; - let base_data = KoalaBear::flatten_to_base(data); - symetric::merkle::merkle_verify::<_, _, _, DIGEST_ELEMS, 16, 8>( - &perm, - &perm, - &merkle_root, - log_max_height, - index, - &base_data, - proof, - ) - } else { - unimplemented!() - } + let base_data = EF::flatten_to_base(data); + symetric::merkle::merkle_verify::<_, _, _, DIGEST_ELEMS, 16, 8>( + &perm, + &perm, + &merkle_root, + log_max_height, + index, + &base_data, + proof, + ) } #[derive(Debug, Clone)] diff --git a/crates/whir/src/verify.rs b/crates/whir/src/verify.rs index d25891a80..40a318ac2 100644 --- a/crates/whir/src/verify.rs +++ b/crates/whir/src/verify.rs @@ -322,7 +322,7 @@ where } for (i, &index) in indices.iter().enumerate() { - if !merkle_verify::, F>(*root, index, dimensions[0], answers[i].clone(), &merkle_proofs[i]) { + if !merkle_verify::(*root, index, dimensions[0], answers[i].clone(), &merkle_proofs[i]) { return Err(ProofError::InvalidProof); } } @@ -344,7 +344,7 @@ where } for (i, &index) in indices.iter().enumerate() { - if !merkle_verify::, EF>(*root, index, dimensions[0], answers[i].clone(), &merkle_proofs[i]) { + if !merkle_verify::(*root, index, dimensions[0], answers[i].clone(), &merkle_proofs[i]) { return Err(ProofError::InvalidProof); } } From 443ccfef5d502640b83a462369e8ec0ffb0cfa19 Mon Sep 17 00:00:00 2001 From: Tom Wambsgans Date: Sun, 26 Jul 2026 22:52:02 +0200 Subject: [PATCH 4/5] field: unblock Aeneas by removing closures from trait declarations Aeneas rejects a trait declaration whose default method bodies contain closures: the closure types are children of the trait, but reference `Self`, so they land in the same dependency group and Aeneas reports "groups of mixed mutually recursive definitions". Two groups blocked the sumcheck verifier. `PackedFieldExtension::to_ext_iter` loses its default body and becomes a required method. `PackedQuinticExtensionField` already provided one; the reflexive `PackedFieldExtension for F::Packing` impl gets an explicit body, which for DIMENSION = 1 reduces to picking lane `i` and so is equivalent to the default it replaces. `RawDataSerializable` is deleted outright rather than rewritten. Every one of its methods and its `NUM_BYTES` constant has zero callers anywhere in the workspace -- it is inherited from Plonky3 and only reached extraction because `Field` listed it as a supertrait. That removes the trait, the two `impl_raw_serializable_primefield{32,64}` macros, and both impls: 307 lines of dead code, and with them the second recursive group. Verified against Aeneas (HEAD 3a8586f, built against Charon 527ea8e3): both recursive groups are gone and Aeneas now proceeds through 62 prepasses into global translation before hitting an unrelated blocker. 66 tests pass, clippy -Dwarnings and rustfmt clean, no-SIMD/AVX2/AVX-512 build. Co-Authored-By: Claude Opus 5 (1M context) --- crates/backend/field/src/field.rs | 103 +------------ crates/backend/field/src/integers.rs | 141 ------------------ .../backend/field/src/packed/packed_traits.rs | 16 +- .../koala-bear/src/monty_31/monty_31.rs | 7 +- .../src/quintic_extension/extension.rs | 60 +------- 5 files changed, 10 insertions(+), 317 deletions(-) diff --git a/crates/backend/field/src/field.rs b/crates/backend/field/src/field.rs index fd3c6ba0a..57403380b 100644 --- a/crates/backend/field/src/field.rs +++ b/crates/backend/field/src/field.rs @@ -6,11 +6,10 @@ use core::fmt::{Debug, Display}; use core::hash::Hash; use core::iter::{Product, Sum}; use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign}; -use core::{array, slice}; +use core::slice; use serde::Serialize; use serde::de::DeserializeOwned; -use utils::iter_array_chunks_padded; use crate::exponentiation::bits_u64; use crate::integers::{QuotientMap, from_integer_types}; @@ -623,112 +622,12 @@ pub trait Algebra: // Every ring is an algebra over itself. impl Algebra for R {} -/// A collection of methods designed to help hash field elements. -/// -/// Most fields will want to reimplement many/all of these methods as the default implementations -/// are slow and involve converting to/from byte representations. -pub trait RawDataSerializable: Sized { - /// The number of bytes which this field element occupies in memory. - /// Must be equal to the length of self.into_bytes(). - const NUM_BYTES: usize; - - /// Convert a field element into a collection of bytes. - #[must_use] - fn into_bytes(self) -> impl IntoIterator; - - /// Convert an iterator of field elements into an iterator of bytes. - #[must_use] - fn into_byte_stream(input: impl IntoIterator) -> impl IntoIterator { - input.into_iter().flat_map(|elem| elem.into_bytes()) - } - - /// Convert an iterator of field elements into an iterator of u32s. - /// - /// If `NUM_BYTES` does not divide `4`, multiple `F`s may be packed together to make a single `u32`. Furthermore, - /// if `NUM_BYTES * input.len()` does not divide `4`, the final `u32` will involve padding bytes which are set to `0`. - #[must_use] - fn into_u32_stream(input: impl IntoIterator) -> impl IntoIterator { - let bytes = Self::into_byte_stream(input); - iter_array_chunks_padded(bytes, 0).map(u32::from_le_bytes) - } - - /// Convert an iterator of field elements into an iterator of u64s. - /// - /// If `NUM_BYTES` does not divide `8`, multiple `F`s may be packed together to make a single `u64`. Furthermore, - /// if `NUM_BYTES * input.len()` does not divide `8`, the final `u64` will involve padding bytes which are set to `0`. - #[must_use] - fn into_u64_stream(input: impl IntoIterator) -> impl IntoIterator { - let bytes = Self::into_byte_stream(input); - iter_array_chunks_padded(bytes, 0).map(u64::from_le_bytes) - } - - /// Convert an iterator of field element arrays into an iterator of byte arrays. - /// - /// Converts an element `[F; N]` into the byte array `[[u8; N]; NUM_BYTES]`. This is - /// intended for use with vectorized hash functions which use vector operations - /// to compute several hashes in parallel. - #[must_use] - fn into_parallel_byte_streams( - input: impl IntoIterator, - ) -> impl IntoIterator { - input.into_iter().flat_map(|vector| { - let bytes = vector.map(|elem| elem.into_bytes().into_iter().collect::>()); - (0..Self::NUM_BYTES).map(move |i| array::from_fn(|j| bytes[j][i])) - }) - } - - /// Convert an iterator of field element arrays into an iterator of u32 arrays. - /// - /// Converts an element `[F; N]` into the u32 array `[[u32; N]; NUM_BYTES/4]`. This is - /// intended for use with vectorized hash functions which use vector operations - /// to compute several hashes in parallel. - /// - /// This function is guaranteed to be equivalent to starting with `Iterator<[F; N]>` performing a transpose - /// operation to get `[Iterator; N]`, calling `into_u32_stream` on each element to get `[Iterator; N]` and then - /// performing another transpose operation to get `Iterator<[u32; N]>`. - /// - /// If `NUM_BYTES` does not divide `4`, multiple `[F; N]`s may be packed together to make a single `[u32; N]`. Furthermore, - /// if `NUM_BYTES * input.len()` does not divide `4`, the final `[u32; N]` will involve padding bytes which are set to `0`. - #[must_use] - fn into_parallel_u32_streams( - input: impl IntoIterator, - ) -> impl IntoIterator { - let bytes = Self::into_parallel_byte_streams(input); - iter_array_chunks_padded(bytes, [0; N]).map(|byte_array: [[u8; N]; 4]| { - array::from_fn(|i| u32::from_le_bytes(array::from_fn(|j| byte_array[j][i]))) - }) - } - - /// Convert an iterator of field element arrays into an iterator of u64 arrays. - /// - /// Converts an element `[F; N]` into the u64 array `[[u64; N]; NUM_BYTES/8]`. This is - /// intended for use with vectorized hash functions which use vector operations - /// to compute several hashes in parallel. - /// - /// This function is guaranteed to be equivalent to starting with `Iterator<[F; N]>` performing a transpose - /// operation to get `[Iterator; N]`, calling `into_u64_stream` on each element to get `[Iterator; N]` and then - /// performing another transpose operation to get `Iterator<[u64; N]>`. - /// - /// If `NUM_BYTES` does not divide `8`, multiple `[F; N]`s may be packed together to make a single `[u64; N]`. Furthermore, - /// if `NUM_BYTES * input.len()` does not divide `8`, the final `[u64; N]` will involve padding bytes which are set to `0`. - #[must_use] - fn into_parallel_u64_streams( - input: impl IntoIterator, - ) -> impl IntoIterator { - let bytes = Self::into_parallel_byte_streams(input); - iter_array_chunks_padded(bytes, [0; N]).map(|byte_array: [[u8; N]; 8]| { - array::from_fn(|i| u64::from_le_bytes(array::from_fn(|j| byte_array[j][i]))) - }) - } -} - /// A field `F`. This permits both modular fields `ℤ/p` along with their field extensions. /// /// A ring is a field if every element `x` has a unique multiplicative inverse `x^{-1}` /// which satisfies `x * x^{-1} = F::ONE`. pub trait Field: Algebra - + RawDataSerializable + Packable + 'static + Copy diff --git a/crates/backend/field/src/integers.rs b/crates/backend/field/src/integers.rs index 4c326e11c..e13cf7c37 100644 --- a/crates/backend/field/src/integers.rs +++ b/crates/backend/field/src/integers.rs @@ -469,144 +469,3 @@ macro_rules! impl_u_i_size { impl_u_i_size!(usize, u8, u16, u32, u64, u128); impl_u_i_size!(isize, i8, i16, i32, i64, i128); - -/// A simple macro which allows us to implement the `RawSerializable` trait for any 32-bit field. -/// The field must implement PrimeField32. -/// -/// This macro doesn't need any inputs as the implementation is identical for all 32-bit fields. -#[macro_export] -macro_rules! impl_raw_serializable_primefield32 { - () => { - const NUM_BYTES: usize = 4; - - #[allow(refining_impl_trait)] - #[inline] - fn into_bytes(self) -> [u8; 4] { - self.to_unique_u32().to_le_bytes() - } - - #[inline] - fn into_u32_stream(input: impl IntoIterator) -> impl IntoIterator { - // As every element is 32 bits, we can just convert the input to a unique u32. - input.into_iter().map(|x| x.to_unique_u32()) - } - - #[inline] - fn into_u64_stream(input: impl IntoIterator) -> impl IntoIterator { - let mut input = input.into_iter(); - iter::from_fn(move || { - // If the first input.next() returns None, we return None. - let a = input.next()?; - // Otherwise we either pack 2 32 bit elements together if the iterator - // gives a second value or just cast the 32 bit element to 64 bits. - if let Some(b) = input.next() { - Some(a.to_unique_u64() | b.to_unique_u64() << 32) - } else { - Some(a.to_unique_u64()) - } - }) - } - - #[inline] - fn into_parallel_byte_streams( - input: impl IntoIterator, - ) -> impl IntoIterator { - input.into_iter().flat_map(|vector| { - let bytes = vector.map(|elem| elem.into_bytes()); - (0..Self::NUM_BYTES).map(move |i| array::from_fn(|j| bytes[j][i])) - }) - } - - #[inline] - fn into_parallel_u32_streams( - input: impl IntoIterator, - ) -> impl IntoIterator { - // As every element is 32 bits, we can just convert the input to a unique u32. - input.into_iter().map(|vec| vec.map(|x| x.to_unique_u32())) - } - - #[inline] - fn into_parallel_u64_streams( - input: impl IntoIterator, - ) -> impl IntoIterator { - let mut input = input.into_iter(); - iter::from_fn(move || { - // If the first input.next() returns None, we return None. - let a = input.next()?; - // Otherwise we either pack pairs of 32 bit elements together if the iterator - // gives two arrays of or just cast the 32 bit elements to 64 bits. - if let Some(b) = input.next() { - let ab = array::from_fn(|i| { - let ai = a[i].to_unique_u64(); - let bi = b[i].to_unique_u64(); - ai | (bi << 32) - }); - Some(ab) - } else { - Some(a.map(|x| x.to_unique_u64())) - } - }) - } - }; -} - -/// A simple macro which allows us to implement the `RawSerializable` trait for any 64-bit field. -/// The field must implement PrimeField64 (and should not implement PrimeField32). -/// -/// This macro doesn't need any inputs as the implementation is identical for all 64-bit fields. -#[macro_export] -macro_rules! impl_raw_serializable_primefield64 { - () => { - const NUM_BYTES: usize = 8; - - #[allow(refining_impl_trait)] - #[inline] - fn into_bytes(self) -> [u8; 8] { - self.to_unique_u64().to_le_bytes() - } - - #[inline] - fn into_u32_stream(input: impl IntoIterator) -> impl IntoIterator { - input.into_iter().flat_map(|x| { - let x_u64 = x.to_unique_u64(); - [x_u64 as u32, (x_u64 >> 32) as u32] - }) - } - - #[inline] - fn into_u64_stream(input: impl IntoIterator) -> impl IntoIterator { - // As every element is 64 bits, we can just convert the input to a unique u64. - input.into_iter().map(|x| x.to_unique_u64()) - } - - #[inline] - fn into_parallel_byte_streams( - input: impl IntoIterator, - ) -> impl IntoIterator { - input.into_iter().flat_map(|vector| { - let bytes = vector.map(|elem| elem.into_bytes()); - (0..Self::NUM_BYTES).map(move |i| array::from_fn(|j| bytes[j][i])) - }) - } - - #[inline] - fn into_parallel_u32_streams( - input: impl IntoIterator, - ) -> impl IntoIterator { - input.into_iter().flat_map(|vec| { - let vec_64 = vec.map(|x| x.to_unique_u64()); - let vec_32_lo = vec_64.map(|x| x as u32); - let vec_32_hi = vec_64.map(|x| (x >> 32) as u32); - [vec_32_lo, vec_32_hi] - }) - } - - #[inline] - fn into_parallel_u64_streams( - input: impl IntoIterator, - ) -> impl IntoIterator { - // As every element is 64 bits, we can just convert the input to a unique u64. - input.into_iter().map(|vec| vec.map(|x| x.to_unique_u64())) - } - }; -} diff --git a/crates/backend/field/src/packed/packed_traits.rs b/crates/backend/field/src/packed/packed_traits.rs index 503777308..cdd81fecd 100644 --- a/crates/backend/field/src/packed/packed_traits.rs +++ b/crates/backend/field/src/packed/packed_traits.rs @@ -339,16 +339,8 @@ pub trait PackedFieldExtension) -> impl Iterator { - iter.into_iter().flat_map(|x| { - (0..BaseField::Packing::WIDTH).map(move |i| { - let packed_coeffs = x.as_basis_coefficients_slice(); - ExtField::from_basis_coefficients_fn(|j| packed_coeffs[j].as_slice()[i]) - }) - }) - } + fn to_ext_iter(iter: impl IntoIterator) -> impl Iterator; /// Given a iterator of packed extension field elements, convert to an iterator of /// extension field elements. @@ -432,6 +424,12 @@ impl PackedFieldExtension for F::Packing { *F::Packing::from_slice(ext_slice) } + #[inline] + fn to_ext_iter(iter: impl IntoIterator) -> impl Iterator { + iter.into_iter() + .flat_map(|x| (0..F::Packing::WIDTH).map(move |i| x.as_slice()[i])) + } + #[inline] fn packed_ext_powers(base: F) -> Powers { F::Packing::packed_powers(base) diff --git a/crates/backend/koala-bear/src/monty_31/monty_31.rs b/crates/backend/koala-bear/src/monty_31/monty_31.rs index 23c561318..f95adfca3 100644 --- a/crates/backend/koala-bear/src/monty_31/monty_31.rs +++ b/crates/backend/koala-bear/src/monty_31/monty_31.rs @@ -9,13 +9,12 @@ use core::hash::Hash; use core::iter::{Product, Sum}; use core::marker::PhantomData; use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign}; -use core::{array, iter}; use field::integers::QuotientMap; use field::op_assign_macros::{impl_add_assign, impl_div_methods, impl_mul_methods, impl_sub_assign}; use field::{ Field, InjectiveMonomial, Packable, PermutationMonomial, PrimeCharacteristicRing, PrimeField, PrimeField32, - PrimeField64, RawDataSerializable, TwoAdicField, impl_raw_serializable_primefield32, quotient_map_small_int, + PrimeField64, TwoAdicField, quotient_map_small_int, }; use rand::Rng; use rand::distr::{Distribution, StandardUniform}; @@ -371,10 +370,6 @@ impl, const D: u64> PermutationMon } } -impl RawDataSerializable for MontyField31 { - impl_raw_serializable_primefield32!(); -} - impl Field for MontyField31 { #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] type Packing = crate::PackedMontyField31Neon; diff --git a/crates/backend/koala-bear/src/quintic_extension/extension.rs b/crates/backend/koala-bear/src/quintic_extension/extension.rs index 81851c2e2..0ca166ddd 100644 --- a/crates/backend/koala-bear/src/quintic_extension/extension.rs +++ b/crates/backend/koala-bear/src/quintic_extension/extension.rs @@ -9,8 +9,7 @@ use core::iter::{Product, Sum}; use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign}; use field::{ - Algebra, BasedVectorSpace, ExtensionField, Field, Packable, PrimeCharacteristicRing, RawDataSerializable, - TwoAdicField, field_to_array, + Algebra, BasedVectorSpace, ExtensionField, Field, Packable, PrimeCharacteristicRing, TwoAdicField, field_to_array, }; use rand::distr::StandardUniform; use rand::prelude::Distribution; @@ -188,63 +187,6 @@ where impl Algebra for QuinticExtensionField {} -impl RawDataSerializable for QuinticExtensionField { - const NUM_BYTES: usize = F::NUM_BYTES * 5; - - #[inline] - fn into_bytes(self) -> impl IntoIterator { - self.value.into_iter().flat_map(|x| x.into_bytes()) - } - - #[inline] - fn into_byte_stream(input: impl IntoIterator) -> impl IntoIterator { - F::into_byte_stream(input.into_iter().flat_map(|x| x.value)) - } - - #[inline] - fn into_u32_stream(input: impl IntoIterator) -> impl IntoIterator { - F::into_u32_stream(input.into_iter().flat_map(|x| x.value)) - } - - #[inline] - fn into_u64_stream(input: impl IntoIterator) -> impl IntoIterator { - F::into_u64_stream(input.into_iter().flat_map(|x| x.value)) - } - - #[inline] - fn into_parallel_byte_streams( - input: impl IntoIterator, - ) -> impl IntoIterator { - F::into_parallel_byte_streams( - input - .into_iter() - .flat_map(|x| (0..5).map(move |i| array::from_fn(|j| x[j].value[i]))), - ) - } - - #[inline] - fn into_parallel_u32_streams( - input: impl IntoIterator, - ) -> impl IntoIterator { - F::into_parallel_u32_streams( - input - .into_iter() - .flat_map(|x| (0..5).map(move |i| array::from_fn(|j| x[j].value[i]))), - ) - } - - #[inline] - fn into_parallel_u64_streams( - input: impl IntoIterator, - ) -> impl IntoIterator { - F::into_parallel_u64_streams( - input - .into_iter() - .flat_map(|x| (0..5).map(move |i| array::from_fn(|j| x[j].value[i]))), - ) - } -} - impl Field for QuinticExtensionField { type Packing = Self; From d062a41e872ed67dcfbf36ea554d6c813ac17af6 Mon Sep 17 00:00:00 2001 From: Tom Wambsgans Date: Sun, 26 Jul 2026 23:03:59 +0200 Subject: [PATCH 5/5] field: drop constructs Aeneas cannot extract from the packing traits Two more Aeneas blockers on the sumcheck verifier, both in the packing traits. `PackedFieldExtension` declared `ExtField: ExtensionField`. Aeneas asserts signatures carry no associated-type equality constraints and reports an internal error otherwise. The equality is not needed to compile: every impl satisfies it anyway, and dropping it leaves the workspace building unchanged. `to_ext_iter(iter: impl IntoIterator) -> impl Iterator` is replaced by `to_ext_lanes(self) -> impl Iterator`. The `impl IntoIterator` argument adds a hidden type parameter to the method, so the `aeneas` preset lifted the return type into an associated type *with* that parameter -- a GAT, which Aeneas cannot extract. Unpacking one element at a time needs no method generics. Eleven of the sixteen call sites passed a single-element array, so they get shorter: `EFPacking::::to_ext_iter([x])` becomes `x.to_ext_lanes()`. The genuine iterator sites use `.flat_map(_::to_ext_lanes)`, which is what the old body did internally. One site in `air_sumcheck` needs the qualified form because with a generic `EF` both `PackedFieldExtension` impls are candidates. Aeneas now gets past both traits. The remaining blocker is not fixable this way: `Field::Packing` and `ExtensionField::ExtensionPacking` are associated types Charon cannot lift, because `Field`/`PackedField` and `ExtensionField`/`PackedFieldExtension` are mutually recursive -- each names the other. Aeneas cannot handle un-lifted associated types, and the 12 errors it reports all trace back to those two. 66 tests pass, clippy -Dwarnings and rustfmt clean, no-SIMD/AVX2/AVX-512 build. Co-Authored-By: Claude Opus 5 (1M context) --- .../backend/field/src/packed/packed_traits.rs | 17 ++++------ .../src/quintic_extension/packed_extension.rs | 11 +++---- crates/backend/poly/src/eq_mle.rs | 33 ++++++++++++------- crates/backend/poly/src/evals.rs | 2 +- .../backend/poly/src/mle/mle_single_owned.rs | 2 +- .../sumcheck/src/product_computation.rs | 12 +++---- crates/backend/sumcheck/src/sc_computation.rs | 4 +-- crates/backend/sumcheck/src/split_eq.rs | 9 +++-- crates/sub_protocols/src/air_sumcheck.rs | 4 +-- .../src/quotient_gkr/sumcheck_utils.rs | 5 ++- 10 files changed, 51 insertions(+), 48 deletions(-) diff --git a/crates/backend/field/src/packed/packed_traits.rs b/crates/backend/field/src/packed/packed_traits.rs index cdd81fecd..bd9a643b0 100644 --- a/crates/backend/field/src/packed/packed_traits.rs +++ b/crates/backend/field/src/packed/packed_traits.rs @@ -326,7 +326,7 @@ pub unsafe trait PackedFieldPow2: PackedField { /// /// This is interpreted by taking a transpose to get `[[F; D]; W]` which can then be reinterpreted /// as `[EF; W]` by making use of the chosen basis `B` again. -pub trait PackedFieldExtension>: +pub trait PackedFieldExtension>: Algebra + Algebra + BasedVectorSpace { /// Given a slice of extension field `EF` elements of length `W`, @@ -335,12 +335,10 @@ pub trait PackedFieldExtension Self; - /// Given a iterator of packed extension field elements, convert to an iterator of - /// extension field elements. - /// - /// This performs the inverse transformation to `from_ext_slice`. + /// Unpack one packed element into the `BaseField::Packing::WIDTH` extension field elements + /// it holds. Use `.flat_map(Self::to_ext_lanes)` to unpack an iterator of packed elements. #[must_use] - fn to_ext_iter(iter: impl IntoIterator) -> impl Iterator; + fn to_ext_lanes(self) -> impl Iterator; /// Given a iterator of packed extension field elements, convert to an iterator of /// extension field elements. @@ -349,7 +347,7 @@ pub trait PackedFieldExtension) -> Vec { - Self::to_ext_iter(iter).collect() + iter.into_iter().flat_map(Self::to_ext_lanes).collect() } /// Similar to `packed_powers`, construct an iterator which returns @@ -425,9 +423,8 @@ impl PackedFieldExtension for F::Packing { } #[inline] - fn to_ext_iter(iter: impl IntoIterator) -> impl Iterator { - iter.into_iter() - .flat_map(|x| (0..F::Packing::WIDTH).map(move |i| x.as_slice()[i])) + fn to_ext_lanes(self) -> impl Iterator { + (0..F::Packing::WIDTH).map(move |i| self.as_slice()[i]) } #[inline] diff --git a/crates/backend/koala-bear/src/quintic_extension/packed_extension.rs b/crates/backend/koala-bear/src/quintic_extension/packed_extension.rs index 352144c19..b118b30a2 100644 --- a/crates/backend/koala-bear/src/quintic_extension/packed_extension.rs +++ b/crates/backend/koala-bear/src/quintic_extension/packed_extension.rs @@ -255,13 +255,10 @@ where } #[inline] - fn to_ext_iter(iter: impl IntoIterator) -> impl Iterator> { - let width = F::Packing::WIDTH; - iter.into_iter().flat_map(move |x| { - (0..width).map(move |i| { - let values = array::from_fn(|j| x.value[j].as_slice()[i]); - QuinticExtensionField::new(values) - }) + fn to_ext_lanes(self) -> impl Iterator> { + (0..F::Packing::WIDTH).map(move |i| { + let values = array::from_fn(|j| self.value[j].as_slice()[i]); + QuinticExtensionField::new(values) }) } diff --git a/crates/backend/poly/src/eq_mle.rs b/crates/backend/poly/src/eq_mle.rs index 288a8fcee..3ab98c6e4 100644 --- a/crates/backend/poly/src/eq_mle.rs +++ b/crates/backend/poly/src/eq_mle.rs @@ -704,21 +704,27 @@ fn eval_eq_with_packed_scalar, const INITIALIZED match eval.len() { 0 => { - let result: Vec = EF::ExtensionPacking::to_ext_iter([scalar]).collect(); + let result: Vec = (scalar).to_ext_lanes().collect(); add_or_set_f::<_, INITIALIZED>(out, &result); } 1 => { // Manually unroll for single variable case let eq_evaluations = eval_eq_1(eval, scalar); - let result: Vec = EF::ExtensionPacking::to_ext_iter(eq_evaluations).collect(); + let result: Vec = eq_evaluations + .into_iter() + .flat_map(EF::ExtensionPacking::to_ext_lanes) + .collect(); add_or_set_f::<_, INITIALIZED>(out, &result); } 2 => { // Manually unroll for two variables case let eq_evaluations = eval_eq_2(eval, scalar); - let result: Vec = EF::ExtensionPacking::to_ext_iter(eq_evaluations).collect(); + let result: Vec = eq_evaluations + .into_iter() + .flat_map(EF::ExtensionPacking::to_ext_lanes) + .collect(); add_or_set_f::<_, INITIALIZED>(out, &result); } 3 => { @@ -735,15 +741,18 @@ fn eval_eq_with_packed_scalar, const INITIALIZED // This avoids the allocation used to accumulate `result` in the other branches. We could // do a similar strategy in those branches but, those branches should only be hit // infrequently in small cases which are already sufficiently fast. - iter_array_chunks_padded::<_, EVAL_LEN>(EF::ExtensionPacking::to_ext_iter(eq_evaluations), EF::ZERO) - .zip(out.as_chunks_mut::().0.iter_mut()) - .for_each(|(res, out_chunk)| { - if INITIALIZED { - EF::add_slices(out_chunk, &res); - } else { - out_chunk.copy_from_slice(&res); - } - }); + iter_array_chunks_padded::<_, EVAL_LEN>( + eq_evaluations.into_iter().flat_map(EF::ExtensionPacking::to_ext_lanes), + EF::ZERO, + ) + .zip(out.as_chunks_mut::().0.iter_mut()) + .for_each(|(res, out_chunk)| { + if INITIALIZED { + EF::add_slices(out_chunk, &res); + } else { + out_chunk.copy_from_slice(&res); + } + }); } _ => { let (&x, tail) = eval.split_first().unwrap(); diff --git a/crates/backend/poly/src/evals.rs b/crates/backend/poly/src/evals.rs index d054a710f..e23264a5a 100644 --- a/crates/backend/poly/src/evals.rs +++ b/crates/backend/poly/src/evals.rs @@ -187,7 +187,7 @@ where for (&a, &b) in part.iter().zip(eq_lo.iter()) { acc += b * a; } - EFPacking::::to_ext_iter([acc]).sum::() * eq_hi[i] + (acc).to_ext_lanes().sum::() * eq_hi[i] }; let work_size: usize = (1 << 15) / std::mem::size_of::>(); diff --git a/crates/backend/poly/src/mle/mle_single_owned.rs b/crates/backend/poly/src/mle/mle_single_owned.rs index 55c1f9234..c993bf7b7 100644 --- a/crates/backend/poly/src/mle/mle_single_owned.rs +++ b/crates/backend/poly/src/mle/mle_single_owned.rs @@ -167,7 +167,7 @@ impl MleOwned { } Self::ExtensionPacked(v) => { assert_eq!(packing_width::(), 1); - EFPacking::::to_ext_iter([v[0]]).nth(0).unwrap() + (v[0]).to_ext_lanes().nth(0).unwrap() } } } diff --git a/crates/backend/sumcheck/src/product_computation.rs b/crates/backend/sumcheck/src/product_computation.rs index 167fe734e..f1b42e3f2 100644 --- a/crates/backend/sumcheck/src/product_computation.rs +++ b/crates/backend/sumcheck/src/product_computation.rs @@ -45,10 +45,10 @@ pub fn run_product_sumcheck( assert!(n_rounds >= 1); let first_sumcheck_poly = match (pol_a, pol_b) { (MleRef::BasePacked(evals), MleRef::ExtensionPacked(weights)) => { - compute_product_sumcheck_polynomial(evals, weights, sum, |e| EFPacking::::to_ext_iter([e]).collect()) + compute_product_sumcheck_polynomial(evals, weights, sum, |e| (e).to_ext_lanes().collect()) } (MleRef::ExtensionPacked(evals), MleRef::ExtensionPacked(weights)) => { - compute_product_sumcheck_polynomial(evals, weights, sum, |e| EFPacking::::to_ext_iter([e]).collect()) + compute_product_sumcheck_polynomial(evals, weights, sum, |e| (e).to_ext_lanes().collect()) } (MleRef::Base(evals), MleRef::Extension(weights)) => { compute_product_sumcheck_polynomial(evals, weights, sum, |e| vec![e]) @@ -85,16 +85,12 @@ pub fn run_product_sumcheck_from_round1( let (second_sumcheck_poly, folded) = match (pol_a, pol_b) { (MleRef::BasePacked(evals), MleRef::ExtensionPacked(weights)) => { let (second_sumcheck_poly, folded) = - fold_and_compute_product_sumcheck_polynomial(evals, weights, r1, sum, |e| { - EFPacking::::to_ext_iter([e]).collect() - }); + fold_and_compute_product_sumcheck_polynomial(evals, weights, r1, sum, |e| (e).to_ext_lanes().collect()); (second_sumcheck_poly, MleGroupOwned::ExtensionPacked(folded)) } (MleRef::ExtensionPacked(evals), MleRef::ExtensionPacked(weights)) => { let (second_sumcheck_poly, folded) = - fold_and_compute_product_sumcheck_polynomial(evals, weights, r1, sum, |e| { - EFPacking::::to_ext_iter([e]).collect() - }); + fold_and_compute_product_sumcheck_polynomial(evals, weights, r1, sum, |e| (e).to_ext_lanes().collect()); (second_sumcheck_poly, MleGroupOwned::ExtensionPacked(folded)) } (MleRef::Base(evals), MleRef::Extension(weights)) => { diff --git a/crates/backend/sumcheck/src/sc_computation.rs b/crates/backend/sumcheck/src/sc_computation.rs index a72a19284..822c054e2 100644 --- a/crates/backend/sumcheck/src/sc_computation.rs +++ b/crates/backend/sumcheck/src/sc_computation.rs @@ -85,12 +85,12 @@ pub(crate) fn identity_decompose(e: EF) -> Vec { #[inline(always)] pub(crate) fn packing_decompose(e: EFPacking) -> Vec { - EFPacking::::to_ext_iter([e]).collect() + (e).to_ext_lanes().collect() } #[inline(always)] pub fn packing_unpack_sum(s: EFPacking) -> EF { - EFPacking::::to_ext_iter([s]).sum::() + (s).to_ext_lanes().sum::() } fn handle_product_computation<'a, EF: KoalaBearExtension>(group: &MleGroupRef<'a, EF>, sum: EF) -> Vec { diff --git a/crates/backend/sumcheck/src/split_eq.rs b/crates/backend/sumcheck/src/split_eq.rs index b784849d3..4e23b0b0d 100644 --- a/crates/backend/sumcheck/src/split_eq.rs +++ b/crates/backend/sumcheck/src/split_eq.rs @@ -54,7 +54,12 @@ impl SplitEq { self.log_packed_hi = new_len.trailing_zeros(); } else { // eq_hi_packed has 0 or 1 element — unpack to remainder and halve - let mut unpacked: ArenaVec = EFPacking::::to_ext_iter(self.eq_hi_packed.iter().copied()).collect(); + let mut unpacked: ArenaVec = self + .eq_hi_packed + .iter() + .copied() + .flat_map(EFPacking::::to_ext_lanes) + .collect(); let scale = self.eq_lo[0]; for v in &mut unpacked { *v *= scale; @@ -98,7 +103,7 @@ impl SplitEq { } else { let width = packing_width::(); let packed_val = self.get_packed(i / width); - EFPacking::::to_ext_iter([packed_val]).nth(i % width).unwrap() + (packed_val).to_ext_lanes().nth(i % width).unwrap() } } } diff --git a/crates/sub_protocols/src/air_sumcheck.rs b/crates/sub_protocols/src/air_sumcheck.rs index 3263cf4a1..83a9b8be0 100644 --- a/crates/sub_protocols/src/air_sumcheck.rs +++ b/crates/sub_protocols/src/air_sumcheck.rs @@ -304,7 +304,7 @@ fn column_evals(multilinears: &MleGroupRef<'_, EF>, i: u MleGroupRef::ExtensionPacked(cols) => { let (packed_i, lane) = (i >> packing_log_width::(), i & (packing_width::() - 1)); cols.iter() - .map(|c| EFPacking::::to_ext_iter([c[packed_i]]).nth(lane).unwrap()) + .map(|c| (c[packed_i]).to_ext_lanes().nth(lane).unwrap()) .collect() } } @@ -325,7 +325,7 @@ where A::ExtraData: AlphaPowers, { let unpack_sum_packed = - |s: EFPacking| -> EF { as PackedFieldExtension, EF>>::to_ext_iter([s]).sum::() }; + |s: EFPacking| -> EF { as PackedFieldExtension, EF>>::to_ext_lanes(s).sum::() }; if let Some((low_degree, low_n_constraints)) = computation.low_degree_air() { match multilinears { diff --git a/crates/sub_protocols/src/quotient_gkr/sumcheck_utils.rs b/crates/sub_protocols/src/quotient_gkr/sumcheck_utils.rs index 30adff7d8..1d59a4385 100644 --- a/crates/sub_protocols/src/quotient_gkr/sumcheck_utils.rs +++ b/crates/sub_protocols/src/quotient_gkr/sumcheck_utils.rs @@ -94,9 +94,8 @@ fn finalize_round( mmf: &mut EF, padding_correction: EF, ) -> EF { - let c0_raw: EF = - EFPacking::::to_ext_iter([coeffs.c0_num + coeffs.c0_den * alpha]).sum::() + padding_correction; - let c2_raw: EF = EFPacking::::to_ext_iter([coeffs.c2_num + coeffs.c2_den * alpha]).sum(); + let c0_raw: EF = (coeffs.c0_num + coeffs.c0_den * alpha).to_ext_lanes().sum::() + padding_correction; + let c2_raw: EF = (coeffs.c2_num + coeffs.c2_den * alpha).to_ext_lanes().sum(); let bare = build_bare_from_coeffs(c0_raw, c2_raw, eq_alpha, *sum, *mmf); prover_state.add_sumcheck_polynomial(&bare.coeffs, Some(eq_alpha)); let r = prover_state.sample();