pyodide_webassembly_runtime_layer/
module.rs

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