Merge pull request #718 from 0xMiden/otrho/fix-mem-store_sw

Fix a bug in `mem::store_sw` intrinsic which used incorrect masks in unaligned writes.

Author: bitwalker

codegen/masm/intrinsics/mem.masm

@@ -293,17 +293,19 @@ end
 # Returns three 32-bit chunks [chunk_lo, chunk_mid, chunk_hi]
 export.offset_dw # [value_hi, value_lo, offset]
     dup.0
-    dup.3 u32shr # [chunk_hi, value_hi, value_lo, offset]
+    dup.3 u32shl # [chunk_hi, value_hi, value_lo, offset]
+
     movdn.3      # [value_hi, value_lo, offset, chunk_hi]
     push.32 dup.3 u32wrapping_sub # [32 - offset, value_hi, value_lo, offset, chunk_hi]
-    u32shl       # [ chunk_mid_hi, value_lo, offset, chunk_hi]
+    u32shr       # [ chunk_mid_hi, value_lo, offset, chunk_hi]
     dup.1        # [ value_lo, chunk_mid_hi, value_lo, offset, chunk_hi]
     dup.3        # [ offset, value_lo, chunk_mid_hi, value_lo, offset, chunk_hi]
-    u32shr       # [ chunk_mid_lo, chunk_mid_hi, value_lo, offset, chunk_hi]
+    u32shl       # [ chunk_mid_lo, chunk_mid_hi, value_lo, offset, chunk_hi]
     u32or        # [ chunk_mid, value_lo, offset, chunk_hi]
+
     movdn.2      # [ value_lo, offset, chunk_mid, chunk_hi]
     push.32 movup.2 u32wrapping_sub # [32 - offset, value_lo, offset, chunk_mid, chunk_hi]
-    u32shl       # [ chunk_lo, chunk_mid, chunk_hi]
+    u32shr       # [ chunk_lo, chunk_mid, chunk_hi]
 end
 
 # Load a machine double-word (64-bit value, in two 32-bit chunks) to the operand stack
@@ -400,15 +402,15 @@ export.store_sw # [addr, offset, value]
         push.32 dup.4 sub # [32 - bit_offset, e0, e1, addr, bit_offset, value]
 
         # compute the masks
-        push.4294967295 dup.1 u32shl  # [mask_hi, rshift, e0, e1, addr, bit_offset, value]
-        dup.0 u32not                  # [mask_lo, mask_hi, rshift, e0, e1, addr, bit_offset, value]
+        push.4294967295 dup.1 u32shr  # [mask_lo, rshift, e0, e1, addr, bit_offset, value]
+        dup.0 u32not                  # [mask_hi, mask_lo, rshift, e0, e1, addr, bit_offset, value]
 
         # manipulate the bits of the two target elements, such that the 32-bit word we're storing
         # is placed at the correct offset from the start of the memory cell when viewing the cell
         # as a set of 4 32-bit chunks
         #
         # First, mask out the bits we're going to overwrite in the two elements we loaded
-        movup.4 u32and         # [e1_masked, mask_hi, rshift, e0, addr, bit_offset, value]
+        movup.4 u32and         # [e1_masked, mask_lo, rshift, e0, addr, bit_offset, value]
         movup.3 movup.2 u32and # [e0_masked, e1_masked, rshift, addr, bit_offset, value]
 
         # now, we need to shift/mask/split the 32-bit value into two elements, then
@@ -454,6 +456,7 @@ export.store_dw # [addr, offset, value_hi, value_lo]
         #
         # convert offset from bytes to bits
         swap.1 push.8 u32wrapping_mul swap.1   # [addr, bit_offset, value_hi, value_lo]
+
         movdn.3        # [bit_offset, value_hi, value_lo, addr]
         dup.0 movdn.4  # [bit_offset, value_hi, value_lo, addr, bit_offset]
         movdn.2        # [value_hi, value_lo, bit_offset, addr, bit_offset]
@@ -463,7 +466,7 @@ export.store_dw # [addr, offset, value_hi, value_lo]
         push.32 movup.5 sub           # [32 - bit_offset, lo, mid, hi, addr]
 
         # compute the masks
-        push.4294967295 swap.1 u32shl # [mask_hi, lo, mid, hi, addr]
+        push.4294967295 swap.1 u32shr # [mask_hi, lo, mid, hi, addr]
         dup.0 u32not                  # [mask_lo, mask_hi, lo, mid, hi, addr]
 
         # combine the bits of the lo and hi chunks with their corresponding elements, so that we

tests/integration/src/codegen/intrinsics/mem.rs

@@ -731,3 +731,231 @@ fn store_u8() {
         _ => panic!("Unexpected test result: {res:?}"),
     }
 }
+
+#[test]
+fn store_unaligned_u32() {
+    // Use the start of the 17th page (1 page after the 16 pages reserved for the Rust stack)
+    let write_to = 17 * 2u32.pow(16);
+    let write_val = 0xddccbbaa_u32; // Little-endian bytes will be [AA BB CC DD].
+
+    // Generate a `test` module with `main` function that stores to a u32 offset.
+    let signature = Signature::new(
+        [AbiParam::new(Type::U32)],
+        [AbiParam::new(Type::U32)], // Return u32 to satisfy test infrastructure
+    );
+
+    // Compile once outside the test loop
+    let (package, context) = compile_test_module(signature, |builder| {
+        let block = builder.current_block();
+        let idx_val = block.borrow().arguments()[0] as ValueRef;
+
+        // Set base pointer, add argument offset to it.
+        let base_addr = builder.u32(write_to, SourceSpan::default());
+        let write_addr = builder.add(base_addr, idx_val, SourceSpan::default()).unwrap();
+        let ptr = builder
+            .inttoptr(write_addr, Type::from(PointerType::new(Type::U32)), SourceSpan::default())
+            .unwrap();
+
+        // Store test value to pointer.
+        let write_val = builder.u32(write_val, SourceSpan::default());
+        builder.store(ptr, write_val, SourceSpan::default()).unwrap();
+
+        // Return a constant to satisfy test infrastructure
+        let result = builder.u32(1, SourceSpan::default());
+        builder.ret(Some(result), SourceSpan::default()).unwrap();
+    });
+
+    let run_test = |offs: u32, expected0: u32, expected1: u32| {
+        // Initialise memory with some known bytes.
+        let initializers = [Initializer::MemoryBytes {
+            addr: write_to,
+            bytes: &[0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88],
+        }];
+
+        let output = eval_package::<u32, _, _>(
+            &package,
+            initializers,
+            &[Felt::new(offs as u64)],
+            context.session(),
+            |trace| {
+                // Get the overwritten words.
+                let word0 = trace.read_from_rust_memory::<u32>(write_to).unwrap();
+                let word1 = trace.read_from_rust_memory::<u32>(write_to + 4).unwrap();
+
+                eprintln!("word0: 0x{word0:0>8x}");
+                eprintln!("word1: 0x{word1:0>8x}");
+
+                assert_eq!(
+                    word0, expected0,
+                    "expected 1st overwritten word to be {expected0}, got {word0}, with offset \
+                     {offs}"
+                );
+
+                assert_eq!(
+                    word1, expected1,
+                    "expected 2nd overwritten word to be {expected1}, got {word1}, with offset \
+                     {offs}"
+                );
+
+                Ok(())
+            },
+        )
+        .unwrap();
+
+        assert_eq!(output, 1);
+    };
+
+    // Overwrite 11 22 33 44 55 66 77 88 with bytes aa bb cc dd at offset 1:
+    //  Expect 11 aa bb cc | dd 66 77 88
+    //  or 0xccbbaa11 and 0x887766dd.
+    run_test(1, 0xccbbaa11, 0x887766dd);
+
+    // Overwrite 11 22 33 44 55 66 77 88 with bytes aa bb cc dd at offset 2:
+    //  Expect 11 22 aa bb | cc dd 77 88
+    //  or 0xbbaa2211 and 0x8877ddcc.
+    run_test(2, 0xbbaa2211, 0x8877ddcc);
+
+    // Overwrite 11 22 33 44 55 66 77 88 with bytes aa bb cc dd at offset 3:
+    //  Expect 11 22 33 aa | bb cc dd 88
+    //  or 0xaa332211 and 0x88ddccbb.
+    run_test(3, 0xaa332211, 0x88ddccbb);
+}
+
+#[test]
+fn store_unaligned_u64() {
+    // Use the start of the 17th page (1 page after the 16 pages reserved for the Rust stack)
+    let write_to = 17 * 2u32.pow(16);
+
+    // STORE_DW writes the high 32bit word to address and low 32bit word to address+1.
+    //   So a .store() of 0xddccbbaa_cdabffee writes 0xddccbbaa to addr and 0xcdabffee to addr+1.
+    //   Which in turn will be little-endian bytes [ AA BB CC DD EE FF AB CD ] at addr.
+    let write_val = 0xddccbbaa_cdabffee_u64;
+
+    // Generate a `test` module with `main` function that stores to a u32 offset.
+    let signature = Signature::new(
+        [AbiParam::new(Type::U32)],
+        [AbiParam::new(Type::U32)], // Return u32 to satisfy test infrastructure
+    );
+
+    // Compile once outside the test loop
+    let (package, context) = compile_test_module(signature, |builder| {
+        let block = builder.current_block();
+        let idx_val = block.borrow().arguments()[0] as ValueRef;
+
+        // Set base pointer, add argument offset to it.
+        let base_addr = builder.u32(write_to, SourceSpan::default());
+        let write_addr = builder.add(base_addr, idx_val, SourceSpan::default()).unwrap();
+        let ptr = builder
+            .inttoptr(write_addr, Type::from(PointerType::new(Type::U64)), SourceSpan::default())
+            .unwrap();
+
+        // Store test value to pointer.
+        let write_val = builder.u64(write_val, SourceSpan::default());
+        builder.store(ptr, write_val, SourceSpan::default()).unwrap();
+
+        // Return a constant to satisfy test infrastructure
+        let result = builder.u32(1, SourceSpan::default());
+        builder.ret(Some(result), SourceSpan::default()).unwrap();
+    });
+
+    let run_test = |offs: u32, expected0: u32, expected1: u32, expected2: u32, expected3: u32| {
+        // Initialise memory with some known bytes.
+        let initializers = [Initializer::MemoryBytes {
+            addr: write_to,
+            bytes: &[
+                0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16,
+                0x17, 0x18,
+            ],
+        }];
+
+        let output = eval_package::<u32, _, _>(
+            &package,
+            initializers,
+            &[Felt::new(offs as u64)],
+            context.session(),
+            |trace| {
+                // Get the overwritten words.
+                let word0 = trace.read_from_rust_memory::<u32>(write_to).unwrap();
+                let word1 = trace.read_from_rust_memory::<u32>(write_to + 4).unwrap();
+                let word2 = trace.read_from_rust_memory::<u32>(write_to + 8).unwrap();
+                let word3 = trace.read_from_rust_memory::<u32>(write_to + 12).unwrap();
+
+                eprintln!("word0: 0x{word0:0>8x}");
+                eprintln!("word1: 0x{word1:0>8x}");
+                eprintln!("word2: 0x{word2:0>8x}");
+                eprintln!("word3: 0x{word3:0>8x}");
+
+                assert_eq!(
+                    word0, expected0,
+                    "expected 1st overwritten word to be {expected0}, got {word0}, with offset \
+                     {offs}"
+                );
+
+                assert_eq!(
+                    word1, expected1,
+                    "expected 2nd overwritten word to be {expected1}, got {word1}, with offset \
+                     {offs}"
+                );
+
+                assert_eq!(
+                    word2, expected2,
+                    "expected 3rd overwritten word to be {expected2}, got {word2}, with offset \
+                     {offs}"
+                );
+
+                assert_eq!(
+                    word3, expected3,
+                    "expected 4th overwritten word to be {expected3}, got {word3}, with offset \
+                     {offs}"
+                );
+
+                Ok(())
+            },
+        )
+        .unwrap();
+
+        assert_eq!(output, 1);
+    };
+
+    // Overwrite    01 02 03 04 05 06 07 08-11 12 13 14 15 16 17 18
+    //   with bytes    aa bb cc dd ee ff ab cd at offset 1:
+    //   Expect     01 aa bb cc dd ee ff ab cd 12 13 14 15 16 17 18
+    //   or         0xccbbaa01, 0xabffeedd, 0x141312cd, 0x18171615
+    run_test(1, 0xccbbaa01, 0xabffeedd, 0x141312cd, 0x18171615);
+
+    // Overwrite    01 02 03 04 05 06 07 08-11 12 13 14 15 16 17 18
+    //   with bytes    aa bb cc dd ee ff ab cd at offset 2:
+    //   Expect     01 02 aa bb cc dd ee ff ab cd 13 14 15 16 17 18
+    //   or         0xbbaa0201, 0xffeeddcc, 0x1413cdab, 0x18171615
+    run_test(2, 0xbbaa0201, 0xffeeddcc, 0x1413cdab, 0x18171615);
+
+    // Overwrite    01 02 03 04 05 06 07 08-11 12 13 14 15 16 17 18
+    //   with bytes    aa bb cc dd ee ff ab cd at offset 3:
+    //   Expect     01 02 03 aa bb cc dd ee ff ab cd 14 15 16 17 18
+    //   or         0xaa030201, 0xeeddccbb, 0x14cdabff, 0x18171615
+    run_test(3, 0xaa030201, 0xeeddccbb, 0x14cdabff, 0x18171615);
+
+    // Overwrite    01 02 03 04 05 06 07 08-11 12 13 14 15 16 17 18
+    //   with bytes    aa bb cc dd ee ff ab cd at offset 4:
+    //   Expect     01 02 03 04 aa bb cc dd ee ff ab cd 15 16 17 18
+    //   or         0x04030201, 0xddccbbaa, 0xcdabffee, 0x18171615
+    run_test(4, 0x04030201, 0xddccbbaa, 0xcdabffee, 0x18171615);
+
+    // Overwrite    01 02 03 04 05 06 07 08-11 12 13 14 15 16 17 18
+    //   with bytes    aa bb cc dd ee ff ab cd at offset 5:
+    //   Expect     01 02 03 04 05 aa bb cc dd ee ff ab cd 16 17 18
+    //   or         0x04030201, 0xccbbaa05, 0xabffeedd, 0x181716cd
+    run_test(5, 0x04030201, 0xccbbaa05, 0xabffeedd, 0x181716cd);
+
+    // Overwrite    01 02 03 04 05 06 07 08-11 12 13 14 15 16 17 18
+    //   with bytes    aa bb cc dd ee ff ab cd at offset 6:
+    //   Expect     01 02 03 04 05 06 aa bb cc dd ee ff ab cd 17 18
+    //   or         0x04030201, 0xbbaa0605, 0xffeeddcc, 0x1817cdab
+    run_test(6, 0x04030201, 0xbbaa0605, 0xffeeddcc, 0x1817cdab);
+
+    // Overwrite    01 02 03 04 05 06 07 08-11 12 13 14 15 16 17 18
+    //   with bytes    aa bb cc dd ee ff ab cd at offset 7:
+    //   Expect     01 02 03 04 05 06 07 aa bb cc dd ee ff ab cd 18
+    //   or         0x04030201, 0xaa070605, 0xeeddccbb, 0x18cdabff
+    run_test(7, 0x04030201, 0xaa070605, 0xeeddccbb, 0x18cdabff);
+}