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