pyodide_webassembly_runtime_layer/
module.rs

1use std::sync::Arc;
2
3use fxhash::FxHashMap;
4use pyo3::{prelude::*, sync::GILOnceCell};
5use wasm_runtime_layer::{
6    backend::WasmModule, ExportType, ExternType, FuncType, GlobalType, ImportType, MemoryType,
7    TableType, ValueType,
8};
9
10use crate::{
11    conversion::js_uint8_array_new, features::UnsupportedWasmFeatureExtensionError, Engine,
12};
13
14#[derive(Debug)]
15/// A WASM module.
16///
17/// This type wraps a [`WebAssembly.Module`] from the JavaScript API.
18///
19/// [`WebAssembly.Module`]: https://developer.mozilla.org/en-US/docs/WebAssembly/JavaScript_interface/Module
20pub struct Module {
21    /// The inner module
22    module: Py<PyAny>,
23    /// The parsed module, containing import and export signatures
24    parsed: Arc<ParsedModule>,
25}
26
27impl Clone for Module {
28    fn clone(&self) -> Self {
29        Python::with_gil(|py| Self {
30            module: self.module.clone_ref(py),
31            parsed: self.parsed.clone(),
32        })
33    }
34}
35
36impl WasmModule<Engine> for Module {
37    fn new(_engine: &Engine, bytes: &[u8]) -> anyhow::Result<Self> {
38        Python::with_gil(|py| {
39            #[cfg(feature = "tracing")]
40            let _span = tracing::debug_span!("Module::new").entered();
41
42            let parsed = ParsedModule::parse(bytes)?;
43
44            let buffer = js_uint8_array_new(py)?.call1((bytes,))?;
45
46            let module = match web_assembly_module_new(py)?.call1((buffer,)) {
47                Ok(module) => module,
48                // check if the error comes from missing feature support
49                // - if so, report the more informative unsupported feature error instead
50                // - if not, bubble up the error that made module instantiation fail
51                Err(err) => match Python::with_gil(|py| {
52                    UnsupportedWasmFeatureExtensionError::check_support(py, bytes)
53                })? {
54                    Ok(()) => anyhow::bail!(err),
55                    Err(unsupported) => anyhow::bail!(unsupported),
56                },
57            };
58
59            let parsed = Arc::new(parsed);
60
61            Ok(Self {
62                module: module.unbind(),
63                parsed,
64            })
65        })
66    }
67
68    fn exports(&self) -> Box<dyn '_ + Iterator<Item = ExportType<'_>>> {
69        Box::new(self.parsed.exports.iter().map(|(name, ty)| ExportType {
70            name: name.as_str(),
71            ty: ty.clone(),
72        }))
73    }
74
75    fn get_export(&self, name: &str) -> Option<ExternType> {
76        self.parsed.exports.get(name).cloned()
77    }
78
79    fn imports(&self) -> Box<dyn '_ + Iterator<Item = ImportType<'_>>> {
80        Box::new(
81            self.parsed
82                .imports
83                .iter()
84                .map(|((module, name), kind)| ImportType {
85                    module,
86                    name,
87                    ty: kind.clone(),
88                }),
89        )
90    }
91}
92
93impl Module {
94    pub(crate) fn module(&self, py: Python) -> Py<PyAny> {
95        self.module.clone_ref(py)
96    }
97}
98
99#[derive(Debug)]
100/// A parsed core module with imports and exports
101struct ParsedModule {
102    /// Import signatures
103    imports: FxHashMap<(String, String), ExternType>,
104    /// Export signatures
105    exports: FxHashMap<String, ExternType>,
106}
107
108impl ParsedModule {
109    #[allow(clippy::too_many_lines)]
110    /// Parses a module from bytes and extracts import and export signatures
111    fn parse(bytes: &[u8]) -> anyhow::Result<Self> {
112        let parser = wasmparser::Parser::new(0);
113
114        let mut imports = FxHashMap::default();
115        let mut exports = FxHashMap::default();
116
117        let mut types = Vec::new();
118
119        let mut functions = Vec::new();
120        let mut memories = Vec::new();
121        let mut tables = Vec::new();
122        let mut globals = Vec::new();
123
124        parser.parse_all(bytes).try_for_each(|payload| {
125            match payload? {
126                wasmparser::Payload::TypeSection(section) => {
127                    for ty in section {
128                        let ty = ty?;
129
130                        let mut subtypes = ty.types();
131                        let subtype = subtypes.next();
132
133                        let ty = match (subtype, subtypes.next()) {
134                            (Some(subtype), None) => match &subtype.composite_type.inner {
135                                wasmparser::CompositeInnerType::Func(func_type) => FuncType::new(
136                                    func_type
137                                        .params()
138                                        .iter()
139                                        .copied()
140                                        .map(ValueType::from_value),
141                                    func_type
142                                        .results()
143                                        .iter()
144                                        .copied()
145                                        .map(ValueType::from_value),
146                                ),
147                                _ => unreachable!(),
148                            },
149                            _ => unimplemented!(),
150                        };
151
152                        types.push(ty);
153                    }
154                },
155                wasmparser::Payload::FunctionSection(section) => {
156                    for type_index in section {
157                        let type_index = type_index?;
158
159                        let ty = &types[type_index as usize];
160
161                        functions.push(ty.clone());
162                    }
163                },
164                wasmparser::Payload::TableSection(section) => {
165                    for table in section {
166                        let table = table?;
167                        tables.push(TableType::from_parsed(&table.ty)?);
168                    }
169                },
170                wasmparser::Payload::MemorySection(section) => {
171                    for memory in section {
172                        let memory = memory?;
173                        memories.push(MemoryType::from_parsed(&memory)?);
174                    }
175                },
176                wasmparser::Payload::GlobalSection(section) => {
177                    for global in section {
178                        let global = global?;
179                        globals.push(GlobalType::from_parsed(global.ty));
180                    }
181                },
182                wasmparser::Payload::TagSection(section) => {
183                    for tag in section {
184                        let tag = tag?;
185
186                        #[cfg(feature = "tracing")]
187                        tracing::trace!(?tag, "tag");
188                        #[cfg(not(feature = "tracing"))]
189                        let _ = tag;
190                    }
191                },
192                wasmparser::Payload::ImportSection(section) => {
193                    for import in section {
194                        let import = import?;
195                        let ty = match import.ty {
196                            wasmparser::TypeRef::Func(index) => {
197                                let sig = types[index as usize].clone().with_name(import.name);
198                                functions.push(sig.clone());
199                                ExternType::Func(sig)
200                            },
201                            wasmparser::TypeRef::Table(ty) => {
202                                tables.push(TableType::from_parsed(&ty)?);
203                                ExternType::Table(TableType::from_parsed(&ty)?)
204                            },
205                            wasmparser::TypeRef::Memory(ty) => {
206                                memories.push(MemoryType::from_parsed(&ty)?);
207                                ExternType::Memory(MemoryType::from_parsed(&ty)?)
208                            },
209                            wasmparser::TypeRef::Global(ty) => {
210                                globals.push(GlobalType::from_parsed(ty));
211                                ExternType::Global(GlobalType::from_parsed(ty))
212                            },
213                            wasmparser::TypeRef::Tag(_) => {
214                                unimplemented!("WebAssembly.Tag is not yet supported")
215                            },
216                        };
217
218                        imports.insert((import.module.to_string(), import.name.to_string()), ty);
219                    }
220                },
221                wasmparser::Payload::ExportSection(section) => {
222                    for export in section {
223                        let export = export?;
224                        let index = export.index as usize;
225                        let ty = match export.kind {
226                            wasmparser::ExternalKind::Func => {
227                                ExternType::Func(functions[index].clone().with_name(export.name))
228                            },
229                            wasmparser::ExternalKind::Table => ExternType::Table(tables[index]),
230                            wasmparser::ExternalKind::Memory => ExternType::Memory(memories[index]),
231                            wasmparser::ExternalKind::Global => ExternType::Global(globals[index]),
232                            wasmparser::ExternalKind::Tag => {
233                                unimplemented!("WebAssembly.Tag is not yet supported")
234                            },
235                        };
236
237                        exports.insert(export.name.to_string(), ty);
238                    }
239                },
240                wasmparser::Payload::ElementSection(section) => {
241                    for element in section {
242                        let element = element?;
243
244                        #[cfg(feature = "tracing")]
245                        match element.kind {
246                            wasmparser::ElementKind::Passive => tracing::debug!("passive"),
247                            wasmparser::ElementKind::Active { .. } => tracing::debug!("active"),
248                            wasmparser::ElementKind::Declared => tracing::debug!("declared"),
249                        }
250                        #[cfg(not(feature = "tracing"))]
251                        let _ = element;
252                    }
253                },
254                _ => (),
255            }
256
257            anyhow::Ok(())
258        })?;
259
260        Ok(Self { imports, exports })
261    }
262}
263
264trait ValueTypeFrom {
265    fn from_value(value: wasmparser::ValType) -> Self;
266    fn from_ref(ty: wasmparser::RefType) -> Self;
267}
268
269impl ValueTypeFrom for ValueType {
270    fn from_value(value: wasmparser::ValType) -> Self {
271        match value {
272            wasmparser::ValType::I32 => Self::I32,
273            wasmparser::ValType::I64 => Self::I64,
274            wasmparser::ValType::F32 => Self::F32,
275            wasmparser::ValType::F64 => Self::F64,
276            wasmparser::ValType::V128 => unimplemented!("v128 is not supported"),
277            wasmparser::ValType::Ref(ty) => Self::from_ref(ty),
278        }
279    }
280
281    fn from_ref(ty: wasmparser::RefType) -> Self {
282        if ty.is_func_ref() {
283            Self::FuncRef
284        } else if ty.is_extern_ref() {
285            Self::ExternRef
286        } else {
287            unimplemented!("unsupported reference type {ty:?}")
288        }
289    }
290}
291
292trait TableTypeFrom: Sized {
293    fn from_parsed(value: &wasmparser::TableType) -> anyhow::Result<Self>;
294}
295
296impl TableTypeFrom for TableType {
297    fn from_parsed(value: &wasmparser::TableType) -> anyhow::Result<Self> {
298        Ok(Self::new(
299            ValueType::from_ref(value.element_type),
300            value.initial.try_into()?,
301            match value.maximum {
302                None => None,
303                Some(maximum) => Some(maximum.try_into()?),
304            },
305        ))
306    }
307}
308
309trait MemoryTypeFrom: Sized {
310    fn from_parsed(value: &wasmparser::MemoryType) -> anyhow::Result<Self>;
311}
312
313impl MemoryTypeFrom for MemoryType {
314    fn from_parsed(value: &wasmparser::MemoryType) -> anyhow::Result<Self> {
315        if value.memory64 {
316            anyhow::bail!("memory64 is not yet supported");
317        }
318
319        if value.shared {
320            anyhow::bail!("shared memory is not yet supported");
321        }
322
323        Ok(Self::new(
324            value.initial.try_into()?,
325            match value.maximum {
326                None => None,
327                Some(maximum) => Some(maximum.try_into()?),
328            },
329        ))
330    }
331}
332
333trait GlobalTypeFrom {
334    fn from_parsed(value: wasmparser::GlobalType) -> Self;
335}
336
337impl GlobalTypeFrom for GlobalType {
338    fn from_parsed(value: wasmparser::GlobalType) -> Self {
339        Self::new(ValueType::from_value(value.content_type), value.mutable)
340    }
341}
342
343fn web_assembly_module_new(py: Python) -> Result<&Bound<PyAny>, PyErr> {
344    static WEB_ASSEMBLY_MODULE: GILOnceCell<Py<PyAny>> = GILOnceCell::new();
345    WEB_ASSEMBLY_MODULE.import(py, "js.WebAssembly.Module", "new")
346}