3 files changed, 39 insertions, 12 deletions

diff --git a/crates/parser/src/conversion.rs b/crates/parser/src/conversion.rs
index d08df9b..82ceddd 100644
--- a/crates/parser/src/conversion.rs
+++ b/crates/parser/src/conversion.rs
@@ -1,5 +1,6 @@
 use alloc::{boxed::Box, format, string::ToString, vec::Vec};
 use tinywasm_types::{BlockArgs, Export, ExternalKind, FuncType, Instruction, MemArg, ValType};
+use wasmparser::{FuncValidator, ValidatorResources};
 
 use crate::{module::CodeSection, Result};
 
@@ -24,23 +25,28 @@ pub(crate) fn convert_module_export(export: wasmparser::Export) -> Result<Export
     })
 }
 
-pub(crate) fn convert_module_code(func: wasmparser::FunctionBody) -> Result<CodeSection> {
+pub(crate) fn convert_module_code(
+    func: wasmparser::FunctionBody,
+    mut validator: FuncValidator<ValidatorResources>,
+) -> Result<CodeSection> {
     let locals_reader = func.get_locals_reader()?;
     let count = locals_reader.get_count();
     let mut locals = Vec::with_capacity(count as usize);
-    locals.extend(
-        locals_reader
-            .into_iter()
-            .filter_map(|l| l.ok())
-            .map(|l| convert_valtype(&l.1)),
-    );
+
+    for (i, local) in locals_reader.into_iter().enumerate() {
+        let local = local?;
+        for _ in 0..local.0 {
+            locals.push(convert_valtype(&local.1));
+        }
+        validator.define_locals(i, local.0, local.1)?;
+    }
 
     if locals.len() != count as usize {
         return Err(crate::ParseError::Other("Invalid local index".to_string()));
     }
 
     let body_reader = func.get_operators_reader()?;
-    let body = process_operators(body_reader.into_iter())?;
+    let body = process_operators(body_reader.into_iter(), validator)?;
 
     Ok(CodeSection {
         locals: locals.into_boxed_slice(),
@@ -106,12 +112,18 @@ pub(crate) fn convert_memarg(memarg: wasmparser::MemArg) -> MemArg {
 
 pub fn process_operators<'a>(
     ops: impl Iterator<Item = Result<wasmparser::Operator<'a>, wasmparser::BinaryReaderError>>,
+    mut validator: FuncValidator<ValidatorResources>,
 ) -> Result<Box<[Instruction]>> {
     let mut instructions = Vec::new();
 
+    let mut offset = 0;
     for op in ops {
+        let op = op?;
+        validator.op(offset, &op)?;
+        offset += 1;
+
         use wasmparser::Operator::*;
-        let res = match op? {
+        let res = match op {
             BrTable { targets } => {
                 let def = targets.default();
                 let targets = targets
@@ -303,5 +315,7 @@ pub fn process_operators<'a>(
         instructions.push(res);
     }
 
+    validator.finish(offset)?;
+
     Ok(instructions.into_boxed_slice())
 }
diff --git a/crates/parser/src/lib.rs b/crates/parser/src/lib.rs
index 31e8cc8..343350a 100644
--- a/crates/parser/src/lib.rs
+++ b/crates/parser/src/lib.rs
@@ -5,6 +5,17 @@
 mod std;
 extern crate alloc;
 
+// log for logging (optional).
+#[cfg(feature = "logging")]
+#[allow(clippy::single_component_path_imports)]
+use log;
+
+#[cfg(not(feature = "logging"))]
+mod log {
+    macro_rules! debug    ( ($($tt:tt)*) => {{}} );
+    pub(crate) use debug;
+}
+
 mod conversion;
 mod error;
 mod module;
diff --git a/crates/parser/src/module.rs b/crates/parser/src/module.rs
index 5bdbcd8..69bb894 100644
--- a/crates/parser/src/module.rs
+++ b/crates/parser/src/module.rs
@@ -1,6 +1,6 @@
+use crate::log::debug;
 use alloc::{boxed::Box, format, vec::Vec};
 use core::fmt::Debug;
-use log::debug;
 use tinywasm_types::{Export, FuncType, Instruction, ValType};
 use wasmparser::{Payload, Validator};
 
@@ -124,9 +124,11 @@ impl ModuleReader {
             }
             CodeSectionEntry(function) => {
                 debug!("Found code section entry");
-                validator.code_section_entry(&function)?;
+                let v = validator.code_section_entry(&function)?;
+                let func_validator = v.into_validator(Default::default());
 
-                self.code_section.push(conversion::convert_module_code(function)?);
+                self.code_section
+                    .push(conversion::convert_module_code(function, func_validator)?);
             }
             ImportSection(_reader) => {
                 return Err(ParseError::UnsupportedSection("Import section".into()));