Skip to main content

pyodide_webassembly_runtime_layer/
module.rs

1use std::sync::Arc;
2
3use fxhash::FxHashMap;
4use pyo3::{prelude::*, sync::PyOnceLock};
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::attach(|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::attach(|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::attach(|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.into_imports() {
194                        let import = import?;
195                        let ty = match import.ty {
196                            wasmparser::TypeRef::Func(index)
197                            | wasmparser::TypeRef::FuncExact(index) => {
198                                let sig = types[index as usize].clone().with_name(import.name);
199                                functions.push(sig.clone());
200                                ExternType::Func(sig)
201                            },
202                            wasmparser::TypeRef::Table(ty) => {
203                                tables.push(TableType::from_parsed(&ty)?);
204                                ExternType::Table(TableType::from_parsed(&ty)?)
205                            },
206                            wasmparser::TypeRef::Memory(ty) => {
207                                memories.push(MemoryType::from_parsed(&ty)?);
208                                ExternType::Memory(MemoryType::from_parsed(&ty)?)
209                            },
210                            wasmparser::TypeRef::Global(ty) => {
211                                globals.push(GlobalType::from_parsed(ty));
212                                ExternType::Global(GlobalType::from_parsed(ty))
213                            },
214                            wasmparser::TypeRef::Tag(_) => {
215                                unimplemented!("WebAssembly.Tag is not yet supported")
216                            },
217                        };
218
219                        imports.insert((import.module.to_string(), import.name.to_string()), ty);
220                    }
221                },
222                wasmparser::Payload::ExportSection(section) => {
223                    for export in section {
224                        let export = export?;
225                        let index = export.index as usize;
226                        let ty = match export.kind {
227                            wasmparser::ExternalKind::Func
228                            | wasmparser::ExternalKind::FuncExact => {
229                                ExternType::Func(functions[index].clone().with_name(export.name))
230                            },
231                            wasmparser::ExternalKind::Table => ExternType::Table(tables[index]),
232                            wasmparser::ExternalKind::Memory => ExternType::Memory(memories[index]),
233                            wasmparser::ExternalKind::Global => ExternType::Global(globals[index]),
234                            wasmparser::ExternalKind::Tag => {
235                                unimplemented!("WebAssembly.Tag is not yet supported")
236                            },
237                        };
238
239                        exports.insert(export.name.to_string(), ty);
240                    }
241                },
242                wasmparser::Payload::ElementSection(section) => {
243                    for element in section {
244                        let element = element?;
245
246                        #[cfg(feature = "tracing")]
247                        match element.kind {
248                            wasmparser::ElementKind::Passive => tracing::debug!("passive"),
249                            wasmparser::ElementKind::Active { .. } => tracing::debug!("active"),
250                            wasmparser::ElementKind::Declared => tracing::debug!("declared"),
251                        }
252                        #[cfg(not(feature = "tracing"))]
253                        let _ = element;
254                    }
255                },
256                _ => (),
257            }
258
259            anyhow::Ok(())
260        })?;
261
262        Ok(Self { imports, exports })
263    }
264}
265
266trait ValueTypeFrom {
267    fn from_value(value: wasmparser::ValType) -> Self;
268    fn from_ref(ty: wasmparser::RefType) -> Self;
269}
270
271impl ValueTypeFrom for ValueType {
272    fn from_value(value: wasmparser::ValType) -> Self {
273        match value {
274            wasmparser::ValType::I32 => Self::I32,
275            wasmparser::ValType::I64 => Self::I64,
276            wasmparser::ValType::F32 => Self::F32,
277            wasmparser::ValType::F64 => Self::F64,
278            wasmparser::ValType::V128 => unimplemented!("v128 is not supported"),
279            wasmparser::ValType::Ref(ty) => Self::from_ref(ty),
280        }
281    }
282
283    fn from_ref(ty: wasmparser::RefType) -> Self {
284        if ty.is_func_ref() {
285            Self::FuncRef
286        } else if ty.is_extern_ref() {
287            Self::ExternRef
288        } else {
289            unimplemented!("unsupported reference type {ty:?}")
290        }
291    }
292}
293
294trait TableTypeFrom: Sized {
295    fn from_parsed(value: &wasmparser::TableType) -> anyhow::Result<Self>;
296}
297
298impl TableTypeFrom for TableType {
299    fn from_parsed(value: &wasmparser::TableType) -> anyhow::Result<Self> {
300        Ok(Self::new(
301            ValueType::from_ref(value.element_type),
302            value.initial.try_into()?,
303            match value.maximum {
304                None => None,
305                Some(maximum) => Some(maximum.try_into()?),
306            },
307        ))
308    }
309}
310
311trait MemoryTypeFrom: Sized {
312    fn from_parsed(value: &wasmparser::MemoryType) -> anyhow::Result<Self>;
313}
314
315impl MemoryTypeFrom for MemoryType {
316    fn from_parsed(value: &wasmparser::MemoryType) -> anyhow::Result<Self> {
317        if value.memory64 {
318            anyhow::bail!("memory64 is not yet supported");
319        }
320
321        if value.shared {
322            anyhow::bail!("shared memory is not yet supported");
323        }
324
325        Ok(Self::new(
326            value.initial.try_into()?,
327            match value.maximum {
328                None => None,
329                Some(maximum) => Some(maximum.try_into()?),
330            },
331        ))
332    }
333}
334
335trait GlobalTypeFrom {
336    fn from_parsed(value: wasmparser::GlobalType) -> Self;
337}
338
339impl GlobalTypeFrom for GlobalType {
340    fn from_parsed(value: wasmparser::GlobalType) -> Self {
341        Self::new(ValueType::from_value(value.content_type), value.mutable)
342    }
343}
344
345fn web_assembly_module_new(py: Python<'_>) -> Result<&Bound<'_, PyAny>, PyErr> {
346    static WEB_ASSEMBLY_MODULE: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
347    WEB_ASSEMBLY_MODULE.import(py, "js.WebAssembly.Module", "new")
348}