feat: update memory operations to use usize for better type safety

Author: Arnaud Riess

crates/herkos-runtime/src/memory.rs

@@ -150,9 +150,9 @@ impl<const MAX_PAGES: usize> IsolatedMemory<MAX_PAGES> {
     ///
     /// Only the low 8 bits of `val` are used (Wasm spec). Traps (`OutOfBounds`)
     /// if the region extends beyond the current active memory.
-    pub fn fill(&mut self, dst: u32, val: u8, len: u32) -> WasmResult<()> {
+    pub fn fill(&mut self, dst: usize, val: u8, len: usize) -> WasmResult<()> {
         let active = self.active_size();
-        fill_inner(self.flat_mut(), active, dst as usize, val, len as usize)
+        fill_inner(self.flat_mut(), active, dst, val, len)
     }
 
     /// Wasm `memory.init` — copy `len` bytes from `data[src_offset..]` into
@@ -164,13 +164,13 @@ impl<const MAX_PAGES: usize> IsolatedMemory<MAX_PAGES> {
     /// beyond active memory.
     pub fn init_data_partial(
         &mut self,
-        dst: u32,
+        dst: usize,
         data: &[u8],
         src_offset: usize,
         len: usize,
     ) -> WasmResult<()> {
         let active = self.active_size();
-        init_data_partial_inner(self.flat_mut(), active, dst as usize, data, src_offset, len)
+        init_data_partial_inner(self.flat_mut(), active, dst, data, src_offset, len)
     }
 
     // ── Bounds-checked (safe) load/store ──────────────────────────────
@@ -826,16 +826,13 @@ mod tests {
     #[test]
     fn fill_out_of_bounds() {
         let mut mem = Mem::try_new(1).unwrap();
-        assert_eq!(
-            mem.fill(PAGE_SIZE as u32 - 3, 0, 10),
-            Err(WasmTrap::OutOfBounds)
-        );
+        assert_eq!(mem.fill(PAGE_SIZE - 3, 0, 10), Err(WasmTrap::OutOfBounds));
     }
 
     #[test]
     fn fill_at_boundary() {
         let mut mem = Mem::try_new(1).unwrap();
-        assert!(mem.fill(PAGE_SIZE as u32 - 4, 0x42, 4).is_ok());
+        assert!(mem.fill(PAGE_SIZE - 4, 0x42, 4).is_ok());
         assert_eq!(mem.load_u8(PAGE_SIZE - 1).unwrap(), 0x42);
     }
 
@@ -878,7 +875,7 @@ mod tests {
     fn init_data_partial_dst_out_of_bounds() {
         let mut mem = Mem::try_new(1).unwrap();
         assert_eq!(
-            mem.init_data_partial(PAGE_SIZE as u32 - 2, b"Hello", 0, 5),
+            mem.init_data_partial(PAGE_SIZE - 2, b"Hello", 0, 5),
             Err(WasmTrap::OutOfBounds)
         );
     }

crates/herkos-tests/data/wat/i32_extend16_s.wat

@@ -0,0 +1,6 @@
+(module
+  (func (export "func_0") (param i32) (result i32)
+    local.get 0
+    i32.extend16_s
+  )
+)

crates/herkos-tests/data/wat/i32_extend8_s.wat

@@ -0,0 +1,6 @@
+(module
+  (func (export "func_0") (param i32) (result i32)
+    local.get 0
+    i32.extend8_s
+  )
+)

crates/herkos-tests/data/wat/i64_extend16_s.wat

@@ -0,0 +1,6 @@
+(module
+  (func (export "func_0") (param i64) (result i64)
+    local.get 0
+    i64.extend16_s
+  )
+)

crates/herkos-tests/data/wat/i64_extend32_s.wat

@@ -0,0 +1,6 @@
+(module
+  (func (export "func_0") (param i64) (result i64)
+    local.get 0
+    i64.extend32_s
+  )
+)

crates/herkos-tests/data/wat/i64_extend8_s.wat

@@ -0,0 +1,6 @@
+(module
+  (func (export "func_0") (param i64) (result i64)
+    local.get 0
+    i64.extend8_s
+  )
+)

crates/herkos-tests/tests/arithmetic.rs

@@ -6,7 +6,8 @@
 //! 3. The semantics match WebAssembly's wrapping arithmetic
 
 use herkos_tests::{
-    abs, add, add_i64, const_i64, const_return, factorial, fibonacci, gcd, mul, nop, sub,
+    abs, add, add_i64, const_i64, const_return, factorial, fibonacci, gcd, i32_extend16_s,
+    i32_extend8_s, i64_extend16_s, i64_extend32_s, i64_extend8_s, mul, nop, sub,
 };
 
 #[test]
@@ -307,3 +308,204 @@ fn test_add_matches_rust_wrapping() {
         );
     }
 }
+
+// Sign-extension operations
+#[test]
+fn test_i32_extend8_s() {
+    let mut extend8_mod = i32_extend8_s::new().unwrap();
+
+    // Positive value: 0x42 sign-extended to i32
+    assert_eq!(
+        extend8_mod.func_0(0x42).unwrap(),
+        0x42,
+        "i32.extend8_s(0x42) should be 0x42"
+    );
+
+    // Negative value: 0xFF sign-extended to i32 (all bits set, -1)
+    assert_eq!(
+        extend8_mod.func_0(0xFF).unwrap(),
+        -1i32,
+        "i32.extend8_s(0xFF) should be -1"
+    );
+
+    // Another negative: 0x80 sign-extended (bit 7 set)
+    assert_eq!(
+        extend8_mod.func_0(0x80).unwrap(),
+        -128i32,
+        "i32.extend8_s(0x80) should be -128"
+    );
+
+    // Zero
+    assert_eq!(
+        extend8_mod.func_0(0).unwrap(),
+        0,
+        "i32.extend8_s(0) should be 0"
+    );
+
+    // 0x7F sign-extended (max positive signed byte)
+    assert_eq!(
+        extend8_mod.func_0(0x7F).unwrap(),
+        127i32,
+        "i32.extend8_s(0x7F) should be 127"
+    );
+}
+
+#[test]
+fn test_i32_extend16_s() {
+    let mut extend16_mod = i32_extend16_s::new().unwrap();
+
+    // Positive value: 0x1234 sign-extended
+    assert_eq!(
+        extend16_mod.func_0(0x1234).unwrap(),
+        0x1234,
+        "i32.extend16_s(0x1234) should be 0x1234"
+    );
+
+    // Negative value: 0xFFFF sign-extended (all bits set, -1)
+    assert_eq!(
+        extend16_mod.func_0(0xFFFF).unwrap(),
+        -1i32,
+        "i32.extend16_s(0xFFFF) should be -1"
+    );
+
+    // Another negative: 0x8000 sign-extended (bit 15 set)
+    assert_eq!(
+        extend16_mod.func_0(0x8000).unwrap(),
+        -32768i32,
+        "i32.extend16_s(0x8000) should be -32768"
+    );
+
+    // Zero
+    assert_eq!(
+        extend16_mod.func_0(0).unwrap(),
+        0,
+        "i32.extend16_s(0) should be 0"
+    );
+
+    // 0x7FFF sign-extended (max positive signed short)
+    assert_eq!(
+        extend16_mod.func_0(0x7FFF).unwrap(),
+        32767i32,
+        "i32.extend16_s(0x7FFF) should be 32767"
+    );
+}
+
+#[test]
+fn test_i64_extend8_s() {
+    let mut extend8_mod = i64_extend8_s::new().unwrap();
+
+    // Positive value
+    assert_eq!(
+        extend8_mod.func_0(0x42).unwrap(),
+        0x42i64,
+        "i64.extend8_s(0x42) should be 0x42"
+    );
+
+    // Negative value: 0xFF sign-extended to i64
+    assert_eq!(
+        extend8_mod.func_0(0xFF).unwrap(),
+        -1i64,
+        "i64.extend8_s(0xFF) should be -1"
+    );
+
+    // Another negative: 0x80 sign-extended
+    assert_eq!(
+        extend8_mod.func_0(0x80).unwrap(),
+        -128i64,
+        "i64.extend8_s(0x80) should be -128"
+    );
+
+    // Zero
+    assert_eq!(
+        extend8_mod.func_0(0).unwrap(),
+        0i64,
+        "i64.extend8_s(0) should be 0"
+    );
+
+    // Max positive signed byte
+    assert_eq!(
+        extend8_mod.func_0(0x7F).unwrap(),
+        127i64,
+        "i64.extend8_s(0x7F) should be 127"
+    );
+}
+
+#[test]
+fn test_i64_extend16_s() {
+    let mut extend16_mod = i64_extend16_s::new().unwrap();
+
+    // Positive value
+    assert_eq!(
+        extend16_mod.func_0(0x1234).unwrap(),
+        0x1234i64,
+        "i64.extend16_s(0x1234) should be 0x1234"
+    );
+
+    // Negative value: 0xFFFF sign-extended to i64
+    assert_eq!(
+        extend16_mod.func_0(0xFFFF).unwrap(),
+        -1i64,
+        "i64.extend16_s(0xFFFF) should be -1"
+    );
+
+    // Another negative: 0x8000 sign-extended
+    assert_eq!(
+        extend16_mod.func_0(0x8000).unwrap(),
+        -32768i64,
+        "i64.extend16_s(0x8000) should be -32768"
+    );
+
+    // Zero
+    assert_eq!(
+        extend16_mod.func_0(0).unwrap(),
+        0i64,
+        "i64.extend16_s(0) should be 0"
+    );
+
+    // Max positive signed short
+    assert_eq!(
+        extend16_mod.func_0(0x7FFF).unwrap(),
+        32767i64,
+        "i64.extend16_s(0x7FFF) should be 32767"
+    );
+}
+
+#[test]
+fn test_i64_extend32_s() {
+    let mut extend32_mod = i64_extend32_s::new().unwrap();
+
+    // Positive value: 0x12345678 sign-extended
+    assert_eq!(
+        extend32_mod.func_0(0x12345678i64).unwrap(),
+        0x12345678i64,
+        "i64.extend32_s(0x12345678) should be 0x12345678"
+    );
+
+    // Negative value: 0xFFFFFFFF sign-extended to i64 (all bits set, -1)
+    assert_eq!(
+        extend32_mod.func_0(0xFFFFFFFFi64).unwrap(),
+        -1i64,
+        "i64.extend32_s(0xFFFFFFFF) should be -1"
+    );
+
+    // Another negative: 0x80000000 sign-extended (bit 31 set)
+    assert_eq!(
+        extend32_mod.func_0(0x80000000i64).unwrap(),
+        -2147483648i64,
+        "i64.extend32_s(0x80000000) should be -2147483648"
+    );
+
+    // Zero
+    assert_eq!(
+        extend32_mod.func_0(0).unwrap(),
+        0i64,
+        "i64.extend32_s(0) should be 0"
+    );
+
+    // Max positive signed int
+    assert_eq!(
+        extend32_mod.func_0(0x7FFFFFFFi64).unwrap(),
+        2147483647i64,
+        "i64.extend32_s(0x7FFFFFFF) should be 2147483647"
+    );
+}

crates/herkos/src/backend/safe.rs

@@ -548,7 +548,7 @@ impl Backend for SafeBackend {
     }
 
     fn emit_memory_fill(&self, dst: VarId, val: VarId, len: VarId) -> String {
-        format!("                memory.fill({dst} as u32, {val} as u8, {len} as u32)?;")
+        format!("                memory.fill({dst} as usize, {val} as u8, {len} as usize)?;")
     }
 
     fn emit_memory_init(
@@ -558,7 +558,7 @@ impl Backend for SafeBackend {
         len: VarId,
         segment_const_name: &str,
     ) -> String {
-        format!("                memory.init_data_partial({dst} as u32, {segment_const_name}, {src_offset} as usize, {len} as usize)?;")
+        format!("                memory.init_data_partial({dst} as usize, {segment_const_name}, {src_offset} as usize, {len} as usize)?;")
     }
 
     fn emit_data_drop(&self, segment: u32) -> String {