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#[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 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 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 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}