air: use ref for column groups in AirWitness (#27)

Author: tcoratger

crates/air/src/prove.rs

@@ -286,7 +286,7 @@ fn eval_unstructured_column_groups<EF: ExtensionField<PF<EF>> + ExtensionField<I
     columns_batching_scalars: &[EF],
 ) -> Vec<Vec<EF>> {
     let mut all_sub_evals = vec![];
-    for group in &witnesses.column_groups {
+    for group in witnesses.column_groups {
         let batched_column = multilinears_linear_combination(
             &witnesses.cols[group.clone()],
             &eval_eq(from_end(
@@ -403,7 +403,7 @@ fn open_structured_columns<EF: ExtensionField<PF<EF>> + ExtensionField<IF>, IF:
     let mut all_inner_sums = vec![];
     let mut all_batched_columns = vec![];
     let mut all_batched_columns_mixed = vec![];
-    for group in &witness.column_groups {
+    for group in witness.column_groups {
         let batched_column = multilinears_linear_combination(
             &witness.cols[group.clone()],
             &eval_eq(from_end(

crates/air/src/witness.rs

@@ -10,7 +10,7 @@ use p3_util::{log2_ceil_usize, log2_strict_usize};
 #[derive(Debug)]
 pub struct AirWitness<'a, F> {
     pub cols: Vec<&'a [F]>,
-    pub column_groups: Vec<Range<usize>>,
+    pub column_groups: &'a [Range<usize>],
 }
 
 impl<'a, F> Deref for AirWitness<'a, F> {
@@ -21,11 +21,8 @@ impl<'a, F> Deref for AirWitness<'a, F> {
 }
 
 impl<'a, F> AirWitness<'a, F> {
-    pub fn new(cols: &'a [impl Borrow<[F]>], column_groups: &[Range<usize>]) -> Self {
-        let cols = cols
-            .iter()
-            .map(std::borrow::Borrow::borrow)
-            .collect::<Vec<_>>();
+    pub fn new(cols: &'a [impl Borrow<[F]>], column_groups: &'a [Range<usize>]) -> Self {
+        let cols = cols.iter().map(Borrow::borrow).collect::<Vec<_>>();
         assert!(
             cols.iter()
                 .all(|col| col.len() == (1 << log2_strict_usize(cols[0].len()))),
@@ -36,7 +33,7 @@ impl<'a, F> AirWitness<'a, F> {
         assert!(column_groups.iter().all(|r| r.start < r.end));
         Self {
             cols,
-            column_groups: column_groups.to_vec(),
+            column_groups,
         }
     }
 

crates/zk_vm/src/prove_execution.rs

@@ -80,7 +80,8 @@ pub fn prove_execution(
             .iter()
             .map(Vec::as_slice),
     );
-    let exec_witness = AirWitness::<PF<EF>>::new(&exec_columns, &exec_column_groups());
+    let exec_column_groups = exec_column_groups();
+    let exec_witness = AirWitness::new(&exec_columns, &exec_column_groups);
     let exec_table = AirTable::<EF, _>::new(VMAir);
 
     #[cfg(test)]
@@ -96,8 +97,10 @@ pub fn prove_execution(
     let dot_product_table = AirTable::<EF, _>::new(DotProductAir);
 
     let (p16_columns, p24_columns) = build_poseidon_columns(&poseidons_16, &poseidons_24);
-    let p16_witness = AirWitness::new(&p16_columns, &poseidon_16_column_groups(&p16_air));
-    let p24_witness = AirWitness::new(&p24_columns, &poseidon_24_column_groups(&p24_air));
+    let p16_column_groups = poseidon_16_column_groups(&p16_air);
+    let p24_column_groups = poseidon_24_column_groups(&p24_air);
+    let p16_witness = AirWitness::new(&p16_columns, &p16_column_groups);
+    let p24_witness = AirWitness::new(&p24_columns, &p24_column_groups);
 
     let (dot_product_columns, dot_product_padding_len) = build_dot_product_columns(&dot_products);
     let dot_product_witness = AirWitness::new(&dot_product_columns, &DOT_PRODUCT_AIR_COLUMN_GROUPS);