_wasmtime/
lib.rs

1use crate::function::{wrap_into_pyfunction, Function};
2use crate::instance::Instance;
3use crate::memory::Memory;
4use crate::module::Module;
5use pyo3::exceptions::Exception;
6use pyo3::prelude::*;
7use pyo3::types::{PyAny, PyBytes, PyDict, PySet};
8use pyo3::wrap_pyfunction;
9
10mod function;
11mod instance;
12mod memory;
13mod module;
14mod value;
15
16fn err2py(err: anyhow::Error) -> PyErr {
17    PyErr::new::<Exception, _>(format!("{:?}", err))
18}
19
20#[pyclass]
21pub struct InstantiateResultObject {
22    instance: Py<Instance>,
23    module: Py<Module>,
24}
25
26#[pymethods]
27impl InstantiateResultObject {
28    #[getter(instance)]
29    fn get_instance(&self) -> PyResult<Py<Instance>> {
30        let gil = Python::acquire_gil();
31        let py = gil.python();
32        Ok(self.instance.clone_ref(py))
33    }
34
35    #[getter(module)]
36    fn get_module(&self) -> PyResult<Py<Module>> {
37        let gil = Python::acquire_gil();
38        let py = gil.python();
39        Ok(self.module.clone_ref(py))
40    }
41}
42
43fn find_export_in(obj: &PyAny, store: &wasmtime::Store, name: &str) -> PyResult<wasmtime::Extern> {
44    let obj = obj.cast_as::<PyDict>()?;
45
46    Ok(if let Some(item) = obj.get_item(name) {
47        if item.is_callable() {
48            if item.get_type().is_subclass::<Function>()? {
49                let wasm_fn = item.cast_as::<Function>()?;
50                wasm_fn.func().into()
51            } else {
52                wrap_into_pyfunction(store, item)?.into()
53            }
54        } else if item.get_type().is_subclass::<Memory>()? {
55            let wasm_mem = item.cast_as::<Memory>()?;
56            wasm_mem.memory.clone().into()
57        } else {
58            return Err(PyErr::new::<Exception, _>(format!(
59                "unsupported import type {}",
60                name
61            )));
62        }
63    } else {
64        return Err(PyErr::new::<Exception, _>(format!(
65            "import {} is not found",
66            name
67        )));
68    })
69}
70
71/// WebAssembly instantiate API method.
72#[pyfunction]
73pub fn instantiate(
74    py: Python,
75    buffer_source: &PyBytes,
76    import_obj: &PyDict,
77) -> PyResult<Py<InstantiateResultObject>> {
78    let wasm_data = buffer_source.as_bytes();
79
80    let engine = wasmtime::Engine::new(&wasmtime::Config::new().wasm_multi_value(true));
81    let store = wasmtime::Store::new(&engine);
82
83    let module = wasmtime::Module::new(&store, wasm_data).map_err(err2py)?;
84
85    // If this module expects to be able to use wasi then go ahead and hook
86    // that up into the imported crates.
87    let cx = wasmtime_wasi::WasiCtxBuilder::new()
88        .build()
89        .map_err(|e| err2py(e.into()))?;
90    let wasi_snapshot_preview1 = wasmtime_wasi::Wasi::new(&store, cx);
91    let cx = wasmtime_wasi::old::snapshot_0::WasiCtxBuilder::new()
92        .build()
93        .map_err(|e| err2py(e.into()))?;
94    let wasi_snapshot = wasmtime_wasi::old::snapshot_0::Wasi::new(&store, cx);
95
96    let mut imports: Vec<wasmtime::Extern> = Vec::new();
97    for i in module.imports() {
98        if i.module() == "wasi_snapshot" {
99            if let Some(func) = wasi_snapshot.get_export(i.name()) {
100                imports.push(func.clone().into());
101                continue;
102            }
103        }
104        if i.module() == "wasi_snapshot_preview1" {
105            if let Some(func) = wasi_snapshot_preview1.get_export(i.name()) {
106                imports.push(func.clone().into());
107                continue;
108            }
109        }
110        let module_name = i.module();
111        if let Some(m) = import_obj.get_item(module_name) {
112            let e = find_export_in(m, &store, i.name())?;
113            imports.push(e);
114        } else {
115            return Err(PyErr::new::<Exception, _>(format!(
116                "imported module {} is not found",
117                module_name
118            )));
119        }
120    }
121
122    let instance = wasmtime::Instance::new(&module, &imports)
123        .map_err(|t| PyErr::new::<Exception, _>(format!("instantiated with trap {:?}", t)))?;
124
125    let module = Py::new(py, Module { module })?;
126
127    let instance = Py::new(py, Instance { instance })?;
128
129    Py::new(py, InstantiateResultObject { instance, module })
130}
131
132#[pyfunction]
133pub fn imported_modules<'p>(py: Python<'p>, buffer_source: &PyBytes) -> PyResult<&'p PyDict> {
134    let wasm_data = buffer_source.as_bytes();
135    let dict = PyDict::new(py);
136    // TODO: error handling
137    let mut parser = wasmparser::ModuleReader::new(wasm_data).unwrap();
138    while !parser.eof() {
139        let section = parser.read().unwrap();
140        match section.code {
141            wasmparser::SectionCode::Import => {}
142            _ => continue,
143        };
144        let reader = section.get_import_section_reader().unwrap();
145        for import in reader {
146            let import = import.unwrap();
147            // Skip over wasi-looking imports since those aren't imported from
148            // Python but rather they're implemented natively.
149            if wasmtime_wasi::is_wasi_module(import.module) {
150                continue;
151            }
152            let set = match dict.get_item(import.module) {
153                Some(set) => set.downcast_ref::<PySet>().unwrap(),
154                None => {
155                    let set = PySet::new::<PyObject>(py, &[])?;
156                    dict.set_item(import.module, set)?;
157                    set
158                }
159            };
160            set.add(import.field)?;
161        }
162    }
163    Ok(dict)
164}
165
166#[pymodule]
167fn lib_wasmtime(_: Python, m: &PyModule) -> PyResult<()> {
168    m.add_class::<Instance>()?;
169    m.add_class::<Memory>()?;
170    m.add_class::<Module>()?;
171    m.add_class::<InstantiateResultObject>()?;
172    m.add_wrapped(wrap_pyfunction!(instantiate))?;
173    m.add_wrapped(wrap_pyfunction!(imported_modules))?;
174    Ok(())
175}