feat: rv64 more cleanup (#2752)

876pol · shuklaayush · commit 56e4f86a4d23 · 2026-04-27T16:04:23.000+02:00
diff --git a/extensions/deferral/circuit/src/call/execution.rs b/extensions/deferral/circuit/src/call/execution.rs
@@ -22,7 +22,7 @@ use crate::{
     poseidon2::deferral_poseidon2_chip,
     utils::{
         byte_commit_to_f, combine_output, join_memory_ops, memory_op_chunk, COMMIT_MEMORY_OPS,
-        COMMIT_NUM_BYTES, DIGEST_MEMORY_OPS, OUTPUT_TOTAL_MEMORY_OPS,
+        COMMIT_NUM_BYTES, DIGEST_MEMORY_OPS, OUTPUT_LEN_NUM_BYTES, OUTPUT_TOTAL_MEMORY_OPS,
     },
     DeferralFn, CALL_AIR_REL_IDX, POSEIDON2_AIR_REL_IDX,
 };
@@ -204,8 +204,13 @@ unsafe fn execute_e12_impl<F: VmField, CTX: ExecutionCtxTrait>(
     let output_f_commit =
         byte_commit_to_f(&output_commit.iter().map(|v| F::from_u8(*v)).collect_vec());
 
-    // (output_commit, output_len) pair, corresponds to guest struct OutputKey
-    let output_key = combine_output(output_commit, output_len.to_le_bytes());
+    let output_len_u32 =
+        u32::try_from(output_len).expect("deferral output length should fit in a u32");
+    let mut output_len_full = [0u8; OUTPUT_LEN_NUM_BYTES];
+    output_len_full[..4].copy_from_slice(&output_len_u32.to_le_bytes());
+
+    // (output_commit, output_len) pair, corresponds to guest struct OutputKey.
+    let output_key = combine_output(output_commit, output_len_full);
 
     let new_input_acc = poseidon2_chip.perm(&old_input_acc, &input_commit, true);
     let new_output_acc = poseidon2_chip.perm(&old_output_acc, &output_f_commit, true);
diff --git a/extensions/deferral/circuit/src/output/execution.rs b/extensions/deferral/circuit/src/output/execution.rs
@@ -163,7 +163,7 @@ unsafe fn execute_e12_impl<F: PrimeField32, CTX: ExecutionCtxTrait>(
     let output_key: [u8; OUTPUT_TOTAL_BYTES] = join_memory_ops(output_key_chunks);
     let (output_commit, output_len) = split_output(output_key);
 
-    let output_len_val = u64::from_le_bytes(output_len) as usize;
+    let output_len_val = rv64_bytes_to_u32(output_len) as usize;
 
     // Bytes are sponge-hashed and constrained against output_commit. The
     // sponge rate is DIGEST_SIZE.
diff --git a/extensions/deferral/circuit/src/output/trace.rs b/extensions/deferral/circuit/src/output/trace.rs
@@ -189,7 +189,7 @@ where
         let output_key: [u8; OUTPUT_TOTAL_BYTES] = join_memory_ops(output_key_chunks);
         let (output_commit, output_len) = split_output(output_key);
 
-        let output_len_val = u64::from_le_bytes(output_len) as usize;
+        let output_len_val = rv64_bytes_to_u32(output_len) as usize;
         let num_rows = output_len_val / DIGEST_SIZE + 1;
         debug_assert!(output_len_val.is_multiple_of(DIGEST_SIZE));
 
diff --git a/extensions/riscv-adapters/src/eq_mod.rs b/extensions/riscv-adapters/src/eq_mod.rs
@@ -29,17 +29,15 @@ use openvm_instructions::{
     riscv::{RV64_MEMORY_AS, RV64_REGISTER_AS, RV64_WORD_NUM_LIMBS},
 };
 use openvm_riscv_circuit::adapters::{
-    abstract_compose, expand_to_rv64_register, tracing_read, tracing_write, RV64_CELL_BITS,
-    RV64_REGISTER_NUM_LIMBS,
+    abstract_compose, expand_to_rv64_register, tracing_read, tracing_read_reg_ptr, tracing_write,
+    RV64_CELL_BITS, RV64_REGISTER_NUM_LIMBS,
 };
 use openvm_stark_backend::{
     interaction::InteractionBuilder,
     p3_air::BaseAir,
     p3_field::{Field, PrimeCharacteristicRing, PrimeField32},
 };
 
-use crate::helpers::tracing_read_reg_ptr;
-
 /// This adapter reads from NUM_READS <= 2 pointers and writes to a register.
 /// * The data is read from the heap (address space 2), and the pointers are read from registers
 ///   (address space 1).
@@ -329,6 +327,7 @@ where
                 memory,
                 record.rs_ptr[i],
                 &mut record.rs_read_aux[i].prev_timestamp,
+                self.pointer_max_bits,
             )
         });
 
diff --git a/extensions/riscv-adapters/src/helpers.rs b/extensions/riscv-adapters/src/helpers.rs
diff --git a/extensions/riscv-adapters/src/lib.rs b/extensions/riscv-adapters/src/lib.rs
@@ -1,5 +1,4 @@
 mod eq_mod;
-mod helpers;
 mod vec_heap;
 mod vec_heap_branch;
 mod vec_to_flat;
diff --git a/extensions/riscv-adapters/src/vec_heap.rs b/extensions/riscv-adapters/src/vec_heap.rs
@@ -30,16 +30,15 @@ use openvm_instructions::{
     riscv::{RV64_MEMORY_AS, RV64_REGISTER_AS, RV64_WORD_NUM_LIMBS},
 };
 use openvm_riscv_circuit::adapters::{
-    abstract_compose, expand_to_rv64_register, tracing_read, tracing_write, RV64_CELL_BITS,
+    abstract_compose, expand_to_rv64_register, tracing_read, tracing_read_reg_ptr, tracing_write,
+    RV64_CELL_BITS,
 };
 use openvm_stark_backend::{
     interaction::InteractionBuilder,
     p3_air::BaseAir,
     p3_field::{Field, PrimeCharacteristicRing, PrimeField32},
 };
 
-use crate::helpers::tracing_read_reg_ptr;
-
 /// This adapter reads from R (R <= 2) pointers and writes to 1 pointer.
 /// * The data is read from the heap (address space 2), and the pointers are read from registers
 ///   (address space 1).
@@ -376,6 +375,7 @@ impl<
                 memory,
                 record.rs_ptrs[i],
                 &mut record.rs_read_aux[i].prev_timestamp,
+                self.pointer_max_bits,
             )
         });
 
@@ -384,6 +384,7 @@ impl<
             memory,
             record.rd_ptr,
             &mut record.rd_read_aux.prev_timestamp,
+            self.pointer_max_bits,
         );
 
         // Read memory values
diff --git a/extensions/riscv-adapters/src/vec_heap_branch.rs b/extensions/riscv-adapters/src/vec_heap_branch.rs
@@ -27,16 +27,14 @@ use openvm_instructions::{
     riscv::{RV64_MEMORY_AS, RV64_REGISTER_AS, RV64_WORD_NUM_LIMBS},
 };
 use openvm_riscv_circuit::adapters::{
-    abstract_compose, expand_to_rv64_register, tracing_read, RV64_CELL_BITS,
+    abstract_compose, expand_to_rv64_register, tracing_read, tracing_read_reg_ptr, RV64_CELL_BITS,
 };
 use openvm_stark_backend::{
     interaction::InteractionBuilder,
     p3_air::BaseAir,
     p3_field::{Field, PrimeCharacteristicRing, PrimeField32},
 };
 
-use crate::helpers::tracing_read_reg_ptr;
-
 /// This adapter reads from NUM_READS <= 2 pointers (for branch operations).
 /// * The data is read from the heap (address space 2), and the pointers are read from registers
 ///   (address space 1).
@@ -273,6 +271,7 @@ impl<
                 memory,
                 record.rs_ptrs[i],
                 &mut record.rs_read_aux[i].prev_timestamp,
+                self.pointer_max_bits,
             )
         });
 
diff --git a/extensions/riscv/circuit/src/adapters/loadstore.rs b/extensions/riscv/circuit/src/adapters/loadstore.rs
@@ -39,8 +39,8 @@ use openvm_stark_backend::{
 
 use super::{RV64_REGISTER_NUM_LIMBS, RV64_WORD_NUM_LIMBS};
 use crate::adapters::{
-    expand_to_rv64_register, memory_read, memory_read_deferral, timed_write, timed_write_deferral,
-    tracing_read, RV64_CELL_BITS,
+    expand_to_rv64_register, memory_read, memory_read_deferral, rv64_bytes_to_u32, timed_write,
+    timed_write_deferral, tracing_read, RV64_CELL_BITS,
 };
 
 /// LoadStore Adapter handles all memory and register operations, so it must be aware
@@ -371,18 +371,12 @@ where
         );
 
         record.rs1_ptr = b.as_canonical_u32();
-        let rs1_u64 = u64::from_le_bytes(tracing_read(
+        record.rs1_val = rv64_bytes_to_u32(tracing_read(
             memory,
             RV64_REGISTER_AS,
             record.rs1_ptr,
             &mut record.rs1_aux_record.prev_timestamp,
         ));
-        assert_eq!(
-            rs1_u64 >> 32,
-            0,
-            "rs1 upper 4 bytes must be zero for load/store adapter"
-        );
-        record.rs1_val = rs1_u64 as u32;
         let rs1_val = record.rs1_val;
 
         record.imm = c.as_canonical_u32() as u16;
@@ -393,12 +387,7 @@ where
         let shift_amount = ptr_val & (RV64_REGISTER_NUM_LIMBS as u32 - 1);
         let ptr_val = ptr_val - shift_amount;
 
-        assert!(
-            ptr_val < (1 << self.pointer_max_bits),
-            "ptr_val: {ptr_val} = rs1_val: {} + imm_extended: {imm_extended} >= 2 ** {}",
-            rs1_val,
-            self.pointer_max_bits
-        );
+        debug_assert!((ptr_val as u64) < (1u64 << self.pointer_max_bits));
 
         // prev_data: We need to keep values of some cells to keep them unchanged when writing to
         // those cells
diff --git a/extensions/riscv/circuit/src/adapters/mod.rs b/extensions/riscv/circuit/src/adapters/mod.rs
@@ -218,6 +218,22 @@ pub fn tracing_read<const N: usize>(
     data
 }
 
+/// Reads an RV64 register, records the memory access, and returns the low 32 bits. Debug-asserts
+/// the returned value fits in `pointer_max_bits` (which, for `pointer_max_bits <= 32`, also
+/// implies the upper 32 bits are zero).
+#[inline(always)]
+pub fn tracing_read_reg_ptr(
+    memory: &mut TracingMemory,
+    ptr: u32,
+    prev_timestamp: &mut u32,
+    pointer_max_bits: usize,
+) -> u32 {
+    let bytes = tracing_read(memory, RV64_REGISTER_AS, ptr, prev_timestamp);
+    let val = rv64_bytes_to_u32(bytes);
+    debug_assert!((val as u64) < (1u64 << pointer_max_bits));
+    val
+}
+
 #[inline(always)]
 pub fn tracing_read_imm(
     memory: &mut TracingMemory,
diff --git a/extensions/riscv/circuit/src/extension/mod.rs b/extensions/riscv/circuit/src/extension/mod.rs
@@ -842,9 +842,7 @@ mod phantom {
     use openvm_stark_backend::p3_field::{Field, PrimeField32};
     use rand::{rngs::StdRng, Rng};
 
-    use crate::adapters::{
-        memory_read, read_rv64_register, read_rv64_register_as_u32, RV64_REGISTER_NUM_LIMBS,
-    };
+    use crate::adapters::{memory_read, read_rv64_register_as_u32, RV64_REGISTER_NUM_LIMBS};
 
     const HINT_DWORD_BYTES: usize = RV64_REGISTER_NUM_LIMBS;
 
@@ -898,12 +896,11 @@ mod phantom {
                 eprintln!("WARNING: Using fixed-seed RNG for deterministic randomness. Consider security implications for your use case.");
             });
 
-            let len = read_rv64_register(memory, a) as usize;
+            let byte_len = read_rv64_register_as_u32(memory, a) as usize * HINT_DWORD_BYTES;
             streams.hint_stream.clear();
-            streams.hint_stream.extend(
-                std::iter::repeat_with(|| F::from_u8(rng.random::<u8>()))
-                    .take(len * HINT_DWORD_BYTES),
-            );
+            streams
+                .hint_stream
+                .extend(std::iter::repeat_with(|| F::from_u8(rng.random::<u8>())).take(byte_len));
             Ok(())
         }
     }
diff --git a/extensions/riscv/circuit/src/hintstore/execution.rs b/extensions/riscv/circuit/src/hintstore/execution.rs
@@ -19,6 +19,7 @@ use openvm_riscv_transpiler::{
 use openvm_stark_backend::p3_field::PrimeField32;
 
 use super::Rv64HintStoreExecutor;
+use crate::adapters::rv64_bytes_to_u32;
 
 #[derive(AlignedBytesBorrow, Clone)]
 #[repr(C)]
@@ -165,26 +166,14 @@ unsafe fn execute_e12_impl<F: PrimeField32, CTX: ExecutionCtxTrait, const IS_HIN
     let pc = exec_state.pc();
     let mem_ptr_limbs =
         exec_state.vm_read::<u8, RV64_REGISTER_NUM_LIMBS>(RV64_REGISTER_AS, pre_compute.b as u32);
-    let mem_ptr_u64 = u64::from_le_bytes(mem_ptr_limbs);
-    assert_eq!(
-        mem_ptr_u64 >> 32,
-        0,
-        "mem_ptr upper 4 bytes must be zero for hintstore"
-    );
-    let mem_ptr = mem_ptr_u64 as u32;
+    let mem_ptr = rv64_bytes_to_u32(mem_ptr_limbs);
 
     let num_words = if IS_HINT_STORED {
         1
     } else {
         let num_words_limbs = exec_state
             .vm_read::<u8, RV64_REGISTER_NUM_LIMBS>(RV64_REGISTER_AS, pre_compute.a as u32);
-        let num_words_u64 = u64::from_le_bytes(num_words_limbs);
-        assert_eq!(
-            num_words_u64 >> 32,
-            0,
-            "num_words upper 4 bytes must be zero"
-        );
-        num_words_u64 as u32
+        rv64_bytes_to_u32(num_words_limbs)
     };
     debug_assert_ne!(num_words, 0);
 
diff --git a/extensions/riscv/circuit/src/hintstore/mod.rs b/extensions/riscv/circuit/src/hintstore/mod.rs
@@ -39,7 +39,10 @@ use openvm_stark_backend::{
     BaseAirWithPublicValues, PartitionedBaseAir,
 };
 
-use crate::adapters::{expand_to_rv64_register, read_rv64_register, tracing_read, tracing_write};
+use crate::adapters::{
+    expand_to_rv64_register, read_rv64_register_as_u32, tracing_read, tracing_read_reg_ptr,
+    tracing_write,
+};
 
 mod execution;
 
@@ -429,13 +432,7 @@ where
         let num_words = if local_opcode == HINT_STORED {
             1
         } else {
-            let num_words_u64 = read_rv64_register(state.memory.data(), a);
-            assert_eq!(
-                num_words_u64 >> 32,
-                0,
-                "num_words upper 4 bytes must be zero"
-            );
-            num_words_u64 as u32
+            read_rv64_register_as_u32(state.memory.data(), a)
         };
 
         // Bounds check: num_words must not exceed MAX_HINT_BUFFER_DWORDS
@@ -455,20 +452,13 @@ where
         record.inner.timestamp = state.memory.timestamp;
         record.inner.mem_ptr_ptr = b;
 
-        let mem_ptr_u64 = u64::from_le_bytes(tracing_read(
+        record.inner.mem_ptr = tracing_read_reg_ptr(
             state.memory,
-            RV64_REGISTER_AS,
             b,
             &mut record.inner.mem_ptr_aux_record.prev_timestamp,
-        ));
-        assert_eq!(
-            mem_ptr_u64 >> 32,
-            0,
-            "mem_ptr upper 4 bytes must be zero for hintstore"
+            self.pointer_max_bits,
         );
-        record.inner.mem_ptr = mem_ptr_u64 as u32;
 
-        debug_assert!(record.inner.mem_ptr <= (1 << self.pointer_max_bits));
         debug_assert_ne!(num_words, 0);
         debug_assert!(num_words <= (1 << self.pointer_max_bits));
 
diff --git a/extensions/riscv/circuit/src/hintstore/tests.rs b/extensions/riscv/circuit/src/hintstore/tests.rs
@@ -343,7 +343,7 @@ fn test_hint_buffer_mem_ptr_range_check() {
 }
 
 #[test]
-#[should_panic(expected = "mem_ptr upper 4 bytes must be zero for hintstore")]
+#[should_panic(expected = "upper 4 bytes must be zero")]
 fn test_hintstore_rs1_upper_bytes_non_zero() {
     let mut rng = create_seeded_rng();
     let mut tester = VmChipTestBuilder::default();
diff --git a/extensions/riscv/circuit/src/jalr/core.rs b/extensions/riscv/circuit/src/jalr/core.rs
@@ -27,8 +27,8 @@ use openvm_stark_backend::{
 };
 
 use crate::adapters::{
-    expand_to_rv64_register, Rv64JalrAdapterExecutor, Rv64JalrAdapterFiller, RV64_CELL_BITS,
-    RV64_REGISTER_NUM_LIMBS, RV64_WORD_NUM_LIMBS,
+    expand_to_rv64_register, rv64_bytes_to_u32, Rv64JalrAdapterExecutor, Rv64JalrAdapterFiller,
+    RV64_CELL_BITS, RV64_REGISTER_NUM_LIMBS, RV64_WORD_NUM_LIMBS,
 };
 
 #[repr(C)]
@@ -251,8 +251,7 @@ where
         let rs1_data = self
             .adapter
             .read(state.memory, instruction, &mut adapter_record);
-        core_record.rs1_val =
-            u32::from_le_bytes([rs1_data[0], rs1_data[1], rs1_data[2], rs1_data[3]]);
+        core_record.rs1_val = rv64_bytes_to_u32(rs1_data);
 
         core_record.imm = c.as_canonical_u32() as u16;
         core_record.imm_sign = g.is_one();
diff --git a/extensions/riscv/circuit/src/jalr/execution.rs b/extensions/riscv/circuit/src/jalr/execution.rs
@@ -13,6 +13,7 @@ use openvm_instructions::{
 use openvm_stark_backend::p3_field::PrimeField32;
 
 use super::core::Rv64JalrExecutor;
+use crate::adapters::rv64_bytes_to_u32;
 #[cfg(feature = "aot")]
 use crate::common::*;
 
@@ -209,11 +210,7 @@ unsafe fn execute_e12_impl<F: PrimeField32, CTX: ExecutionCtxTrait, const ENABLE
     let pc = exec_state.pc();
     let rs1 =
         exec_state.vm_read::<u8, RV64_REGISTER_NUM_LIMBS>(RV64_REGISTER_AS, pre_compute.b as u32);
-    debug_assert!(
-        rs1[4..].iter().all(|&b| b == 0),
-        "upper bytes of rs1 must be zero"
-    );
-    let rs1 = u32::from_le_bytes([rs1[0], rs1[1], rs1[2], rs1[3]]);
+    let rs1 = rv64_bytes_to_u32(rs1);
     let to_pc = rs1.wrapping_add(pre_compute.imm_extended);
     let to_pc = to_pc - (to_pc & 1);
     debug_assert!(to_pc < (1 << PC_BITS));
diff --git a/extensions/riscv/circuit/src/jalr/tests.rs b/extensions/riscv/circuit/src/jalr/tests.rs
diff --git a/extensions/riscv/circuit/src/load_sign_extend/tests.rs b/extensions/riscv/circuit/src/load_sign_extend/tests.rs
diff --git a/extensions/sha2/circuit/src/sha2_chips/trace.rs b/extensions/sha2/circuit/src/sha2_chips/trace.rs

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,4 @@`
`1`	`1`	`mod eq_mod;`
`2`		`-mod helpers;`
`3`	`2`	`mod vec_heap;`
`4`	`3`	`mod vec_heap_branch;`
`5`	`4`	`mod vec_to_flat;`