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