prover: enforce load/store memory alignment in AIR

prover/src/chips/instructions/i/load_store.rs

@@ -120,6 +120,22 @@ impl MachineChip for LoadStoreChip {
             add_with_carries(value_a, offset)
         };
         traces.fill_columns(row_idx, ram_base_address, Column::RamBaseAddr);
+
+        let alignment_divisor: Option<u8> = match vm_step.step.instruction.opcode.builtin() {
+            Some(BuiltinOpcode::LH) | Some(BuiltinOpcode::LHU) | Some(BuiltinOpcode::SH) => Some(2),
+            Some(BuiltinOpcode::LW) | Some(BuiltinOpcode::SW) => Some(4),
+            _ => None,
+        };
+        if let Some(divisor) = alignment_divisor {
+            // Witness data only: malformed traces must be rejected by AIR
+            // constraints, not by panicking during trace construction.
+            traces.fill_columns(
+                row_idx,
+                ram_base_address[0] / divisor,
+                Column::RamBaseAddrAlignmentQuotient,
+            );
+        }
+
         let carry_bits = [carry_bits[1], carry_bits[3]];
         traces.fill_columns(row_idx, carry_bits, Column::CarryFlag);
         let clk = row_idx as u32 + 1;
@@ -447,6 +463,32 @@ impl MachineChip for LoadStoreChip {
         eval.add_constraint(helper4[WORD_SIZE - 1].clone() * ram3_4_accessed.clone());
 
         let ram_base_addr = trace_eval!(trace_eval, Column::RamBaseAddr);
+        {
+            let [is_lh] = trace_eval!(trace_eval, IsLh);
+            let [is_lhu] = trace_eval!(trace_eval, IsLhu);
+            let [is_lw] = trace_eval!(trace_eval, IsLw);
+            let [is_sh] = trace_eval!(trace_eval, IsSh);
+            let [is_sw] = trace_eval!(trace_eval, IsSw);
+            let [alignment_quotient] =
+                trace_eval!(trace_eval, Column::RamBaseAddrAlignmentQuotient);
+
+            let is_halfword_access = is_lh + is_lhu + is_sh;
+            let is_word_access = is_lw + is_sw;
+
+            // Close the VM/AIR semantic gap: the runtime memory layer rejects
+            // unaligned halfword and word accesses, and these constraints force
+            // the proved effective address to obey the same rule.
+            eval.add_constraint(
+                is_halfword_access
+                    * (ram_base_addr[0].clone()
+                        - alignment_quotient.clone() * BaseField::from(2u32)),
+            );
+
+            eval.add_constraint(
+                is_word_access
+                    * (ram_base_addr[0].clone() - alignment_quotient * BaseField::from(4u32)),
+            );
+        }
         let carry_flag = trace_eval!(trace_eval, Column::CarryFlag);
         let value_b = trace_eval!(trace_eval, Column::ValueB);
         let value_c = trace_eval!(trace_eval, Column::ValueC);
@@ -1079,6 +1121,73 @@ mod test {
         .unwrap();
     }
 
+    fn alignment_trace(op: Column, address: u32, quotient: u8) -> TracesBuilder {
+        let mut traces = TracesBuilder::new(LOG_SIZE);
+
+        traces.fill_columns(0, true, op);
+        traces.fill_columns(0, address, Column::RamBaseAddr);
+        traces.fill_columns(0, quotient, Column::RamBaseAddrAlignmentQuotient);
+
+        match op {
+            Column::IsSh | Column::IsSw => {
+                traces.fill_columns(0, address, Column::ValueA);
+                traces.fill_columns(0, 0u32, Column::ValueB);
+            }
+            _ => {
+                traces.fill_columns(0, 0u32, Column::ValueA);
+                traces.fill_columns(0, address, Column::ValueB);
+            }
+        }
+
+        traces
+    }
+
+    fn assert_alignment_constraints_reject(op: Column, address: u32, quotient: u8) {
+        let traces = alignment_trace(op, address, quotient);
+        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || {
+            assert_chip::<LoadStoreChip>(traces, None);
+        }));
+
+        assert!(
+            result.is_err(),
+            "unaligned load/store trace unexpectedly satisfied LoadStoreChip constraints"
+        );
+    }
+
+    #[test]
+    fn test_unaligned_lh_rejected_in_air() {
+        assert_alignment_constraints_reject(Column::IsLh, 0x0000_1001, 0);
+    }
+
+    #[test]
+    fn test_unaligned_lhu_rejected_in_air() {
+        assert_alignment_constraints_reject(Column::IsLhu, 0x0000_1001, 0);
+    }
+
+    #[test]
+    fn test_unaligned_lw_rejected_in_air() {
+        assert_alignment_constraints_reject(Column::IsLw, 0x0000_1002, 0);
+    }
+
+    #[test]
+    fn test_issue_605_unaligned_halfword_store_rejected_in_air() {
+        assert_alignment_constraints_reject(Column::IsSh, 0x0000_1001, 0);
+    }
+
+    #[test]
+    fn test_issue_605_unaligned_word_store_rejected_in_air() {
+        assert_alignment_constraints_reject(Column::IsSw, 0x0000_1002, 0);
+    }
+
+    #[test]
+    fn test_aligned_halfword_and_word_accesses_accepted_in_air() {
+        assert_chip::<LoadStoreChip>(alignment_trace(Column::IsLh, 0x0000_1002, 1), None);
+        assert_chip::<LoadStoreChip>(alignment_trace(Column::IsLhu, 0x0000_1002, 1), None);
+        assert_chip::<LoadStoreChip>(alignment_trace(Column::IsLw, 0x0000_1004, 1), None);
+        assert_chip::<LoadStoreChip>(alignment_trace(Column::IsSh, 0x0000_1002, 1), None);
+        assert_chip::<LoadStoreChip>(alignment_trace(Column::IsSw, 0x0000_1004, 1), None);
+    }
+
     #[test]
     fn test_invalid_store() {
         let basic_block = BasicBlock::new(vec![

prover/src/chips/range_check/range128.rs

@@ -71,6 +71,19 @@ impl MachineChip for Range128Chip {
         fill_main_col(qt_aux, is_lh, side_note);
         let [is_lb] = traces.column(row_idx, Column::IsLb);
         fill_main_col(qt_aux, is_lb, side_note);
+
+        let [alignment_quotient] = traces.column(row_idx, Column::RamBaseAddrAlignmentQuotient);
+        let [is_lh_alignment] = traces.column(row_idx, Column::IsLh);
+        let [is_lhu] = traces.column(row_idx, Column::IsLhu);
+        let [is_lw] = traces.column(row_idx, Column::IsLw);
+        let [is_sh] = traces.column(row_idx, Column::IsSh);
+        let [is_sw] = traces.column(row_idx, Column::IsSw);
+        let alignment_selector = is_lh_alignment + is_lhu + is_lw + is_sh + is_sw;
+
+        // q is used by LoadStoreChip to prove addr_low = 2*q or addr_low = 4*q.
+        // Keep q in [0, 127], so unaligned addresses cannot be represented by
+        // arbitrary field elements such as 1/2 or 3/4.
+        fill_main_col(alignment_quotient, alignment_selector, side_note);
     }
     /// Fills the whole interaction trace in one-go using SIMD in the stwo-usual way
     ///
@@ -125,6 +138,23 @@ impl MachineChip for Range128Chip {
             logup_trace_gen,
             lookup_element,
         );
+
+        let [alignment_quotient] =
+            original_traces.get_base_column(Column::RamBaseAddrAlignmentQuotient);
+        let [is_lh_alignment] = original_traces.get_base_column(Column::IsLh);
+        let [is_lhu] = original_traces.get_base_column(Column::IsLhu);
+        let [is_lw] = original_traces.get_base_column(Column::IsLw);
+        let [is_sh] = original_traces.get_base_column(Column::IsSh);
+        let [is_sw] = original_traces.get_base_column(Column::IsSw);
+        let alignment_selectors = &[is_lh_alignment, is_lhu, is_lw, is_sh, is_sw];
+
+        check_col(
+            alignment_quotient,
+            alignment_selectors,
+            original_traces.log_size(),
+            logup_trace_gen,
+            lookup_element,
+        );
     }
 
     fn add_constraints<E: stwo_constraint_framework::EvalAtRow>(
@@ -180,6 +210,20 @@ impl MachineChip for Range128Chip {
             numerator.into(),
             &[qt_aux],
         ));
+
+        let [alignment_quotient] = trace_eval.column_eval(Column::RamBaseAddrAlignmentQuotient);
+        let [is_lh_alignment] = trace_eval.column_eval(Column::IsLh);
+        let [is_lhu] = trace_eval.column_eval(Column::IsLhu);
+        let [is_lw] = trace_eval.column_eval(Column::IsLw);
+        let [is_sh] = trace_eval.column_eval(Column::IsSh);
+        let [is_sw] = trace_eval.column_eval(Column::IsSw);
+        let alignment_selector = is_lh_alignment + is_lhu + is_lw + is_sh + is_sw;
+
+        eval.add_to_relation(RelationEntry::new(
+            lookup_elements,
+            alignment_selector.into(),
+            &[alignment_quotient],
+        ));
     }
 }
 
@@ -313,4 +357,54 @@ mod test {
             ext.generate_interaction_trace(component_trace, &side_note, &lookup_elements);
         assert_ne!(claimed_sum + claimed_sum_2, SecureField::zero());
     }
+
+    #[test]
+    fn test_alignment_quotient_is_range_checked() {
+        const LOG_SIZE: u32 = PreprocessedBuilder::MIN_LOG_SIZE;
+        let (config, twiddles) = test_params(LOG_SIZE);
+        let mut traces = TracesBuilder::new(LOG_SIZE);
+        let program_info = ProgramInfo::dummy();
+        let program_trace_ref = ProgramTraceRef {
+            program_memory: &program_info,
+            init_memory: Default::default(),
+            exit_code: Default::default(),
+            public_output: Default::default(),
+        };
+        let program_traces = ProgramTracesBuilder::new(LOG_SIZE, program_trace_ref);
+        let mut side_note = SideNote::new(&program_traces, &HarvardEmulator::default().finalize());
+
+        for row_idx in 0..(1 << LOG_SIZE) {
+            traces.fill_columns(row_idx, true, Column::IsLw);
+            traces.fill_columns(row_idx, 0u8, Column::RamBaseAddrAlignmentQuotient);
+
+            Range128Chip::fill_main_trace(
+                &mut traces,
+                row_idx,
+                &Some(ProgramStep::default()),
+                &mut side_note,
+                &ExtensionsConfig::default(),
+            );
+        }
+
+        *traces.column_mut::<{ Column::RamBaseAddrAlignmentQuotient.size() }>(
+            0,
+            Column::RamBaseAddrAlignmentQuotient,
+        )[0] = BaseField::from(128u32);
+
+        let CommittedTraces {
+            claimed_sum,
+            lookup_elements,
+            ..
+        } = commit_traces::<Range128Chip>(config, &twiddles, &traces.finalize(), None);
+
+        let ext = ExtensionComponent::multiplicity128();
+        let component_trace = ext.generate_component_trace(
+            128u32.trailing_zeros(),
+            program_trace_ref,
+            &mut side_note,
+        );
+        let (_, claimed_sum_2) =
+            ext.generate_interaction_trace(component_trace, &side_note, &lookup_elements);
+        assert_ne!(claimed_sum + claimed_sum_2, SecureField::zero());
+    }
 }

prover/src/column.rs

@@ -474,6 +474,11 @@ pub enum Column {
     /// The starting address of the read-write memory access
     #[size = 4]
     RamBaseAddr,
+    /// Quotient of the least-significant byte of `RamBaseAddr` by the required
+    /// alignment. Used only by load/store alignment AIR constraints: divided by
+    /// 2 on halfword accesses and by 4 on word accesses.
+    #[size = 1]
+    RamBaseAddrAlignmentQuotient,
     /// The new value of the read-write memory at RamBaseAddr, if accessed
     #[size = 1]
     Ram1ValCur,