fix: fix errors in intrinsics::mem::replace_element,

`intrinsics::mem::store_felt_unchecked`,
`intrinsics::mem::load_sw` #230

- assert in `intrinsics::mem::replace_element`,
- storing not in the word address in `intrinsics::mem::store_felt_unchecked`,
-loading from not from a pointer, but a value(partial) in
`intrinsics::mem::load_sw` causing "out of bounds memory access".
Add test for `mem::intrinsics::store_sw/load_sw`

codegen/masm/intrinsics/mem.masm

@@ -126,7 +126,7 @@ export.load_sw # [waddr, index, offset]
                     dup.0  # [waddr, waddr, offset]
                     u32overflowing_add.1 assertz # [waddr + 1, waddr, offset]
                     # load the word and drop the unused elements
-                    padw movup.4 mem_loadw movdn.4 drop drop drop # [w0, waddr, offset]
+                    padw movup.4 mem_loadw movdn.3 drop drop drop # [w0, waddr, offset]
                     # shift the low bits
                     push.32 dup.3  # [offset, 32, w0, waddr, offset]
                     u32overflowing_sub assertz # [32 - offset, w0, waddr, offset]
@@ -364,12 +364,11 @@ proc.store_felt_unchecked # [waddr, index, value]
     mem_loadw  # [w0, w1, w2, w3, waddr, index, value]
 
     # rewrite the desired element
-    movup.6
-    movup.5
-    exec.replace_element
+    movup.6 # [value, w0, w1, w2, w3, waddr, index]
+    movup.6 # [index, value, w0, w1, w2, w3, waddr]
+    exec.replace_element # [w0', w1', w2', w3', waddr]
 
     # store the updated word
-    movup.4
     mem_storew
     dropw
 end

codegen/masm/src/tests.rs

@@ -6,6 +6,7 @@ use midenc_hir::{
     AbiParam, CallConv, Felt, FieldElement, FunctionIdent, Immediate, InstBuilder, Linkage,
     OperandStack, ProgramBuilder, Signature, SourceSpan, Stack, Type,
 };
+use prop::test_runner::{Config, TestRunner};
 use proptest::prelude::*;
 use smallvec::{smallvec, SmallVec};
 
@@ -208,6 +209,10 @@ impl TestByEmulationHarness {
     pub fn step_over(&mut self) -> Result<EmulatorEvent, EmulationError> {
         self.emulator.step_over()
     }
+
+    fn reset(&mut self) {
+        self.emulator.reset();
+    }
 }
 
 #[test]
@@ -453,6 +458,69 @@ fn i32_checked_neg() {
     harness.invoke(neg, &[min]).expect("execution failed");
 }
 
+#[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");
+    let id = {
+        let mut fb = mb
+            .function(
+                "store_load_sw",
+                Signature::new(
+                    [AbiParam::new(Type::U32), AbiParam::new(Type::U32)],
+                    [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::U32.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 mut compiler = MasmCompiler::new(&context.session);
+    let program = compiler.compile(program).expect("compilation failed").freeze();
+
+    // eprintln!("{}", program);
+
+    fn test(program: Arc<Program>, ptr: u32, value: u32) -> u32 {
+        eprintln!("---------------------------------");
+        eprintln!("testing store_sw/load_sw 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 stack = harness
+            .execute_program(program.clone(), &[Felt::new(ptr as u64), Felt::new(value as u64)])
+            .expect("execution failed");
+        stack.pop().unwrap().as_int() as u32
+    }
+
+    TestRunner::new(Config::with_cases(1024))
+        .run(&(0u32..MEMORY_SIZE_BYTES - 4, any::<u32>()), move |(ptr, value)| {
+            let out = test(program.clone(), ptr, value);
+            prop_assert_eq!(out, value);
+            Ok(())
+        })
+        .unwrap();
+}
+
 macro_rules! proptest_unary_numeric_op {
     ($ty_name:ident :: $op:ident, $ty:ty => $ret:ty, $rust_op:ident) => {
         proptest_unary_numeric_op_impl!($ty_name :: $op, $ty => $ret, $rust_op, 0..$ty_name::MAX);