fix skipping immediates for memory instructions

crates/wasm_interp/src/instance.rs

@@ -2,7 +2,7 @@ use bumpalo::{collections::Vec, Bump};
 use std::fmt::{self, Write};
 use std::iter::{self, once, Iterator};
 
-use roc_wasm_module::opcodes::OpCode;
+use roc_wasm_module::opcodes::{MemoryInstruction, OpCode};
 use roc_wasm_module::parse::{Parse, SkipBytes};
 use roc_wasm_module::sections::{ImportDesc, MemorySection, SignatureParamsIter};
 use roc_wasm_module::{ExportType, WasmModule};
@@ -978,44 +978,35 @@ impl<'a, I: ImportDispatcher> Instance<'a, I> {
             }
             MEMORY => {
                 // the first argument determines exactly which memory operation we have
-                let op = module.code.bytes[self.program_counter];
+                match MemoryInstruction::try_from(module.code.bytes[self.program_counter]) {
-                self.program_counter += 1;
+                    Ok(op) => match op {
+                        MemoryInstruction::MemoryInit => todo!("WASM instruction: memory.init"),
+                        MemoryInstruction::DataDrop => todo!("WASM instruction: data.drop"),
+                        MemoryInstruction::MemoryCopy => {
+                            let size = self.value_store.pop_u32()? as usize;
+                            let source = self.value_store.pop_u32()? as usize;
+                            let destination = self.value_store.pop_u32()? as usize;
 
-                match op {
+                            // skip the op byte and an extra two zero bytes.
-                    8 => {
+                            // in future versions of WebAssembly this byte may be used to index additional memories
-                        // memory.init
+                            self.program_counter += 1 + 2;
-                        todo!("WASM instruction: memory.init")
-                    }
-                    9 => {
-                        // data.drop x
-                        todo!("WASM instruction: data.drop")
-                    }
-                    10 => {
-                        // memory.copy
-                        let size = self.value_store.pop_u32()? as usize;
-                        let source = self.value_store.pop_u32()? as usize;
-                        let destination = self.value_store.pop_u32()? as usize;
 
-                        // skip two zero bytes; in future versions of WebAssembly this byte may be
+                            self.memory.copy_within(source..source + size, destination)
-                        // used to index additional memories
+                        }
-                        self.program_counter += 2;
+                        MemoryInstruction::MemoryFill => {
+                            let size = self.value_store.pop_u32()? as usize;
+                            let byte_value = self.value_store.pop_u32()? as u8;
+                            let destination = self.value_store.pop_u32()? as usize;
 
-                        self.memory.copy_within(source..source + size, destination)
+                            // skip the op byte and an extra zero byte.
-                    }
+                            // in future versions of WebAssembly this byte may be used to index additional memories
-                    11 => {
+                            self.program_counter += 1 + 1;
-                        // memory.fill
-                        let size = self.value_store.pop_u32()? as usize;
-                        let byte_value = self.value_store.pop_u32()? as u8;
-                        let destination = self.value_store.pop_u32()? as usize;
 
-                        // skip a zero byte; in future versions of WebAssembly this byte may be
+                            self.memory[destination..][..size].fill(byte_value);
-                        // used to index additional memories
+                        }
-                        self.program_counter += 1;
+                    },
-
+                    Err(other) => unreachable!("invalid memory instruction {other:?}"),
-                        self.memory[destination..][..size].fill(byte_value);
+                };
-                    }
-                    other => unreachable!("invalid memory instruction {:?}", other),
-                }
             }
             I32CONST => {
                 let value = i32::parse((), &module.code.bytes, &mut self.program_counter).unwrap();

crates/wasm_module/src/opcodes.rs

@@ -192,6 +192,29 @@ impl From<u8> for OpCode {
     }
 }
 
+#[repr(u8)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum MemoryInstruction {
+    MemoryInit = 8,
+    DataDrop = 9,
+    MemoryCopy = 10,
+    MemoryFill = 11,
+}
+
+impl TryFrom<u8> for MemoryInstruction {
+    type Error = u8;
+
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        match value {
+            8 => Ok(Self::MemoryInit),
+            9 => Ok(Self::DataDrop),
+            10 => Ok(Self::MemoryCopy),
+            11 => Ok(Self::MemoryFill),
+            _ => Err(value),
+        }
+    }
+}
+
 /// The format of the *immediate* operands of an operator
 /// Immediates appear directly in the byte stream after the opcode,
 /// rather than being popped off the value stack. These are the possible forms.
@@ -205,6 +228,7 @@ enum OpImmediates {
     Leb64x1,
     Leb32x2,
     BrTable,
+    Memory,
 }
 
 fn immediates_for(op: OpCode) -> Result<OpImmediates, String> {
@@ -233,7 +257,7 @@ fn immediates_for(op: OpCode) -> Result<OpImmediates, String> {
         | I64STORE32 => Leb32x2,
 
         CURRENTMEMORY | GROWMEMORY => Byte1,
-        MEMORY => Leb32x2,
+        MEMORY => Memory,
 
         I32CONST => Leb32x1,
         I64CONST => Leb64x1,
@@ -314,6 +338,29 @@ impl SkipBytes for OpCode {
                     u32::skip_bytes(bytes, cursor)?;
                 }
             }
+            Memory => {
+                match MemoryInstruction::try_from(bytes[*cursor + 1]) {
+                    Ok(op) => match op {
+                        MemoryInstruction::MemoryInit => {
+                            // memory.init
+                            todo!("WASM instruction: memory.init")
+                        }
+                        MemoryInstruction::DataDrop => {
+                            // data.drop x
+                            todo!("WASM instruction: data.drop")
+                        }
+                        MemoryInstruction::MemoryCopy => {
+                            // memory.copy
+                            *cursor += 1 + 1 + 2;
+                        }
+                        MemoryInstruction::MemoryFill => {
+                            // memory.fill
+                            *cursor += 1 + 1 + 1;
+                        }
+                    },
+                    Err(other) => unreachable!("invalid memory instruction {other:?}"),
+                }
+            }
         }
         Ok(())
     }