refactor: decompose composition poly via squared-coset intermediate

crypto/stark/src/batched_layout.rs

@@ -0,0 +1,76 @@
+/// Tracks per-table column offsets within a shared batched Merkle tree.
+/// Each phase (main, aux, composition) has its own layout.
+#[derive(Debug, Clone)]
+pub struct BatchedLayout {
+    /// (col_start, col_end) for each table in the concatenated row.
+    pub table_ranges: Vec<(usize, usize)>,
+    /// Total number of columns across all tables.
+    pub total_columns: usize,
+    /// Domain size (LDE size for main/aux, trace size for composition).
+    pub domain_size: usize,
+}
+
+impl BatchedLayout {
+    /// Build layout from per-table column counts.
+    pub fn new(column_counts: &[usize], domain_size: usize) -> Self {
+        let mut ranges = Vec::with_capacity(column_counts.len());
+        let mut offset = 0;
+        for &count in column_counts {
+            ranges.push((offset, offset + count));
+            offset += count;
+        }
+        BatchedLayout {
+            table_ranges: ranges,
+            total_columns: offset,
+            domain_size,
+        }
+    }
+
+    pub fn num_tables(&self) -> usize {
+        self.table_ranges.len()
+    }
+
+    /// Extract one table's columns from a full row of opened values.
+    pub fn extract_table<T: Clone>(&self, table_idx: usize, row: &[T]) -> Vec<T> {
+        let (start, end) = self.table_ranges[table_idx];
+        row[start..end].to_vec()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_layout_construction() {
+        let layout = BatchedLayout::new(&[5, 3, 7], 1024);
+        assert_eq!(layout.total_columns, 15);
+        assert_eq!(layout.table_ranges, vec![(0, 5), (5, 8), (8, 15)]);
+        assert_eq!(layout.num_tables(), 3);
+        assert_eq!(layout.domain_size, 1024);
+    }
+
+    #[test]
+    fn test_extract_table() {
+        let layout = BatchedLayout::new(&[2, 3], 1024);
+        let row = vec![10, 20, 30, 40, 50];
+        assert_eq!(layout.extract_table(0, &row), vec![10, 20]);
+        assert_eq!(layout.extract_table(1, &row), vec![30, 40, 50]);
+    }
+
+    #[test]
+    fn test_empty_tables() {
+        let layout = BatchedLayout::new(&[0, 3, 0], 512);
+        assert_eq!(layout.total_columns, 3);
+        assert_eq!(layout.table_ranges, vec![(0, 0), (0, 3), (3, 3)]);
+        assert_eq!(layout.extract_table::<i32>(0, &[1, 2, 3]), vec![]);
+        assert_eq!(layout.extract_table(1, &[1, 2, 3]), vec![1, 2, 3]);
+    }
+
+    #[test]
+    fn test_single_table() {
+        let layout = BatchedLayout::new(&[4], 2048);
+        assert_eq!(layout.total_columns, 4);
+        assert_eq!(layout.table_ranges, vec![(0, 4)]);
+    }
+}

crypto/stark/src/fri/mod.rs

@@ -85,6 +85,89 @@ where
     (last_value, fri_layer_list)
 }
 
+/// FRI commit phase with optional injection after the first fold.
+///
+/// After the first fold reduces 2N evaluations to N, the `injection` vector
+/// (if provided) is added elementwise. This allows committing auxiliary
+/// polynomials (e.g. composition DEEP quotient) at a smaller domain size
+/// without extending them to the full LDE domain.
+///
+/// The injection vector must be in bit-reversed order on the squared coset,
+/// matching the output order of the first fold.
+pub fn commit_phase_from_evaluations_with_injection<
+    F: IsFFTField + IsSubFieldOf<E>,
+    E: IsField,
+>(
+    number_layers: usize,
+    mut evals: Vec<FieldElement<E>>,
+    injection: Option<Vec<FieldElement<E>>>,
+    transcript: &mut impl IsStarkTranscript<E, F>,
+    coset_offset: &FieldElement<F>,
+    domain_size: usize,
+) -> (
+    FieldElement<E>,
+    Vec<FriLayer<E, FriLayerMerkleTreeBackend<E>>>,
+)
+where
+    FieldElement<F>: AsBytes + Sync + Send,
+    FieldElement<E>: AsBytes + Sync + Send,
+{
+    let mut inv_twiddles = compute_coset_twiddles_inv(coset_offset, domain_size);
+
+    let mut fri_layer_list = Vec::with_capacity(number_layers);
+    let mut current_coset_offset = coset_offset.clone();
+    let mut current_domain_size = domain_size;
+    let mut first_fold = true;
+
+    for _ in 1..number_layers {
+        let zeta = transcript.sample_field_element();
+        current_coset_offset = current_coset_offset.square();
+        current_domain_size /= 2;
+
+        fold_evaluations_in_place(&mut evals, &zeta, &inv_twiddles);
+
+        // After the first fold, inject composition DEEP values if provided.
+        if first_fold {
+            if let Some(ref inj) = injection {
+                debug_assert_eq!(
+                    evals.len(),
+                    inj.len(),
+                    "Injection vector size must match folded evaluation size"
+                );
+                for (e, v) in evals.iter_mut().zip(inj.iter()) {
+                    *e = &*e + v;
+                }
+            }
+            first_fold = false;
+        }
+
+        let leaves: Vec<[FieldElement<E>; 2]> = evals
+            .chunks_exact(2)
+            .map(|chunk| [chunk[0].clone(), chunk[1].clone()])
+            .collect();
+        let merkle_tree = FriLayerMerkleTree::build(&leaves)
+            .expect("FRI commit: Merkle tree construction must succeed");
+        let root = merkle_tree.root;
+        fri_layer_list.push(FriLayer::new(
+            &evals,
+            merkle_tree,
+            current_coset_offset.clone().to_extension(),
+            current_domain_size,
+        ));
+
+        transcript.append_bytes(&root);
+        update_twiddles_in_place(&mut inv_twiddles);
+    }
+
+    let zeta = transcript.sample_field_element();
+    fold_evaluations_in_place(&mut evals, &zeta, &inv_twiddles);
+
+    let last_value = evals.first().unwrap_or(&FieldElement::zero()).clone();
+    transcript.append_field_element(&last_value);
+
+    (last_value, fri_layer_list)
+}
+
 pub fn query_phase<F: IsField>(
     fri_layers: &Vec<FriLayer<F, FriLayerMerkleTreeBackend<F>>>,
     iotas: &[usize],

crypto/stark/src/lib.rs

@@ -1,5 +1,6 @@
 #[cfg(feature = "debug-checks")]
 pub mod bus_debug;
+pub mod batched_layout;
 pub mod constraints;
 pub mod context;
 pub mod debug;

crypto/stark/src/proof/stark.rs

@@ -78,3 +78,71 @@ pub struct StarkProof<F: IsSubFieldOf<E>, E: IsField, PI> {
 pub struct MultiProof<F: IsSubFieldOf<E>, E: IsField, PI> {
     pub proofs: Vec<StarkProof<F, E, PI>>,
 }
+
+/// Per-table data in a batched proof (OOD evaluations + public inputs).
+#[derive(Debug, Clone)]
+pub struct TableProofData<F: IsField, E: IsField, PI> {
+    pub trace_length: usize,
+    pub public_inputs: PI,
+    pub trace_ood_evaluations: Table<E>,
+    pub composition_poly_parts_ood_evaluation: Vec<FieldElement<E>>,
+    pub bus_public_inputs: Option<BusPublicInputs<E>>,
+    /// Number of composition poly parts for this table.
+    pub num_composition_parts: usize,
+    _phantom: core::marker::PhantomData<F>,
+}
+
+impl<F: IsField, E: IsField, PI> TableProofData<F, E, PI> {
+    pub fn new(
+        trace_length: usize,
+        public_inputs: PI,
+        trace_ood_evaluations: Table<E>,
+        composition_poly_parts_ood_evaluation: Vec<FieldElement<E>>,
+        bus_public_inputs: Option<BusPublicInputs<E>>,
+        num_composition_parts: usize,
+    ) -> Self {
+        Self {
+            trace_length,
+            public_inputs,
+            trace_ood_evaluations,
+            composition_poly_parts_ood_evaluation,
+            bus_public_inputs,
+            num_composition_parts,
+            _phantom: core::marker::PhantomData,
+        }
+    }
+}
+
+/// Per-query opening from all shared Merkle trees.
+#[derive(Debug, Clone)]
+pub struct BatchedQueryOpening<F: IsField, E: IsField> {
+    /// Main trace: opened row and symmetric row + Merkle proofs.
+    pub main_trace: PolynomialOpenings<F>,
+    /// Aux trace: opened row and symmetric row + Merkle proofs (if any table has aux).
+    pub aux_trace: Option<PolynomialOpenings<E>>,
+    /// Composition poly: opened values + Merkle proof.
+    pub composition_poly: PolynomialOpenings<E>,
+}
+
+/// Proof format with shared Merkle trees across all tables.
+#[derive(Debug, Clone)]
+pub struct BatchedProof<F: IsField, E: IsField, PI> {
+    /// Shared Merkle root for all tables' main trace columns.
+    pub main_merkle_root: Commitment,
+    /// Shared Merkle root for all tables' aux trace columns.
+    pub aux_merkle_root: Option<Commitment>,
+    /// Shared Merkle root for all tables' composition poly evaluations.
+    pub composition_merkle_root: Commitment,
+    /// Per-table OOD data and public inputs.
+    pub tables: Vec<TableProofData<F, E, PI>>,
+    /// Shared FRI layer roots.
+    pub fri_layers_merkle_roots: Vec<Commitment>,
+    /// Shared FRI last value.
+    pub fri_last_value: FieldElement<E>,
+    /// Shared FRI decommitments (one per query).
+    pub query_list: Vec<FriDecommitment<E>>,
+    /// Per-query openings from the shared trees.
+    pub query_openings: Vec<BatchedQueryOpening<F, E>>,
+    /// Grinding nonce.
+    pub nonce: Option<u64>,
+}