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 tables = Vec::new();
121        let mut memories = 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.into_iter_err_on_gc_types() {
128                        types.push(Type::Func(Partial::Lazy(ty?)));
129                    }
130                },
131                wasmparser::Payload::ImportSection(section) => {
132                    for import in section.into_imports() {
133                        let import = import?;
134                        let ty = match import.ty {
135                            wasmparser::TypeRef::Func(index)
136                            | wasmparser::TypeRef::FuncExact(index) => {
137                                let Type::Func(ty) = &mut types[index as usize];
138                                let ty = ty.force(|ty| FuncType::from_parsed(ty))?;
139                                let func = ty.clone().with_name(import.name);
140                                functions.push(Partial::Eager(func.clone()));
141                                ExternType::Func(func)
142                            },
143                            wasmparser::TypeRef::Table(ty) => {
144                                let table = TableType::from_parsed(&ty)?;
145                                tables.push(Partial::Eager(table));
146                                ExternType::Table(table)
147                            },
148                            wasmparser::TypeRef::Memory(ty) => {
149                                let memory = MemoryType::from_parsed(&ty)?;
150                                memories.push(Partial::Eager(memory));
151                                ExternType::Memory(memory)
152                            },
153                            wasmparser::TypeRef::Global(ty) => {
154                                let global = GlobalType::from_parsed(&ty)?;
155                                globals.push(Partial::Eager(global));
156                                ExternType::Global(global)
157                            },
158                            wasmparser::TypeRef::Tag(_) => {
159                                anyhow::bail!(
160                                    "tag imports are not yet supported in the wasm_runtime_layer"
161                                )
162                            },
163                        };
164
165                        imports.insert((import.module.to_string(), import.name.to_string()), ty);
166                    }
167                },
168                wasmparser::Payload::FunctionSection(section) => {
169                    for type_index in section {
170                        let type_index = type_index?;
171                        let Type::Func(ty) = &types[type_index as usize];
172                        functions.push(ty.clone());
173                    }
174                },
175                wasmparser::Payload::TableSection(section) => {
176                    for table in section {
177                        tables.push(Partial::Lazy(table?.ty));
178                    }
179                },
180                wasmparser::Payload::MemorySection(section) => {
181                    for memory in section {
182                        memories.push(Partial::Lazy(memory?));
183                    }
184                },
185                wasmparser::Payload::GlobalSection(section) => {
186                    for global in section {
187                        globals.push(Partial::Lazy(global?.ty));
188                    }
189                },
190                wasmparser::Payload::ExportSection(section) => {
191                    for export in section {
192                        let export = export?;
193                        let index = export.index as usize;
194                        let ty = match export.kind {
195                            wasmparser::ExternalKind::Func
196                            | wasmparser::ExternalKind::FuncExact => {
197                                let ty = functions[index].force(|ty| FuncType::from_parsed(ty))?;
198                                let func = ty.clone().with_name(export.name);
199                                ExternType::Func(func)
200                            },
201                            wasmparser::ExternalKind::Table => {
202                                let table = tables[index].force(|ty| TableType::from_parsed(ty))?;
203                                ExternType::Table(*table)
204                            },
205                            wasmparser::ExternalKind::Memory => {
206                                let memory =
207                                    memories[index].force(|ty| MemoryType::from_parsed(ty))?;
208                                ExternType::Memory(*memory)
209                            },
210                            wasmparser::ExternalKind::Global => {
211                                let global =
212                                    globals[index].force(|ty| GlobalType::from_parsed(ty))?;
213                                ExternType::Global(*global)
214                            },
215                            wasmparser::ExternalKind::Tag => {
216                                anyhow::bail!(
217                                    "tag exports are not yet supported in the wasm_runtime_layer"
218                                )
219                            },
220                        };
221
222                        exports.insert(export.name.to_string(), ty);
223                    }
224                },
225                _ => (),
226            }
227
228            anyhow::Ok(())
229        })?;
230
231        Ok(Self { imports, exports })
232    }
233}
234
235enum Type<T> {
236    Func(T),
237}
238
239#[derive(Clone)]
240enum Partial<L, E> {
241    Lazy(L),
242    Eager(E),
243}
244
245impl<L, E> Partial<L, E> {
246    fn force(&mut self, eval: impl FnOnce(&mut L) -> anyhow::Result<E>) -> anyhow::Result<&mut E> {
247        match self {
248            Self::Eager(x) => Ok(x),
249            Self::Lazy(x) => {
250                *self = Self::Eager(eval(x)?);
251                // Safety: we have the only mutable reference and have just
252                //         overridden the variant to Self::Eager(...)
253                let Self::Eager(x) = self else {
254                    unsafe { std::hint::unreachable_unchecked() }
255                };
256                Ok(x)
257            },
258        }
259    }
260}
261
262trait ValueTypeFrom: Sized {
263    fn from_value(value: wasmparser::ValType) -> anyhow::Result<Self>;
264    fn from_ref(ty: wasmparser::RefType) -> anyhow::Result<Self>;
265}
266
267impl ValueTypeFrom for ValueType {
268    fn from_value(value: wasmparser::ValType) -> anyhow::Result<Self> {
269        match value {
270            wasmparser::ValType::I32 => Ok(Self::I32),
271            wasmparser::ValType::I64 => Ok(Self::I64),
272            wasmparser::ValType::F32 => Ok(Self::F32),
273            wasmparser::ValType::F64 => Ok(Self::F64),
274            wasmparser::ValType::V128 => {
275                anyhow::bail!("v128 is not yet supported in the wasm_runtime_layer")
276            },
277            wasmparser::ValType::Ref(ty) => Self::from_ref(ty),
278        }
279    }
280
281    fn from_ref(ty: wasmparser::RefType) -> anyhow::Result<Self> {
282        if ty.is_func_ref() {
283            Ok(Self::FuncRef)
284        } else if ty.is_extern_ref() {
285            Ok(Self::ExternRef)
286        } else {
287            anyhow::bail!("reference type {ty:?} is not yet supported in the wasm_runtime_layer")
288        }
289    }
290}
291
292trait FuncTypeFrom: Sized {
293    fn from_parsed(value: &wasmparser::FuncType) -> anyhow::Result<Self>;
294}
295
296impl FuncTypeFrom for FuncType {
297    fn from_parsed(value: &wasmparser::FuncType) -> anyhow::Result<Self> {
298        let params = value
299            .params()
300            .iter()
301            .copied()
302            .map(ValueType::from_value)
303            .collect::<anyhow::Result<Vec<_>>>()?;
304        let results = value
305            .results()
306            .iter()
307            .copied()
308            .map(ValueType::from_value)
309            .collect::<anyhow::Result<Vec<_>>>()?;
310
311        Ok(Self::new(params, results))
312    }
313}
314
315trait TableTypeFrom: Sized {
316    fn from_parsed(value: &wasmparser::TableType) -> anyhow::Result<Self>;
317}
318
319impl TableTypeFrom for TableType {
320    fn from_parsed(value: &wasmparser::TableType) -> anyhow::Result<Self> {
321        Ok(Self::new(
322            ValueType::from_ref(value.element_type)?,
323            value.initial.try_into()?,
324            match value.maximum {
325                None => None,
326                Some(maximum) => Some(maximum.try_into()?),
327            },
328        ))
329    }
330}
331
332trait MemoryTypeFrom: Sized {
333    fn from_parsed(value: &wasmparser::MemoryType) -> anyhow::Result<Self>;
334}
335
336impl MemoryTypeFrom for MemoryType {
337    fn from_parsed(value: &wasmparser::MemoryType) -> anyhow::Result<Self> {
338        if value.memory64 {
339            anyhow::bail!("memory64 is not yet supported in the wasm_runtime_layer");
340        }
341
342        if value.shared {
343            anyhow::bail!("shared memory is not yet supported in the wasm_runtime_layer");
344        }
345
346        Ok(Self::new(
347            value.initial.try_into()?,
348            match value.maximum {
349                None => None,
350                Some(maximum) => Some(maximum.try_into()?),
351            },
352        ))
353    }
354}
355
356trait GlobalTypeFrom: Sized {
357    fn from_parsed(value: &wasmparser::GlobalType) -> anyhow::Result<Self>;
358}
359
360impl GlobalTypeFrom for GlobalType {
361    fn from_parsed(value: &wasmparser::GlobalType) -> anyhow::Result<Self> {
362        Ok(Self::new(
363            ValueType::from_value(value.content_type)?,
364            value.mutable,
365        ))
366    }
367}
368
369fn web_assembly_module_new(py: Python<'_>) -> Result<&Bound<'_, PyAny>, PyErr> {
370    static WEB_ASSEMBLY_MODULE: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
371    WEB_ASSEMBLY_MODULE.import(py, "js.WebAssembly.Module", "new")
372}