fix: implement store_dw, fix load_dw, add roundtrip test

codegen/masm/intrinsics/mem.masm

@@ -286,7 +286,7 @@ export.realign_dw # [chunk_hi, chunk_mid, chunk_lo, offset]
     dup.3 u32shl # [x_hi_hi, chunk_mid, chunk__lo, offset]
 
     # Move the value below the other chunks temporarily
-    movdn.2 # [chunk_mid, chunk_lo, x_hi_hi, offset]
+    movdn.3 # [chunk_mid, chunk_lo, offset, x_hi_hi]
 
     # We must split the middle chunk into two parts,
     # one containing the bits to be combined with the
@@ -301,7 +301,7 @@ export.realign_dw # [chunk_hi, chunk_mid, chunk_lo, offset]
     # Then, we shift the chunk right by 32 - offset bits,
     # re-aligning the low bits of the first word, and
     # isolating them.
-    push.32 dup.4 u32shr  # [x_hi_lo, chunk_mid, chunk_lo, offset, x_hi_hi]
+    push.32 dup.4 u32wrapping_sub u32shr  # [x_hi_lo, chunk_mid, chunk_lo, offset, x_hi_hi]
 
     # Move the high bits back to the top
     #
@@ -327,7 +327,7 @@ export.realign_dw # [chunk_hi, chunk_mid, chunk_lo, offset]
     swap.1
 
     # Shift the value right, as done previously for the middle chunk
-    push.32 movup.4 u32shr # [x_lo_lo, x_lo_hi, x_hi]
+    push.32 movup.4 u32wrapping_sub u32shr # [x_lo_lo, x_lo_hi, x_hi]
 
     # OR the two halves together, giving us our second word, `x_lo`
     u32or # [x_lo, x_hi]
@@ -336,46 +336,67 @@ export.realign_dw # [chunk_hi, chunk_mid, chunk_lo, offset]
     swap.1 # [x_hi, x_lo]
 end
 
+# Shift a double-word (64-bit) value by the given offset
+export.offset_dw # [value_hi, value_lo, offset]
+    dup.0
+    dup.3 u32shr # [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]
+    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]
+    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]
+    swap.2 # [ chunk_hi, chunk_mid, chunk_lo]
+end
+
 
 # Load a pair of machine words (32-bit elements) to the operand stack
 export.load_dw # [waddr, index, offset]
     # check for alignment and offset validity
     dup.2 eq.0
     dup.3 push.8 u32lt assert # offset must be < 8
+    # convert offset from bytes to bits 
+    movup.3 push.8 u32wrapping_mul movdn.3 # [waddr, index, offset, value_hi, value_lo]
     # if the pointer is naturally aligned..
     if.true
         # drop byte offset
-        movup.2 drop
+        movup.2 drop # [waddr, index]
         # check which element to start at
         dup.1 eq.0
         if.true
             # drop index
-            swap.1 drop
+            swap.1 drop # [waddr]
             # load first two elements
-            padw movup.4 mem_loadw
+            padw movup.4 mem_loadw # [w0, w1, w2, w3]
             # drop last two elements, and we're done
-            movup.4 movup.4 drop drop
+            movup.3 movup.3 drop drop # [w0, w1]
         else
           dup.1 eq.1
           if.true
               # drop index
-              swap.1 drop
+              swap.1 drop # [waddr]
               # load second and third elements
-              padw movup.4 mem_loadw
+              padw movup.4 mem_loadw # [w0, w1, w2, w3]
               # drop unused elements, and we're done
-              drop movup.3 drop
+              drop movup.2 drop # [w1, w2]
           else
               swap.1 eq.2
               if.true
                 # load third and fourth elements, drop unused, and we're done
                 padw movup.4 mem_loadw drop drop
               else
                 # load first element of next word, drop the rest
-                dup.0 u32overflowing_add.1 assertz padw movup.4 mem_loadw
-                movup.4 movup.4 movup.4
-                drop drop drop
+                dup.0 u32overflowing_add.1 assertz padw movup.4 # [waddr + 1, 0, 0, 0, 0, waddr]
+                mem_loadw # [w0, w1, w2, w3, waddr]
+                movup.3 movup.3 movup.3 # [w1, w2, w3, w0, waddr]
+                drop drop drop # [w0, waddr]
                 # load fourth element, and we're done
-                movup.4 padw movup.4 mem_loadw drop drop drop
+                swap.1 padw movup.4 # [waddr, 0, 0, 0, 0, w0]
+                mem_loadw drop drop drop
               end
           end
         end
@@ -386,7 +407,7 @@ export.load_dw # [waddr, index, offset]
             # drop the index
             swap.1 drop
             # load three elements containing the double-word on the stack
-            padw movup.4 mem_loadw movup.4 drop
+            padw movup.4 mem_loadw movup.3 drop
             # re-align it, and we're done; realign_dw gets [w0, w1, w2, offset]
             exec.realign_dw
         else
@@ -660,8 +681,69 @@ end
 # start of the data; an element index, which indicates which element of
 # the word the data starts in; and a byte offset, which indicates which
 # byte is the start of the data.
-export.store_dw # [waddr, index, offset, value]
-    # TODO: implement
-    # cleanup the operand stack
-    dropw
+export.store_dw # [waddr, index, offset, value_hi, value_lo]
+    # check for alignment and offset validity
+    dup.2 eq.0
+    dup.3 push.8 u32lt assert # offset must be < 8
+    # convert offset from bytes to bits 
+    movup.3 push.8 u32wrapping_mul movdn.3 # [waddr, index, offset, value_hi, value_lo]
+    # if the pointer is naturally aligned..
+    if.true
+        # drop byte offset
+        movup.2 drop # [waddr, index, value_hi, value_lo]
+        # check which element to start at
+        dup.1 eq.0
+        if.true
+            # drop index
+            swap.1 drop # [waddr, value_hi, value_lo]
+            push.0 movdn.3 push.0 movdn.3 # [waddr, value_hi, value_lo, 0, 0]
+            mem_storew 
+            # cleanup the operand stack
+            dropw
+        else
+            dup.1 eq.1
+            if.true
+                # drop index
+                swap.1 drop # [waddr, value_hi, value_lo]
+                # store as the second and third elements of the word
+                push.0 swap.1 push.0 movdn.4 # [waddr, 0, value_hi, value_lo, 0]
+                mem_storew 
+                # cleanup the operand stack
+                dropw
+            else
+                swap.1 eq.2
+                if.true
+                    # store as the third and fourth elements of the word
+                    push.0 swap.1 push.0 swap.1 # [waddr, 0, 0, value_hi, value_lo]
+                    mem_storew 
+                    # cleanup the operand stack
+                    dropw
+                else
+                    # store the first element of the next word
+                    dup.0 u32overflowing_add.1 assertz # [waddr + 1, waddr, value_hi, value_lo]
+                    padw drop movup.6 movup.4 # [waddr + 1, value_lo, 0, 0, 0, waddr, value_hi]
+                    mem_storew dropw # [waddr, value_hi]
+                    # store the forth element
+                    padw drop movup.3 # [waddr, 0, 0, 0, value_hi]
+                    mem_storew dropw
+                end
+            end
+        end
+    else # unaligned; an unaligned double-word spans three elements
+        # check if we start in the first element
+        dup.1 eq.0
+        if.true
+            # drop the index
+            swap.1 drop # [waddr, offset, value_hi, value_lo]
+            swap.3 # [value_lo, offset, value_hi, waddr]
+            movup.2 # [value_hi value_lo, offset, waddr]
+            exec.offset_dw # [chunk_hi, chunk_mid, chunk_lo, waddr]
+            push.0 swap.4 # [waddr, chunk_hi, chunk_mid, chunk_lo, 0]
+            mem_storew
+            dropw
+        else
+            # TODO: implement
+            push.1 assertz
+        end 
+    end  
 end

codegen/masm/src/tests.rs

@@ -1,5 +1,5 @@
 #![allow(unused_imports)]
-use std::{cell::RefCell, sync::Arc};
+use std::sync::Arc;
 
 use midenc_hir::{
     self as hir,
@@ -14,6 +14,9 @@ use smallvec::{smallvec, SmallVec};
 
 use super::*;
 
+const MEMORY_SIZE_BYTES: u32 = 1048576 * 2; // Twice the size of the default Rust shadow stack size
+const MEMORY_SIZE_VM_WORDS: u32 = MEMORY_SIZE_BYTES / 16;
+
 #[cfg(test)]
 #[allow(unused_macros)]
 macro_rules! assert_masm_output {
@@ -496,8 +499,6 @@ fn i32_checked_neg() {
 
 #[test]
 fn codegen_mem_store_sw_load_sw() {
-    const MEMORY_SIZE_BYTES: u32 = 1048576 * 2; // Twice the size of the default Rust shadow stack size
-    const MEMORY_SIZE_VM_WORDS: u32 = MEMORY_SIZE_BYTES / 16;
     let context = TestContext::default();
     let mut builder = ProgramBuilder::new(&context.session.diagnostics);
     let mut mb = builder.module("test");
@@ -536,7 +537,7 @@ fn codegen_mem_store_sw_load_sw() {
 
     // eprintln!("{}", program);
 
-    fn test(program: Arc<Program>, ptr: u32, value: u32) -> u32 {
+    fn roundtrip(program: Arc<Program>, ptr: u32, value: u32) -> u32 {
         eprintln!("---------------------------------");
         eprintln!("testing store_sw/load_sw ptr: {ptr}, value: {value}");
         eprintln!("---------------------------------");
@@ -554,13 +555,117 @@ fn codegen_mem_store_sw_load_sw() {
 
     TestRunner::new(Config::with_cases(1024))
         .run(&(0u32..MEMORY_SIZE_BYTES - 4, any::<u32>()), move |(ptr, value)| {
-            let out = test(program.clone(), ptr, value);
+            let out = roundtrip(program.clone(), ptr, value);
             prop_assert_eq!(out, value);
             Ok(())
         })
         .unwrap();
 }
 
+#[test]
+fn codegen_mem_store_dw_load_dw() {
+    let context = TestContext::default();
+    let mut builder = ProgramBuilder::new(&context.session.diagnostics);
+    let mut mb = builder.module("test");
+    let id = {
+        let mut fb = mb
+            .function(
+                "store_load_dw",
+                Signature::new(
+                    [AbiParam::new(Type::U32), AbiParam::new(Type::U64)],
+                    [AbiParam::new(Type::U32)],
+                ),
+            )
+            .expect("unexpected symbol conflict");
+        let entry = fb.current_block();
+        let (ptr_u32, value) = {
+            let args = fb.block_params(entry);
+            (args[0], args[1])
+        };
+        let ptr = fb.ins().inttoptr(ptr_u32, Type::Ptr(Type::U64.into()), SourceSpan::UNKNOWN);
+        fb.ins().store(ptr, value, SourceSpan::UNKNOWN);
+        let loaded_value = fb.ins().load(ptr, SourceSpan::UNKNOWN);
+        fb.ins().ret(Some(loaded_value), SourceSpan::UNKNOWN);
+        fb.build().expect("unexpected error building function")
+    };
+
+    mb.build().expect("unexpected error constructing test module");
+
+    let program = builder.with_entrypoint(id).link().expect("failed to link program");
+
+    let ir_module = program
+        .modules()
+        .iter()
+        .take(1)
+        .collect::<Vec<&midenc_hir::Module>>()
+        .first()
+        .expect("no module in IR program")
+        .to_string();
+
+    eprintln!("{}", ir_module.as_str());
+
+    let mut compiler = MasmCompiler::new(&context.session);
+    let program = compiler
+        .compile(program)
+        .expect("compilation failed")
+        .unwrap_executable()
+        .freeze();
+
+    eprintln!("{}", program);
+
+    fn roundtrip(program: Arc<Program>, ptr: u32, value: u64) -> u64 {
+        eprintln!("---------------------------------");
+        eprintln!("testing store_dw/load_dw ptr: {ptr}, value: {value}");
+        eprintln!("---------------------------------");
+        let mut harness = TestByEmulationHarness::with_emulator_config(
+            MEMORY_SIZE_VM_WORDS as usize,
+            Emulator::DEFAULT_HEAP_START as usize,
+            Emulator::DEFAULT_LOCALS_START as usize,
+            true,
+        );
+        let mut args: SmallVec<[Felt; 4]> = smallvec!(Felt::new(ptr as u64));
+        args.extend(value.canonicalize());
+        let mut stack = harness.execute_program(program.clone(), &args).expect("execution failed");
+        u64::from_stack(&mut stack)
+    }
+
+    fn test(program: Arc<Program>, ptr: u32, value: u64) {
+        let out = roundtrip(program.clone(), ptr, value);
+        assert_eq!(out, value);
+    }
+
+    // test zero address
+    test(program.clone(), 0, 42);
+    // test first and second element
+    test(program.clone(), 16, 42);
+    test(program.clone(), 16, u64::MAX);
+    // test second and third element
+    test(program.clone(), 20, 42);
+    test(program.clone(), 20, u64::MAX);
+    // test third and fourth element
+    test(program.clone(), 24, 42);
+    test(program.clone(), 24, u64::MAX);
+    // test fourth element and first element of the next word
+    test(program.clone(), 28, 42);
+    test(program.clone(), 28, u64::MAX);
+
+    // unaligned
+    test(program.clone(), 33, 42);
+    test(program.clone(), 33, u64::MAX);
+    test(program.clone(), 34, 42);
+    test(program.clone(), 35, 42);
+    test(program.clone(), 36, 42);
+    test(program.clone(), 37, 42);
+
+    // TestRunner::new(Config::with_cases(1024))
+    //     .run(&(0u32..MEMORY_SIZE_BYTES - 4, any::<u64>()), move |(ptr, value)| {
+    //         let out = test(program.clone(), ptr, value);
+    //         prop_assert_eq!(out, value);
+    //         Ok(())
+    //     })
+    //     .unwrap();
+}
+
 #[allow(unused)]
 macro_rules! proptest_unary_numeric_op {
     ($ty_name:ident :: $op:ident, $ty:ty => $ret:ty, $rust_op:ident) => {