feat(prover): slim MEMW table and add MEMW_A aligned fast path

prover/src/lib.rs

@@ -36,7 +36,8 @@ use crate::tables::trace_builder::Traces;
 use crate::test_utils::{
     E, F, VmAir, create_bitwise_air, create_branch_air, create_commit_air, create_cpu_air,
     create_decode_air, create_dvrm_air, create_halt_air, create_load_air, create_lt_air,
-    create_memw_air, create_mul_air, create_page_air, create_register_air, create_shift_air,
+    create_memw_air, create_memw_aligned_air, create_mul_air, create_page_air, create_register_air,
+    create_shift_air,
 };
 
 use stark::proof::options::{GoldilocksCubicProofOptions, ProofOptions};
@@ -61,6 +62,7 @@ pub struct TableCounts {
     pub cpu: usize,
     pub lt: usize,
     pub memw: usize,
+    pub memw_aligned: usize,
     pub load: usize,
     pub mul: usize,
     pub dvrm: usize,
@@ -75,7 +77,15 @@ impl TableCounts {
     /// allowing a malicious prover to bypass soundness checks.
     /// Sum of all chunk counts across split tables.
     pub fn total(&self) -> usize {
-        self.cpu + self.lt + self.memw + self.load + self.mul + self.dvrm + self.shift + self.branch
+        self.cpu
+            + self.lt
+            + self.memw
+            + self.memw_aligned
+            + self.load
+            + self.mul
+            + self.dvrm
+            + self.shift
+            + self.branch
     }
 
     /// Validate that all required tables have at least one chunk.
@@ -87,6 +97,7 @@ impl TableCounts {
             ("cpu", self.cpu),
             ("lt", self.lt),
             ("memw", self.memw),
+            ("memw_aligned", self.memw_aligned),
             ("load", self.load),
             ("mul", self.mul),
             ("dvrm", self.dvrm),
@@ -167,6 +178,7 @@ pub(crate) struct VmAirs {
     pub lts: Vec<VmAir>,
     pub shifts: Vec<VmAir>,
     pub memws: Vec<VmAir>,
+    pub memw_aligneds: Vec<VmAir>,
     pub loads: Vec<VmAir>,
     pub decode: VmAir,
     pub muls: Vec<VmAir>,
@@ -201,6 +213,13 @@ impl VmAirs {
         for (air, trace) in self.memws.iter().zip(traces.memws.iter_mut()) {
             pairs.push((air, trace, &()));
         }
+        for (air, trace) in self
+            .memw_aligneds
+            .iter()
+            .zip(traces.memw_aligneds.iter_mut())
+        {
+            pairs.push((air, trace, &()));
+        }
         for (air, trace) in self.loads.iter().zip(traces.loads.iter_mut()) {
             pairs.push((air, trace, &()));
         }
@@ -242,6 +261,9 @@ impl VmAirs {
         for air in &self.memws {
             refs.push(air);
         }
+        for air in &self.memw_aligneds {
+            refs.push(air);
+        }
         for air in &self.loads {
             refs.push(air);
         }
@@ -294,6 +316,9 @@ impl VmAirs {
         let memws: Vec<_> = (0..table_counts.memw)
             .map(|i| create_memw_air(proof_options).with_name(&format!("MEMW[{}]", i)))
             .collect();
+        let memw_aligneds: Vec<_> = (0..table_counts.memw_aligned)
+            .map(|i| create_memw_aligned_air(proof_options).with_name(&format!("MEMW_A[{}]", i)))
+            .collect();
         let loads: Vec<_> = (0..table_counts.load)
             .map(|i| create_load_air(proof_options).with_name(&format!("LOAD[{}]", i)))
             .collect();
@@ -336,6 +361,7 @@ impl VmAirs {
             lts,
             shifts,
             memws,
+            memw_aligneds,
             loads,
             decode,
             muls,