2 files changed, 62 insertions, 57 deletions

diff --git a/crates/parser/src/lib.rs b/crates/parser/src/lib.rs
index 0b5af2e..7b3dfa2 100644
--- a/crates/parser/src/lib.rs
+++ b/crates/parser/src/lib.rs
@@ -39,7 +39,7 @@ pub use error::*;
 use module::ModuleReader;
 use wasmparser::{Validator, WasmFeatures};
 
-pub use tinywasm_types::TinyWasmModule;
+pub use tinywasm_types::Module;
 
 /// Parser optimization and lowering options.
 #[non_exhaustive]
@@ -113,8 +113,8 @@ impl Parser {
         Validator::new_with_features(features)
     }
 
-    /// Parse a [`TinyWasmModule`] from bytes
-    pub fn parse_module_bytes(&self, wasm: impl AsRef<[u8]>) -> Result<TinyWasmModule> {
+    /// Parse a [`Module`] from bytes
+    pub fn parse_module_bytes(&self, wasm: impl AsRef<[u8]>) -> Result<Module> {
         let wasm = wasm.as_ref();
         let mut validator = Self::create_validator(self.options.clone());
         let mut reader = ModuleReader::default();
@@ -131,16 +131,16 @@ impl Parser {
     }
 
     #[cfg(feature = "std")]
-    /// Parse a [`TinyWasmModule`] from a file. Requires `std` feature.
-    pub fn parse_module_file(&self, path: impl AsRef<crate::std::path::Path> + Clone) -> Result<TinyWasmModule> {
+    /// Parse a [`Module`] from a file. Requires `std` feature.
+    pub fn parse_module_file(&self, path: impl AsRef<crate::std::path::Path> + Clone) -> Result<Module> {
         let file = crate::std::fs::File::open(&path)
             .map_err(|e| ParseError::Other(alloc::format!("Error opening file {:?}: {}", path.as_ref(), e)))?;
         self.parse_module_stream(&mut crate::std::io::BufReader::new(file))
     }
 
     #[cfg(feature = "std")]
-    /// Parse a [`TinyWasmModule`] from a stream. Requires `std` feature.
-    pub fn parse_module_stream(&self, mut stream: impl std::io::Read) -> Result<TinyWasmModule> {
+    /// Parse a [`Module`] from a stream. Requires `std` feature.
+    pub fn parse_module_stream(&self, mut stream: impl std::io::Read) -> Result<Module> {
         let mut validator = Self::create_validator(self.options.clone());
         let mut reader = ModuleReader::default();
         let mut buffer = alloc::vec::Vec::new();
@@ -170,10 +170,30 @@ impl Parser {
     }
 }
 
-impl TryFrom<ModuleReader> for TinyWasmModule {
+impl TryFrom<ModuleReader> for Module {
     type Error = ParseError;
 
     fn try_from(reader: ModuleReader) -> Result<Self> {
         reader.into_module(&ParserOptions::default())
     }
 }
+
+/// Parse a module from bytes
+pub fn parse_bytes(wasm: &[u8]) -> Result<Module> {
+    let data = Parser::new().parse_module_bytes(wasm)?;
+    Ok(data)
+}
+
+#[cfg(feature = "std")]
+/// Parse a module from a file. Requires the `std` feature.
+pub fn parse_file(path: impl AsRef<crate::std::path::Path> + Clone) -> Result<Module> {
+    let data = Parser::new().parse_module_file(path)?;
+    Ok(data)
+}
+
+#[cfg(feature = "std")]
+/// Parse a module from a stream. Requires `parser` and `std` features.
+pub fn parse_stream(stream: impl crate::std::io::Read) -> Result<Module> {
+    let data = Parser::new().parse_module_stream(stream)?;
+    Ok(data)
+}
diff --git a/crates/parser/src/module.rs b/crates/parser/src/module.rs
index 9c4d6b4..47257fe 100644
--- a/crates/parser/src/module.rs
+++ b/crates/parser/src/module.rs
@@ -1,5 +1,6 @@
 use crate::log::debug;
 use crate::{ParseError, ParserOptions, Result, conversion, optimize};
+use alloc::sync::Arc;
 use alloc::{format, string::ToString, vec::Vec};
 use tinywasm_types::*;
 use wasmparser::{FuncValidatorAllocations, Payload, Validator};
@@ -12,7 +13,7 @@ pub(crate) struct ModuleReader {
 
     pub(crate) version: Option<u16>,
     pub(crate) start_func: Option<u32>,
-    pub(crate) func_types: Vec<FuncType>,
+    pub(crate) func_types: Vec<Arc<FuncType>>,
     pub(crate) code_type_addrs: Vec<u32>,
     pub(crate) exports: Vec<Export>,
     pub(crate) code: Vec<Code>,
@@ -53,8 +54,8 @@ impl ModuleReader {
                 validator.type_section(&reader)?;
                 self.func_types = reader
                     .into_iter()
-                    .map(|t| conversion::convert_module_type(t?))
-                    .collect::<Result<Vec<FuncType>>>()?;
+                    .map(|t| conversion::convert_module_type(t?).map(Arc::new))
+                    .collect::<Result<Vec<Arc<FuncType>>>>()?;
             }
 
             Payload::GlobalSection(reader) => {
@@ -167,7 +168,7 @@ impl ModuleReader {
         Ok(())
     }
 
-    pub(crate) fn into_module(self, options: &ParserOptions) -> Result<TinyWasmModule> {
+    pub(crate) fn into_module(self, options: &ParserOptions) -> Result<Module> {
         if !self.end_reached {
             return Err(ParseError::EndNotReached);
         }
@@ -176,73 +177,57 @@ impl ModuleReader {
             return Err(ParseError::Other("Code and code type address count mismatch".to_string()));
         }
 
-        let Self {
-            start_func,
-            func_types,
-            code_type_addrs,
-            exports,
-            code,
-            globals,
-            table_types,
-            memory_types,
-            imports,
-            data,
-            elements,
-            ..
-        } = self;
-
-        let imported_func_count = imports.iter().filter(|i| matches!(&i.kind, ImportKind::Function(_))).count();
-        let imported_memory_count = imports.iter().filter(|i| matches!(&i.kind, ImportKind::Memory(_))).count() as u32;
-        let has_local_memory_export =
-            exports.iter().any(|export| export.kind == ExternalKind::Memory && export.index >= imported_memory_count);
-        let has_active_data_segment_on_local_memory = data.iter().any(|data| match &data.kind {
-            DataKind::Active { mem, .. } => *mem >= imported_memory_count,
+        let imported_func_count = self.imports.iter().filter(|i| matches!(&i.kind, ImportKind::Function(_))).count();
+        let import_mem_count = self.imports.iter().filter(|i| matches!(&i.kind, ImportKind::Memory(_))).count() as u32;
+        let has_local_mem_export =
+            self.exports.iter().any(|export| export.kind == ExternalKind::Memory && export.index >= import_mem_count);
+        let has_active_data_segment_on_local_memory = self.data.iter().any(|data| match &data.kind {
+            DataKind::Active { mem, .. } => *mem >= import_mem_count,
             DataKind::Passive => false,
         });
         let optimize_local_memory_allocation = options.optimize_local_memory_allocation();
-        let mut local_memory_allocation = if memory_types.is_empty() {
+        let mut local_memory_allocation = if self.memory_types.is_empty() {
             LocalMemoryAllocation::Skip
         } else if !optimize_local_memory_allocation || has_active_data_segment_on_local_memory {
             LocalMemoryAllocation::Eager
-        } else if has_local_memory_export {
+        } else if has_local_mem_export {
             LocalMemoryAllocation::Lazy
         } else {
             LocalMemoryAllocation::Skip
         };
-        let mut funcs = Vec::with_capacity(code.len());
+        let mut funcs = Vec::with_capacity(self.code.len());
 
-        for (func_idx, ((instructions, mut data, locals), ty_idx)) in code.into_iter().zip(code_type_addrs).enumerate()
+        for (func_idx, ((instructions, mut data, locals), ty_idx)) in
+            self.code.into_iter().zip(self.code_type_addrs).enumerate()
         {
-            let ty = func_types.get(ty_idx as usize).expect("No func type for func, this is a bug").clone();
+            let ty = self.func_types.get(ty_idx as usize).expect("No func type for func, this is a bug").clone();
             let params = ValueCounts::from_iter(ty.params());
             let self_func = (imported_func_count + func_idx) as u32;
-            let optimized = optimize::optimize_instructions(
-                instructions,
-                &mut data,
-                self_func,
-                imported_memory_count,
-                optimize_local_memory_allocation && local_memory_allocation != LocalMemoryAllocation::Eager,
-            );
+            let local_mem_alloc =
+                optimize_local_memory_allocation && local_memory_allocation != LocalMemoryAllocation::Eager;
+            let optimized =
+                optimize::optimize_instructions(instructions, &mut data, self_func, import_mem_count, local_mem_alloc);
 
             if optimized.uses_local_memory {
                 local_memory_allocation = LocalMemoryAllocation::Eager;
             }
 
-            funcs.push(WasmFunction { instructions: ArcSlice::from(optimized.instructions), data, locals, params, ty });
+            funcs.push(WasmFunction { instructions: optimized.instructions.into(), data, locals, params, ty }.into());
         }
 
-        Ok(TinyWasmModule {
+        Ok(ModuleInner {
             funcs: funcs.into(),
-            func_types: func_types.into(),
-            globals: globals.into(),
-            table_types: table_types.into(),
-            imports: imports.into(),
-            start_func,
-            data: data.into(),
-            exports: exports.into(),
-            elements: elements.into(),
-            memory_types: memory_types.into(),
+            func_types: self.func_types.into(),
+            globals: self.globals.into(),
+            table_types: self.table_types.into(),
+            imports: self.imports.into(),
+            start_func: self.start_func,
+            data: self.data.into(),
+            exports: self.exports.into(),
+            elements: self.elements.into(),
+            memory_types: self.memory_types.into(),
             local_memory_allocation,
-        })
+        }
+        .into())
     }
 }