core: attach types to functions

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

2 files changed, 34 insertions, 29 deletions

diff --git a/crates/parser/src/lib.rs b/crates/parser/src/lib.rs
index 0618683..ae6cc4d 100644
--- a/crates/parser/src/lib.rs
+++ b/crates/parser/src/lib.rs
@@ -20,7 +20,7 @@ mod log {
 mod conversion;
 mod error;
 mod module;
-use alloc::vec::Vec;
+use alloc::{string::ToString, vec::Vec};
 pub use error::*;
 use module::ModuleReader;
 use tinywasm_types::WasmFunction;
@@ -103,15 +103,26 @@ impl TryFrom<ModuleReader> for TinyWasmModule {
             return Err(ParseError::EndNotReached);
         }
 
-        let func_types = reader.func_addrs;
+        let code_type_addrs = reader.code_type_addrs;
+        let local_function_count = reader.code.len();
+
+        if code_type_addrs.len() != local_function_count {
+            return Err(ParseError::Other("Code and code type address count mismatch".to_string()));
+        }
+
         let funcs = reader
             .code
             .into_iter()
-            .zip(func_types)
-            .map(|(f, ty)| WasmFunction {
-                instructions: f.body,
-                locals: f.locals,
-                ty_addr: ty,
+            .zip(code_type_addrs)
+            .map(|(f, ty_idx)| {
+                (
+                    ty_idx,
+                    WasmFunction {
+                        instructions: f.body,
+                        locals: f.locals,
+                        ty: reader.func_types.get(ty_idx as usize).unwrap().clone(),
+                    },
+                )
             })
             .collect::<Vec<_>>();
 
@@ -119,17 +130,17 @@ impl TryFrom<ModuleReader> for TinyWasmModule {
         let table_types = reader.table_types;
 
         Ok(TinyWasmModule {
-            version: reader.version,
-            start_func: reader.start_func,
-            func_types: reader.func_types.into_boxed_slice(),
             funcs: funcs.into_boxed_slice(),
-            exports: reader.exports.into_boxed_slice(),
+            func_types: reader.func_types.into_boxed_slice(),
             globals: globals.into_boxed_slice(),
             table_types: table_types.into_boxed_slice(),
-            memory_types: reader.memory_types.into_boxed_slice(),
             imports: reader.imports.into_boxed_slice(),
+            version: reader.version,
+            start_func: reader.start_func,
             data: reader.data.into_boxed_slice(),
+            exports: reader.exports.into_boxed_slice(),
             elements: reader.elements.into_boxed_slice(),
+            memory_types: reader.memory_types.into_boxed_slice(),
         })
     }
 }
diff --git a/crates/parser/src/module.rs b/crates/parser/src/module.rs
index 660a702..811c51f 100644
--- a/crates/parser/src/module.rs
+++ b/crates/parser/src/module.rs
@@ -17,7 +17,10 @@ pub struct ModuleReader {
     pub start_func: Option<u32>,
 
     pub func_types: Vec<FuncType>,
-    pub func_addrs: Vec<u32>,
+
+    // map from local function index to type index
+    pub code_type_addrs: Vec<u32>,
+
     pub exports: Vec<Export>,
     pub code: Vec<CodeSection>,
     pub globals: Vec<Global>,
@@ -36,7 +39,7 @@ impl Debug for ModuleReader {
         f.debug_struct("ModuleReader")
             .field("version", &self.version)
             .field("func_types", &self.func_types)
-            .field("func_addrs", &self.func_addrs)
+            .field("func_addrs", &self.code_type_addrs)
             .field("code", &self.code)
             .field("exports", &self.exports)
             .field("globals", &self.globals)
@@ -88,13 +91,13 @@ impl ModuleReader {
                     .collect::<Result<Vec<FuncType>>>()?;
             }
             FunctionSection(reader) => {
-                if !self.func_addrs.is_empty() {
+                if !self.code_type_addrs.is_empty() {
                     return Err(ParseError::DuplicateSection("Function section".into()));
                 }
 
                 debug!("Found function section");
                 validator.function_section(&reader)?;
-                self.func_addrs = reader.into_iter().map(|f| Ok(f?)).collect::<Result<Vec<_>>>()?;
+                self.code_type_addrs = reader.into_iter().map(|f| Ok(f?)).collect::<Result<Vec<_>>>()?;
             }
             GlobalSection(reader) => {
                 if !self.globals.is_empty() {
@@ -148,9 +151,7 @@ impl ModuleReader {
                 debug!("Found code section entry");
                 let v = validator.code_section_entry(&function)?;
                 let func_validator = v.into_validator(Default::default());
-
-                self.code
-                    .push(conversion::convert_module_code(function, func_validator)?);
+                self.code.push(conversion::convert_module_code(function, func_validator)?);
             }
             ImportSection(reader) => {
                 if !self.imports.is_empty() {
@@ -168,10 +169,8 @@ impl ModuleReader {
 
                 debug!("Found export section");
                 validator.export_section(&reader)?;
-                self.exports = reader
-                    .into_iter()
-                    .map(|e| conversion::convert_module_export(e?))
-                    .collect::<Result<Vec<_>>>()?;
+                self.exports =
+                    reader.into_iter().map(|e| conversion::convert_module_export(e?)).collect::<Result<Vec<_>>>()?;
             }
             End(offset) => {
                 debug!("Reached end of module");
@@ -191,12 +190,7 @@ impl ModuleReader {
             //     validator.tag_section(&tag)?;
             // }
             UnknownSection { .. } => return Err(ParseError::UnsupportedSection("Unknown section".into())),
-            section => {
-                return Err(ParseError::UnsupportedSection(format!(
-                    "Unsupported section: {:?}",
-                    section
-                )))
-            }
+            section => return Err(ParseError::UnsupportedSection(format!("Unsupported section: {:?}", section))),
         };
 
         Ok(())