chore: refactoring

Signed-off-by: Henry Gressmann <mail@henrygressmann.de>

6 files changed, 196 insertions, 118 deletions

diff --git a/crates/cli/bin.rs b/crates/cli/bin.rs
index ce33544..181b2e1 100644
--- a/crates/cli/bin.rs
+++ b/crates/cli/bin.rs
@@ -82,9 +82,8 @@ fn run(wasm: &[u8]) -> Result<()> {
     let mut store = tinywasm::Store::default();
 
     let module = tinywasm::Module::parse_bytes(wasm)?;
-    let mut instance = module.instantiate(&mut store)?;
-
-    let mut func = instance.get_func(&mut store, "add")?;
+    let instance = module.instantiate(&mut store)?;
+    let func = instance.get_func(&store, "add")?;
     let params = vec![WasmValue::I32(2), WasmValue::I32(2)];
     let res = func.call(&mut store, params)?;
     info!("{res:?}");
diff --git a/crates/parser/src/lib.rs b/crates/parser/src/lib.rs
index 0d12990..d91694b 100644
--- a/crates/parser/src/lib.rs
+++ b/crates/parser/src/lib.rs
@@ -15,6 +15,7 @@ use module::ModuleReader;
 use tinywasm_types::{Function, TinyWasmModule};
 use wasmparser::Validator;
 
+#[derive(Default)]
 pub struct Parser {}
 
 impl Parser {
@@ -102,8 +103,7 @@ impl TryFrom<ModuleReader> for TinyWasmModule {
                 locals: f.locals,
                 ty,
             })
-            .collect::<Vec<_>>()
-            .into_boxed_slice();
+            .collect::<Vec<_>>();
 
         Ok(TinyWasmModule {
             version: reader.version,
diff --git a/crates/tinywasm/src/module/mod.rs b/crates/tinywasm/src/module/mod.rs
index 2a5141c..db0c876 100644
--- a/crates/tinywasm/src/module/mod.rs
+++ b/crates/tinywasm/src/module/mod.rs
@@ -2,18 +2,16 @@ use alloc::{
     boxed::Box,
     format,
     string::{String, ToString},
+    sync::Arc,
     vec,
     vec::Vec,
 };
-use log::info;
 use tinywasm_types::{
-    Export, ExternalKind, FuncAddr, FuncType, TinyWasmModule, ValType, WasmValue,
+    Export, ExternalKind, FuncAddr, FuncType, ModuleInstanceAddr, TinyWasmModule, ValType,
+    WasmValue,
 };
 
-use crate::{
-    store::{self, StoreData},
-    Error, Result, Store,
-};
+use crate::{runtime::Stack, store, Error, Result, Store};
 
 #[derive(Debug)]
 pub struct Module {
@@ -56,21 +54,37 @@ impl Module {
         store: &mut Store,
         // imports: Option<()>,
     ) -> Result<ModuleInstance> {
-        let mut i = ModuleInstance::new(store, self)?;
-        let _ = i.start(store)?;
-        Ok(i)
+        let idx = store.next_module_instance_idx();
+
+        let func_addrs = store.add_funcs(self.data.funcs, idx);
+        let instance = ModuleInstance::new(
+            self.data.types,
+            self.data.start_func,
+            self.data.exports,
+            func_addrs,
+            idx,
+        );
+
+        store.add_instance(instance.clone())?;
+        // let _ = instance.start(store)?;
+        Ok(instance)
     }
 }
 
 /// A WebAssembly Module Instance.
 /// Addrs are indices into the store's data structures.
 /// See https://webassembly.github.io/spec/core/exec/runtime.html#module-instances
+#[derive(Debug, Clone)]
+pub struct ModuleInstance(Arc<ModuleInstanceInner>);
+
 #[derive(Debug)]
-pub struct ModuleInstance {
+struct ModuleInstanceInner {
+    pub(crate) _idx: ModuleInstanceAddr,
     pub(crate) func_start: Option<FuncAddr>,
     pub(crate) types: Box<[FuncType]>,
-    pub(crate) exports: Box<[Export]>,
-    // pub(crate) func_addrs: Vec<FuncAddr>,
+    pub exports: ExportInstance,
+
+    pub(crate) func_addrs: Vec<FuncAddr>,
     // pub table_addrs: Vec<TableAddr>,
     // pub mem_addrs: Vec<MemAddr>,
     // pub global_addrs: Vec<GlobalAddr>,
@@ -78,21 +92,45 @@ pub struct ModuleInstance {
     // pub data_addrs: Vec<DataAddr>,
 }
 
-impl ModuleInstance {
-    /// Get an exported function by name
-    pub fn get_func(&mut self, store: &store::Store, name: &str) -> Result<FuncHandle> {
-        let export = self
-            .exports
+#[derive(Debug)]
+pub struct ExportInstance(Box<[Export]>);
+
+impl ExportInstance {
+    pub fn func(&self, name: &str) -> Result<&Export> {
+        self.0
             .iter()
             .find(|e| e.name == name.into() && e.kind == ExternalKind::Func)
-            .ok_or(Error::Other(format!("export {} not found", name)))?;
+            .ok_or(Error::Other(format!("export {} not found", name)))
+    }
+}
 
-        let func = store.get_func(export.index as usize)?;
-        let ty = self.types[func.ty as usize].clone();
+impl ModuleInstance {
+    fn new(
+        types: Box<[FuncType]>,
+        func_start: Option<FuncAddr>,
+        exports: Box<[Export]>,
+        func_addrs: Vec<FuncAddr>,
+        idx: ModuleInstanceAddr,
+    ) -> Self {
+        Self(Arc::new(ModuleInstanceInner {
+            _idx: idx,
+            types,
+            func_start,
+            func_addrs,
+            exports: ExportInstance(exports),
+        }))
+    }
+
+    /// Get an exported function by name
+    pub fn get_func(&self, store: &store::Store, name: &str) -> Result<FuncHandle> {
+        let export = self.0.exports.func(name)?;
+        let func_addr = self.0.func_addrs[export.index as usize];
+        let func = store.get_func(func_addr as usize)?;
+        let ty = self.0.types[func.ty_addr() as usize].clone();
 
         Ok(FuncHandle {
             addr: export.index,
-            module: self,
+            _module: self.clone(),
             name: Some(name.to_string()),
             ty,
         })
@@ -100,44 +138,27 @@ impl ModuleInstance {
 
     /// Get the start  function of the module
     pub fn get_start_func(&mut self, store: &store::Store) -> Result<Option<FuncHandle>> {
-        let Some(addr) = self.func_start else {
+        let Some(func_index) = self.0.func_start else {
             return Ok(None);
         };
 
-        let func = store.get_func(addr as usize)?;
-        let ty = self.types[func.ty as usize].clone();
+        let func_addr = self.0.func_addrs[func_index as usize];
+        let func = store.get_func(func_addr as usize)?;
+        let ty = self.0.types[func.ty_addr() as usize].clone();
 
         Ok(Some(FuncHandle {
-            module: self,
-            addr,
+            _module: self.clone(),
+            addr: func_addr,
             ty,
             name: None,
         }))
     }
 
-    pub fn new(store: &mut Store, mut module: Module) -> Result<Self> {
-        let store_data = StoreData {
-            funcs: module.data.funcs,
-        };
-
-        store.initialize(store_data)?;
-        Ok(Self {
-            types: module.data.types,
-            func_start: module.data.start_func,
-            // table_addrs,
-            // mem_addrs,
-            // global_addrs,
-            // elem_addrs,
-            // data_addrs,
-            exports: module.data.exports,
-        })
-    }
-
     /// Invoke the start function of the module
     /// Returns None if the module has no start function
     /// https://webassembly.github.io/spec/core/syntax/modules.html#syntax-start
     pub fn start(&mut self, store: &mut store::Store) -> Result<Option<()>> {
-        let Some(mut func) = self.get_start_func(store)? else {
+        let Some(func) = self.get_start_func(store)? else {
             return Ok(None);
         };
 
@@ -147,16 +168,16 @@ impl ModuleInstance {
 }
 
 #[derive(Debug)]
-pub struct FuncHandle<'a> {
-    module: &'a mut ModuleInstance,
+pub struct FuncHandle {
+    _module: ModuleInstance,
     addr: FuncAddr,
     ty: FuncType,
     pub name: Option<String>,
 }
 
-impl<'a> FuncHandle<'a> {
+impl FuncHandle {
     /// Call a function
-    pub fn call(&mut self, store: &mut Store, params: Vec<WasmValue>) -> Result<Vec<WasmValue>> {
+    pub fn call(&self, store: &mut Store, params: Vec<WasmValue>) -> Result<Vec<WasmValue>> {
         let func = store
             .data
             .funcs
@@ -177,56 +198,25 @@ impl<'a> FuncHandle<'a> {
 
         let mut local_types: Vec<ValType> = Vec::new();
         local_types.extend(func_ty.params.iter());
-        local_types.extend(func.locals.iter());
+        local_types.extend(func.locals().iter());
 
-        let runtime = &mut store.runtime;
-        let stack = &mut runtime.stack;
-        let locals = &mut stack.locals;
-        locals.extend(params);
+        // let runtime = &mut store.runtime;
 
-        let mut instrs = func.instructions.iter();
-        while let Some(instr) = instrs.next() {
-            use tinywasm_types::Instruction::*;
-            match instr {
-                LocalGet(local_index) => {
-                    let val = &locals[*local_index as usize];
-                    info!("local: {:#?}", val);
-                    stack.value_stack.push(val.clone());
-                }
-                I64Add => {
-                    let a = stack.value_stack.pop().unwrap();
-                    let b = stack.value_stack.pop().unwrap();
-                    let (WasmValue::I64(a), WasmValue::I64(b)) = (a, b) else {
-                        panic!("Invalid type");
-                    };
-                    let c = WasmValue::I64(a + b);
-                    stack.value_stack.push(c);
-                }
-                I32Add => {
-                    let a = stack.value_stack.pop().unwrap();
-                    let b = stack.value_stack.pop().unwrap();
-                    let (WasmValue::I32(a), WasmValue::I32(b)) = (a, b) else {
-                        panic!("Invalid type");
-                    };
-                    let c = WasmValue::I32(a + b);
-                    stack.value_stack.push(c);
-                }
-                End => {
-                    let res = func_ty
-                        .results
-                        .iter()
-                        .map(|_| runtime.stack.value_stack.pop())
-                        .collect::<Option<Vec<_>>>()
-                        .ok_or(Error::Other(
-                            "function did not return the correct number of values".into(),
-                        ))?;
+        let mut stack = Stack::default();
+        stack.locals.extend(params);
 
-                    return Ok(res);
-                }
-                _ => todo!(),
-            }
-        }
+        let instrs = func.instructions().iter();
+        store.runtime.exec(&mut stack, instrs)?;
+
+        let res = func_ty
+            .results
+            .iter()
+            .map(|_| stack.value_stack.pop())
+            .collect::<Option<Vec<_>>>()
+            .ok_or(Error::Other(
+                "function did not return the correct number of values".into(),
+            ))?;
 
-        Err(Error::FuncDidNotReturn)
+        Ok(res)
     }
 }
diff --git a/crates/tinywasm/src/runtime/mod.rs b/crates/tinywasm/src/runtime/mod.rs
index 76686f7..c4c1892 100644
--- a/crates/tinywasm/src/runtime/mod.rs
+++ b/crates/tinywasm/src/runtime/mod.rs
@@ -2,11 +2,57 @@ mod executer;
 mod stack;
 
 pub use executer::*;
+use log::info;
 pub use stack::*;
+use tinywasm_types::{Instruction, WasmValue};
+
+use crate::{Error, Result};
 
 /// A WebAssembly Runtime.
 /// See https://webassembly.github.io/spec/core/exec/runtime.html
 #[derive(Debug, Default)]
-pub struct Runtime {
-    pub stack: Stack,
+pub struct Runtime {}
+
+impl Runtime {
+    pub(crate) fn exec(
+        &self,
+        stack: &mut Stack,
+        instrs: core::slice::Iter<Instruction>,
+    ) -> Result<()> {
+        let locals = &mut stack.locals;
+        for instr in instrs {
+            use tinywasm_types::Instruction::*;
+            match instr {
+                LocalGet(local_index) => {
+                    let val = &locals[*local_index as usize];
+                    info!("local: {:#?}", val);
+                    stack.value_stack.push(val.clone());
+                }
+                I64Add => {
+                    let a = stack.value_stack.pop().unwrap();
+                    let b = stack.value_stack.pop().unwrap();
+                    let (WasmValue::I64(a), WasmValue::I64(b)) = (a, b) else {
+                        panic!("Invalid type");
+                    };
+                    let c = WasmValue::I64(a + b);
+                    stack.value_stack.push(c);
+                }
+                I32Add => {
+                    let a = stack.value_stack.pop().unwrap();
+                    let b = stack.value_stack.pop().unwrap();
+                    let (WasmValue::I32(a), WasmValue::I32(b)) = (a, b) else {
+                        panic!("Invalid type");
+                    };
+                    let c = WasmValue::I32(a + b);
+                    stack.value_stack.push(c);
+                }
+                End => {
+                    return Ok(());
+                }
+                _ => todo!(),
+            }
+        }
+
+        Err(Error::FuncDidNotReturn)
+    }
 }
diff --git a/crates/tinywasm/src/store.rs b/crates/tinywasm/src/store.rs
index fc64bc4..d0e4f14 100644
--- a/crates/tinywasm/src/store.rs
+++ b/crates/tinywasm/src/store.rs
@@ -1,25 +1,46 @@
-use alloc::{boxed::Box, format};
-use tinywasm_types::Function;
+use alloc::{boxed::Box, format, vec::Vec};
+use tinywasm_types::{FuncAddr, Function, Instruction, ModuleInstanceAddr, TypeAddr, ValType};
 
-use crate::{runtime::Runtime, Error, Result};
+use crate::{runtime::Runtime, Error, ModuleInstance, Result};
 
 /// global state that can be manipulated by WebAssembly programs
 /// https://webassembly.github.io/spec/core/exec/runtime.html#store
 #[derive(Debug, Default)]
 pub struct Store {
+    module_instances: Vec<ModuleInstance>,
+    module_instance_count: usize,
+
     pub(crate) data: StoreData,
     pub(crate) runtime: Runtime,
 }
 
-// #[derive(Debug)]
-// pub struct FunctionInstance {
-//     pub(crate) func: Function,
-//     pub(crate) module: usize,
-// }
+#[derive(Debug)]
+pub struct FunctionInstance {
+    pub(crate) func: Function,
+    pub(crate) module_instance: ModuleInstanceAddr, // index into store.module_instances
+}
+
+impl FunctionInstance {
+    pub(crate) fn module_instance_addr(&self) -> ModuleInstanceAddr {
+        self.module_instance
+    }
+
+    pub(crate) fn locals(&self) -> &[ValType] {
+        &self.func.locals
+    }
+
+    pub(crate) fn instructions(&self) -> &[Instruction] {
+        &self.func.instructions
+    }
+
+    pub(crate) fn ty_addr(&self) -> TypeAddr {
+        self.func.ty
+    }
+}
 
 #[derive(Debug, Default)]
 pub struct StoreData {
-    pub funcs: Box<[Function]>,
+    pub(crate) funcs: Vec<FunctionInstance>,
     // pub tables: Vec<TableAddr>,
     // pub mems: Vec<MemAddr>,
     // pub globals: Vec<GlobalAddr>,
@@ -28,16 +49,37 @@ pub struct StoreData {
 }
 
 impl Store {
+    pub(crate) fn next_module_instance_idx(&self) -> ModuleInstanceAddr {
+        self.module_instance_count as ModuleInstanceAddr
+    }
+
     /// Initialize the store with global state from the given module
-    pub(crate) fn initialize(&mut self, data: StoreData) -> Result<()> {
-        self.data = data;
+    pub(crate) fn add_instance(&mut self, instance: ModuleInstance) -> Result<()> {
+        self.module_instances.push(instance);
+        self.module_instance_count += 1;
         Ok(())
     }
 
-    pub(crate) fn get_func(&self, index: usize) -> Result<&Function> {
+    pub(crate) fn add_funcs(
+        &mut self,
+        funcs: Vec<Function>,
+        idx: ModuleInstanceAddr,
+    ) -> Vec<FuncAddr> {
+        let mut func_addrs = Vec::with_capacity(funcs.len());
+        for func in funcs {
+            self.data.funcs.push(FunctionInstance {
+                func: func,
+                module_instance: idx,
+            });
+            func_addrs.push(idx as FuncAddr);
+        }
+        func_addrs
+    }
+
+    pub(crate) fn get_func(&self, addr: usize) -> Result<&FunctionInstance> {
         self.data
             .funcs
-            .get(index)
-            .ok_or_else(|| Error::Other(format!("function {} not found", index)))
+            .get(addr)
+            .ok_or_else(|| Error::Other(format!("function {} not found", addr)))
     }
 }
diff --git a/crates/types/src/lib.rs b/crates/types/src/lib.rs
index 9477dc7..a354203 100644
--- a/crates/types/src/lib.rs
+++ b/crates/types/src/lib.rs
@@ -10,8 +10,8 @@ pub struct TinyWasmModule {
     pub version: Option<u16>,
     pub start_func: Option<FuncAddr>,
 
+    pub funcs: Vec<Function>,
     pub types: Box<[FuncType]>,
-    pub funcs: Box<[Function]>,
     pub exports: Box<[Export]>,
     // pub tables: Option<TableType>,
     // pub memories: Option<MemoryType>,
@@ -85,6 +85,7 @@ pub enum ExternalKind {
 /// These are indexes into the respective stores.
 /// See https://webassembly.github.io/spec/core/exec/runtime.html#addresses
 pub type Addr = u32;
+pub type ModuleInstanceAddr = Addr;
 pub type FuncAddr = Addr;
 pub type TableAddr = Addr;
 pub type MemAddr = Addr;