feat: basic working wasm adder (experiment)

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

10 files changed, 455 insertions, 162 deletions

diff --git a/crates/cli/Cargo.toml b/crates/cli/Cargo.toml
index 4a9974b..bfe0391 100644
--- a/crates/cli/Cargo.toml
+++ b/crates/cli/Cargo.toml
@@ -13,3 +13,5 @@ path="bin.rs"
 tinywasm={path="../tinywasm"}
 argh="0.1"
 color-eyre="0.6"
+tracing="0.1"
+tracing-subscriber="0.3"
diff --git a/crates/cli/bin.rs b/crates/cli/bin.rs
index 4e85481..38e30da 100644
--- a/crates/cli/bin.rs
+++ b/crates/cli/bin.rs
@@ -1,6 +1,9 @@
 use argh::FromArgs;
 use color_eyre::eyre::Result;
-use tinywasm::{self, Module};
+use tinywasm::{self, module::WasmValue, Module};
+use util::install_tracing;
+
+mod util;
 
 #[derive(FromArgs)]
 /// TinyWasm CLI
@@ -25,6 +28,9 @@ struct Run {
 }
 
 fn main() -> Result<()> {
+    color_eyre::install()?;
+    install_tracing(None);
+
     let args: TinyWasmCli = argh::from_env();
 
     match args.nested {
@@ -37,7 +43,14 @@ fn main() -> Result<()> {
 }
 
 fn run(wasm: &[u8]) -> Result<()> {
-    let module = Module::new(wasm)?;
-    println!("{:#?}", module);
+    let mut module = Module::new(wasm)?;
+    let args = [WasmValue::I32(1), WasmValue::I32(2)];
+    let res = module.run("add", &args)?;
+    println!("res: {:?}", res);
+
+    let args = [WasmValue::I64(1), WasmValue::I64(2)];
+    let res = module.run("add_64", &args)?;
+    println!("res: {:?}", res);
+
     Ok(())
 }
diff --git a/crates/cli/util.rs b/crates/cli/util.rs
new file mode 100644
index 0000000..2a8160a
--- /dev/null
+++ b/crates/cli/util.rs
@@ -0,0 +1,14 @@
+pub fn install_tracing(log_level: Option<tracing::Level>) {
+    use tracing_subscriber::filter::LevelFilter;
+    use tracing_subscriber::prelude::*;
+
+    tracing_subscriber::registry()
+        .with(
+            tracing_subscriber::fmt::layer()
+                .compact()
+                .with_filter(LevelFilter::from(
+                    log_level.unwrap_or(tracing::Level::DEBUG),
+                )),
+        )
+        .init();
+}
diff --git a/crates/tinywasm/Cargo.toml b/crates/tinywasm/Cargo.toml
index 2f71aa9..807a19f 100644
--- a/crates/tinywasm/Cargo.toml
+++ b/crates/tinywasm/Cargo.toml
@@ -10,6 +10,7 @@ path="src/lib.rs"
 thiserror={version="1.0", package="thiserror-core", default-features=false}
 wasmparser={version="0.100", package="wasmparser-nostd", default-features=false}
 tracing={version="0.1.38", default-features=false} # logging
+hashbrown="0.14"
 
 [features]
 default=["std"]
diff --git a/crates/tinywasm/src/error.rs b/crates/tinywasm/src/error.rs
index b9c93c4..a253880 100644
--- a/crates/tinywasm/src/error.rs
+++ b/crates/tinywasm/src/error.rs
@@ -6,6 +6,9 @@ pub enum Error {
     #[error("error parsing module")]
     ParseError { message: String, offset: usize },
 
+    #[error("unsupported feature: {0}")]
+    UnsupportedFeature(String),
+
     #[error("unknown error: {0}")]
     Other(String),
 }
@@ -14,6 +17,10 @@ impl Error {
     pub fn other<T>(message: &str) -> Result<T, Self> {
         Err(Self::Other(message.to_string()))
     }
+
+    pub fn unsupported<T>(feature: &str) -> Result<T, Self> {
+        Err(Self::UnsupportedFeature(feature.to_string()))
+    }
 }
 
 impl From<wasmparser::BinaryReaderError> for Error {
diff --git a/crates/tinywasm/src/instance.rs b/crates/tinywasm/src/instance.rs
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/crates/tinywasm/src/instance.rs
diff --git a/crates/tinywasm/src/lib.rs b/crates/tinywasm/src/lib.rs
index 3d572a2..0f97863 100644
--- a/crates/tinywasm/src/lib.rs
+++ b/crates/tinywasm/src/lib.rs
@@ -19,15 +19,12 @@ pub struct Instance {}
 mod tests {
     use super::*;
     use crate::{error::Result, Module};
-    use std::dbg;
 
     #[test]
     fn it_works() -> Result<()> {
-        let wasm = include_bytes!("../../../examples/wasm/helloworld.wasm");
+        let wasm = include_bytes!("../../../examples/wasm/add.wasm");
         let module = Module::new(wasm)?;
 
-        dbg!(module);
-
         Ok(())
     }
 }
diff --git a/crates/tinywasm/src/module.rs b/crates/tinywasm/src/module.rs
deleted file mode 100644
index 1ed42c4..0000000
--- a/crates/tinywasm/src/module.rs
+++ /dev/null
@@ -1,155 +0,0 @@
-use crate::error::{Error, Result};
-use alloc::{format, vec::Vec};
-use tracing::error;
-use wasmparser::*;
-
-pub struct Module<'a> {
-    reader: ModuleReader<'a>,
-}
-
-impl<'a> core::fmt::Debug for Module<'a> {
-    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
-        f.debug_struct("Module")
-            .field("version", &self.reader.version)
-            .field("type_section", &self.reader.type_section)
-            .field("function_section", &self.reader.function_section)
-            .field("table_section", &self.reader.table_section)
-            .field("memory_section", &self.reader.memory_section)
-            .field("global_section", &self.reader.global_section)
-            .field("element_section", &self.reader.element_section)
-            .field("data_section", &self.reader.data_section)
-            .field("code_section", &self.reader.code_section)
-            .field("import_section", &self.reader.import_section)
-            .field("export_section", &self.reader.export_section)
-            .finish()
-    }
-}
-
-#[derive(Default)]
-pub struct ModuleReader<'a> {
-    pub version: Option<u16>,
-    pub type_section: Option<TypeSectionReader<'a>>,
-    pub function_section: Option<FunctionSectionReader<'a>>,
-    pub table_section: Option<TableSectionReader<'a>>,
-    pub memory_section: Option<MemorySectionReader<'a>>,
-    pub global_section: Option<GlobalSectionReader<'a>>,
-    pub element_section: Option<ElementSectionReader<'a>>,
-    pub data_section: Option<DataSectionReader<'a>>,
-    pub code_section: Option<CodeSection<'a>>,
-    pub import_section: Option<ImportSectionReader<'a>>,
-    pub export_section: Option<ExportSectionReader<'a>>,
-}
-
-#[derive(Debug)]
-pub struct CodeSection<'a> {
-    pub(crate) functions: Vec<FunctionBody<'a>>,
-}
-
-impl<'a> CodeSection<'a> {
-    fn new() -> Self {
-        Self {
-            functions: Vec::new(),
-        }
-    }
-}
-
-impl<'a> Module<'a> {
-    pub fn new(wasm: &'a [u8]) -> Result<Self> {
-        let mut validator = Validator::new();
-        let mut reader = ModuleReader::new();
-
-        for payload in wasmparser::Parser::new(0).parse_all(wasm) {
-            reader.process_payload(payload?, &mut validator)?;
-        }
-
-        Ok(Self { reader })
-    }
-}
-
-impl<'a> ModuleReader<'a> {
-    pub fn new() -> Self {
-        Self::default()
-    }
-
-    pub fn process_payload(
-        &mut self,
-        payload: Payload<'a>,
-        validator: &mut Validator,
-    ) -> Result<()> {
-        use wasmparser::Payload::*;
-        match payload {
-            Version {
-                num,
-                encoding,
-                range,
-            } => {
-                validator.version(num, encoding, &range)?;
-                self.version = Some(num);
-                match encoding {
-                    wasmparser::Encoding::Module => {}
-                    wasmparser::Encoding::Component => return Error::other("Component"),
-                }
-            }
-            TypeSection(reader) => {
-                validator.type_section(&reader)?;
-                self.type_section = Some(reader);
-            }
-            FunctionSection(reader) => {
-                validator.function_section(&reader)?;
-                self.function_section = Some(reader);
-            }
-            TableSection(reader) => {
-                validator.table_section(&reader)?;
-                self.table_section = Some(reader);
-            }
-            MemorySection(reader) => {
-                validator.memory_section(&reader)?;
-                self.memory_section = Some(reader);
-            }
-            GlobalSection(reader) => {
-                validator.global_section(&reader)?;
-                self.global_section = Some(reader);
-            }
-            ElementSection(reader) => {
-                validator.element_section(&reader)?;
-                self.element_section = Some(reader);
-            }
-            DataSection(reader) => {
-                validator.data_section(&reader)?;
-                self.data_section = Some(reader);
-            }
-            CodeSectionStart { count, range, .. } => {
-                validator.code_section_start(count, &range)?;
-
-                self.code_section = Some(CodeSection::new());
-            }
-            CodeSectionEntry(function) => {
-                validator.code_section_entry(&function)?;
-
-                if let Some(code_section) = &mut self.code_section {
-                    code_section.functions.push(function);
-                } else {
-                    return Error::other("Empty code section");
-                }
-            }
-            ImportSection(reader) => {
-                validator.import_section(&reader)?;
-                self.import_section = Some(reader);
-            }
-            ExportSection(reader) => {
-                validator.export_section(&reader)?;
-                self.export_section = Some(reader);
-            }
-
-            End(offset) => {
-                validator.end(offset)?;
-                return Ok(());
-            }
-            x => Error::other(&format!("Unknown payload: {:?}", x))?,
-        };
-
-        error!("Missing end");
-
-        Ok(())
-    }
-}
diff --git a/crates/tinywasm/src/module/mod.rs b/crates/tinywasm/src/module/mod.rs
new file mode 100644
index 0000000..9f69b32
--- /dev/null
+++ b/crates/tinywasm/src/module/mod.rs
@@ -0,0 +1,231 @@
+use core::fmt::Debug;
+
+use crate::error::{Error, Result};
+use alloc::{format, string::String, vec, vec::Vec};
+use hashbrown::{HashMap, HashSet};
+use tracing::{error, info};
+use wasmparser::*;
+
+mod reader;
+use self::reader::ModuleReader;
+
+#[derive(Debug)]
+pub struct ModuleMetadata {
+    pub version: u16,
+}
+
+pub struct Module<'data> {
+    pub meta: ModuleMetadata,
+
+    pub types: Vec<FuncType>,
+    pub functions: Vec<u32>,
+    pub exports: Vec<Export<'data>>,
+    pub code: Vec<FunctionBody<'data>>,
+
+    marker: core::marker::PhantomData<&'data ()>,
+}
+
+impl Debug for Module<'_> {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.debug_struct("Module")
+            .field("meta", &self.meta)
+            .field("types", &self.types)
+            .field("functions", &self.functions)
+            .field("exports", &self.exports)
+            .field("code", &self.code)
+            .finish()
+    }
+}
+
+#[derive(Debug)]
+#[non_exhaustive]
+pub enum WasmValue {
+    I32(i32),
+    I64(i64),
+}
+
+impl WasmValue {
+    fn to_bytes(&self) -> (Vec<u8>, ValType) {
+        match self {
+            Self::I32(val) => {
+                let val = val.to_le_bytes().to_vec();
+                (val, ValType::I32)
+            }
+            Self::I64(val) => {
+                let val = val.to_le_bytes().to_vec();
+                (val, ValType::I64)
+            } // _ => {
+              //     panic!("Unsupported return type");
+              // }
+        }
+    }
+
+    fn from_bytes(bytes: &[u8], ty: &ValType) -> Self {
+        match ty {
+            ValType::I32 => {
+                let val = i32::from_le_bytes(bytes.try_into().unwrap());
+                Self::I32(val)
+            }
+            ValType::I64 => {
+                let val = i64::from_le_bytes(bytes.try_into().unwrap());
+                Self::I64(val)
+            }
+            _ => {
+                panic!("Unsupported return type");
+            }
+        }
+    }
+}
+
+impl<'data> Module<'data> {
+    pub fn new(wasm: &'data [u8]) -> Result<Self> {
+        let mut validator = Validator::new();
+        let mut reader = ModuleReader::new();
+
+        for payload in wasmparser::Parser::new(0).parse_all(wasm) {
+            reader.process_payload(payload?, &mut validator)?;
+        }
+
+        if !reader.end_reached {
+            return Error::other("End not reached");
+        }
+
+        Self::from_reader(reader)
+    }
+
+    pub fn run(&mut self, func_name: &str, args: &[WasmValue]) -> Result<Vec<WasmValue>> {
+        let func = self
+            .exports
+            .iter()
+            .find(|e| e.name == func_name)
+            .ok_or_else(|| Error::Other(format!("Function {} not found", func_name)))?;
+
+        let func_type_index = self.functions[func.index as usize];
+        let func_type = &self.types[func_type_index as usize];
+
+        info!("func_type: {:#?}", func_type);
+        let code = &mut self.code[func.index as usize];
+        code.allow_memarg64(false);
+
+        let mut locals = vec![];
+        for ty in func_type.params() {
+            locals.push(ty.clone());
+        }
+
+        let mut returns = vec![];
+        for ty in func_type.results() {
+            returns.push(ty.clone());
+        }
+
+        let locals_reader = code.get_locals_reader().unwrap();
+        for local in locals_reader.into_iter() {
+            let local = local.unwrap();
+            if locals.len() != local.0 as usize {
+                panic!("Invalid local index");
+            }
+            locals.push(local.1);
+        }
+
+        let mut body = code.get_operators_reader().unwrap().into_iter();
+
+        let mut local_values = vec![];
+        for (i, arg) in args.iter().enumerate() {
+            let (val, ty) = arg.to_bytes();
+            if locals[i] != ty {
+                return Error::other(&format!(
+                    "Invalid argument type for {}, index {}: expected {:?}, got {:?}",
+                    func_name, i, locals[i], ty
+                ));
+            }
+
+            local_values.push(val);
+        }
+
+        let mut stack: Vec<Vec<u8>> = Vec::new();
+        while let Some(op) = body.next() {
+            let op = op.unwrap();
+            info!("op: {:#?}", op);
+
+            match op {
+                Operator::LocalGet { local_index } => {
+                    let local = locals.get(local_index as usize).unwrap();
+                    let val = &local_values[local_index as usize];
+                    info!("local: {:#?}", local);
+                    stack.push(val.clone());
+                }
+                Operator::I64Add => {
+                    let a = stack.pop().unwrap();
+                    let b = stack.pop().unwrap();
+                    let a = i64::from_le_bytes(a.try_into().unwrap());
+                    let b = i64::from_le_bytes(b.try_into().unwrap());
+                    let c = (a + b).to_le_bytes().to_vec();
+                    stack.push(c);
+                }
+                Operator::I32Add => {
+                    let a = stack.pop().unwrap();
+                    let b = stack.pop().unwrap();
+                    let a = i32::from_le_bytes(a.try_into().unwrap());
+                    let b = i32::from_le_bytes(b.try_into().unwrap());
+                    let c = (a + b).to_le_bytes().to_vec();
+                    stack.push(c);
+                }
+                Operator::End => {
+                    info!("stack: {:#?}", stack);
+                    let res = returns
+                        .iter()
+                        .map(|ty| {
+                            let val = stack.pop().unwrap();
+                            WasmValue::from_bytes(&val, ty)
+                        })
+                        .collect::<Vec<_>>();
+                    return Ok(res);
+                }
+                _ => {}
+            }
+        }
+
+        return Error::other("End not reached");
+    }
+
+    fn from_reader(reader: ModuleReader<'data>) -> Result<Self> {
+        let types = reader
+            .type_section
+            .map(|s| {
+                s.into_iter()
+                    .map(|ty| {
+                        let Type::Func(func) = ty?;
+                        Ok(func)
+                    })
+                    .collect::<Result<Vec<_>>>()
+            })
+            .transpose()?
+            .unwrap_or_default();
+
+        let functions = reader
+            .function_section
+            .map(|s| s.into_iter().map(|f| Ok(f?)).collect::<Result<Vec<_>>>())
+            .transpose()?
+            .unwrap_or_default();
+
+        let exports = reader
+            .export_section
+            .map(|s| s.into_iter().map(|e| Ok(e?)).collect::<Result<Vec<_>>>())
+            .transpose()?
+            .unwrap_or_default();
+
+        let code = reader.code_section.map(|s| s.functions).unwrap_or_default();
+
+        let meta = ModuleMetadata {
+            version: reader.version.unwrap_or(1),
+        };
+
+        Ok(Self {
+            marker: core::marker::PhantomData,
+            meta,
+            types,
+            exports,
+            functions,
+            code,
+        })
+    }
+}
diff --git a/crates/tinywasm/src/module/reader.rs b/crates/tinywasm/src/module/reader.rs
new file mode 100644
index 0000000..7ecfddf
--- /dev/null
+++ b/crates/tinywasm/src/module/reader.rs
@@ -0,0 +1,183 @@
+use alloc::{format, vec::Vec};
+use core::fmt::Debug;
+use tracing::debug;
+use wasmparser::{
+    DataSectionReader, ElementSectionReader, ExportSectionReader, FunctionBody,
+    FunctionSectionReader, GlobalSectionReader, ImportSectionReader, MemorySectionReader, Payload,
+    TableSectionReader, TypeSectionReader, Validator,
+};
+
+use crate::{Error, Result};
+
+#[derive(Default)]
+pub struct ModuleReader<'a> {
+    pub version: Option<u16>,
+
+    pub type_section: Option<TypeSectionReader<'a>>,
+    pub function_section: Option<FunctionSectionReader<'a>>,
+    pub export_section: Option<ExportSectionReader<'a>>,
+    pub code_section: Option<CodeSection<'a>>,
+
+    pub table_section: Option<TableSectionReader<'a>>,
+    pub memory_section: Option<MemorySectionReader<'a>>,
+    pub global_section: Option<GlobalSectionReader<'a>>,
+    pub element_section: Option<ElementSectionReader<'a>>,
+    pub data_section: Option<DataSectionReader<'a>>,
+    pub import_section: Option<ImportSectionReader<'a>>,
+    pub end_reached: bool,
+}
+
+impl Debug for ModuleReader<'_> {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.debug_struct("ModuleReader")
+            .field("version", &self.version)
+            .field("type_section", &self.type_section)
+            .field("function_section", &self.function_section)
+            .field("table_section", &self.table_section)
+            .field("memory_section", &self.memory_section)
+            .field("global_section", &self.global_section)
+            .field("element_section", &self.element_section)
+            .field("data_section", &self.data_section)
+            .field("code_section", &self.code_section)
+            .field("import_section", &self.import_section)
+            .field("export_section", &self.export_section)
+            .finish()
+    }
+}
+
+impl<'a> ModuleReader<'a> {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn process_payload(
+        &mut self,
+        payload: Payload<'a>,
+        validator: &mut Validator,
+    ) -> Result<()> {
+        use wasmparser::Payload::*;
+
+        match payload {
+            Version {
+                num,
+                encoding,
+                range,
+            } => {
+                validator.version(num, encoding, &range)?;
+                self.version = Some(num);
+                match encoding {
+                    wasmparser::Encoding::Module => {}
+                    wasmparser::Encoding::Component => return Error::other("Component"),
+                }
+            }
+            TypeSection(reader) => {
+                debug!("Found type section");
+                validator.type_section(&reader)?;
+                self.type_section = Some(reader);
+            }
+            FunctionSection(reader) => {
+                debug!("Found function section");
+                validator.function_section(&reader)?;
+                self.function_section = Some(reader);
+            }
+            TableSection(_reader) => {
+                return Error::unsupported("Table section");
+                // debug!("Found table section");
+                // validator.table_section(&reader)?;
+                // self.table_section = Some(reader);
+            }
+            MemorySection(_reader) => {
+                return Error::unsupported("Memory section");
+                // debug!("Found memory section");
+                // validator.memory_section(&reader)?;
+                // self.memory_section = Some(reader);
+            }
+            GlobalSection(_reader) => {
+                return Error::unsupported("Global section");
+                // debug!("Found global section");
+                // validator.global_section(&reader)?;
+                // self.global_section = Some(reader);
+            }
+            ElementSection(_reader) => {
+                return Error::unsupported("Element section");
+                // debug!("Found element section");
+                // validator.element_section(&reader)?;
+                // self.element_section = Some(reader);
+            }
+            DataSection(_reader) => {
+                return Error::unsupported("Data section");
+                // debug!("Found data section");
+                // validator.data_section(&reader)?;
+                // self.data_section = Some(reader);
+            }
+            CodeSectionStart { count, range, .. } => {
+                debug!("Found code section ({} functions)", count);
+                validator.code_section_start(count, &range)?;
+                self.code_section = Some(CodeSection::new());
+            }
+            CodeSectionEntry(function) => {
+                debug!("Found code section entry");
+                validator.code_section_entry(&function)?;
+
+                if let Some(code_section) = &mut self.code_section {
+                    code_section.functions.push(function);
+                } else {
+                    return Error::other("Empty code section");
+                }
+            }
+            ImportSection(_reader) => {
+                return Error::unsupported("Import section");
+
+                // debug!("Found import section");
+                // validator.import_section(&reader)?;
+                // self.import_section = Some(reader);
+            }
+            ExportSection(reader) => {
+                debug!("Found export section");
+                validator.export_section(&reader)?;
+                self.export_section = Some(reader);
+            }
+            End(offset) => {
+                debug!("Reached end of module");
+                if self.end_reached {
+                    return Error::other("End reached twice");
+                }
+
+                validator.end(offset)?;
+                self.end_reached = true;
+            }
+            x => Error::other(&format!("Unknown payload: {:?}", x))?,
+        };
+
+        return Ok(());
+    }
+
+    // fn exports(&mut self) -> Result<Vec<Export>> {
+    //     let mut exports = Vec::new();
+
+    //     if let Some(export_section) = self.export_section {
+    //         for export in export_section.into_iter() {
+    //             let export = export?;
+    //             let name = export.name;
+    //             let kind = export.kind;
+    //             let index = export.index;
+    //             exports.push(Export { name, kind, index });
+    //         }
+    //     }
+
+    //     Ok(exports)
+    // }
+}
+
+#[derive(Debug)]
+pub struct CodeSection<'a> {
+    pub(crate) functions: Vec<FunctionBody<'a>>,
+}
+
+impl<'a> CodeSection<'a> {
+    fn new() -> Self {
+        Self {
+            functions: Vec::new(),
+        }
+    }
+}