Skip to content

Commit 34ceee9

Browse files
coolreader18coolreader18
committed
Refactor wasm function storage
1 parent a34951d commit 34ceee9

File tree

5 files changed

+87
-116
lines changed

5 files changed

+87
-116
lines changed

‎Cargo.lock

Lines changed: 10 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎wasm/lib/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ serde-wasm-bindgen = "0.1"
2525
serde = "1.0"
2626
js-sys = "0.3"
2727
futures = "0.1"
28+
generational-arena = "0.2"
2829

2930
[dependencies.web-sys]
3031
version = "0.3"

‎wasm/lib/src/convert.rs

Lines changed: 38 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
use generational_arena::Arena;
12
use std::cell::RefCell;
23

34
use js_sys::{Array, ArrayBuffer, Object, Promise, Reflect, SyntaxError, Uint8Array};
@@ -16,7 +17,29 @@ use crate::browser_module;
1617
use crate::vm_class::{stored_vm_from_wasm, WASMVirtualMachine};
1718

1819
// Currently WASM do not support multithreading. We should change this once it is enabled.
19-
thread_local!(static JS_FUNCS: RefCell<Vec<js_sys::Function>> = RefCell::new(vec![]));
20+
thread_local!(static JS_HANDLES: RefCell<Arena<JsValue>> = RefCell::new(Arena::new()));
21+
22+
pub struct JsHandle(generational_arena::Index);
23+
impl JsHandle {
24+
pub fn new(js: JsValue) -> Self {
25+
let idx = JS_HANDLES.with(|arena| arena.borrow_mut().insert(js));
26+
JsHandle(idx)
27+
}
28+
pub fn get(&self) -> JsValue {
29+
JS_HANDLES.with(|arena| {
30+
arena
31+
.borrow()
32+
.get(self.0)
33+
.expect("index was removed")
34+
.clone()
35+
})
36+
}
37+
}
38+
impl Drop for JsHandle {
39+
fn drop(&mut self) {
40+
JS_HANDLES.with(|arena| arena.borrow_mut().remove(self.0));
41+
}
42+
}
2043

2144
#[wasm_bindgen(inline_js = r"
2245
export class PyError extends Error {
@@ -195,32 +218,22 @@ pub fn js_to_py(vm: &VirtualMachine, js_val: JsValue) -> PyObjectRef {
195218
dict.into_object()
196219
}
197220
} else if js_val.is_function() {
198-
let func = js_sys::Function::from(js_val);
199-
let idx = JS_FUNCS.with(|funcs| {
200-
let mut funcs = funcs.borrow_mut();
201-
funcs.push(func);
202-
funcs.len() - 1
203-
});
221+
let js_handle = JsHandle::new(js_val);
204222
vm.ctx
205223
.new_method(move |vm: &VirtualMachine, args: PyFuncArgs| -> PyResult {
206-
JS_FUNCS.with(|funcs| {
207-
let this = Object::new();
208-
for (k, v) in args.kwargs {
209-
Reflect::set(&this, &k.into(), &py_to_js(vm, v))
210-
.expect("property to be settable");
211-
}
212-
let js_args = Array::new();
213-
for v in args.args {
214-
js_args.push(&py_to_js(vm, v));
215-
}
216-
funcs
217-
.borrow()
218-
.get(idx)
219-
.unwrap()
220-
.apply(&this, &js_args)
221-
.map(|val| js_to_py(vm, val))
222-
.map_err(|err| js_err_to_py_err(vm, &err))
223-
})
224+
let this = Object::new();
225+
for (k, v) in args.kwargs {
226+
Reflect::set(&this, &k.into(), &py_to_js(vm, v))
227+
.expect("property to be settable");
228+
}
229+
let js_args = Array::new();
230+
for v in args.args {
231+
js_args.push(&py_to_js(vm, v));
232+
}
233+
let func = js_sys::Function::from(js_handle.get());
234+
func.apply(&this, &js_args)
235+
.map(|val| js_to_py(vm, val))
236+
.map_err(|err| js_err_to_py_err(vm, &err))
224237
})
225238
} else if let Some(err) = js_val.dyn_ref::<js_sys::Error>() {
226239
js_err_to_py_err(vm, err).into_object()

‎wasm/lib/src/vm_class.rs

Lines changed: 15 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -6,15 +6,12 @@ use js_sys::{Object, TypeError};
66
use wasm_bindgen::prelude::*;
77

88
use rustpython_compiler::compile;
9-
use rustpython_vm::function::PyFuncArgs;
10-
use rustpython_vm::pyobject::{
11-
ItemProtocol, PyObject, PyObjectPayload, PyObjectRef, PyResult, PyValue,
12-
};
9+
use rustpython_vm::pyobject::{ItemProtocol, PyObject, PyObjectPayload, PyObjectRef, PyValue};
1310
use rustpython_vm::scope::{NameProtocol, Scope};
1411
use rustpython_vm::{InitParameter, PySettings, VirtualMachine};
1512

1613
use crate::browser_module::setup_browser_module;
17-
use crate::convert::{self, PyResultExt};
14+
use crate::convert::{self, JsHandle, PyResultExt};
1815
use crate::js_module;
1916
use crate::wasm_builtins;
2017
use rustpython_compiler::mode::Mode;
@@ -67,7 +64,6 @@ impl StoredVirtualMachine {
6764
// https://rustwasm.github.io/2018/10/24/multithreading-rust-and-wasm.html#atomic-instructions
6865
thread_local! {
6966
static STORED_VMS: RefCell<HashMap<String, Rc<StoredVirtualMachine>>> = RefCell::default();
70-
static JS_PRINT_FUNC: RefCell<Option<js_sys::Function>> = RefCell::new(None);
7167
}
7268

7369
pub fn get_vm_id(vm: &VirtualMachine) -> &str {
@@ -223,39 +219,28 @@ impl WASMVirtualMachine {
223219
fn error() -> JsValue {
224220
TypeError::new("Unknown stdout option, please pass a function or 'console'").into()
225221
}
226-
let print_fn: PyObjectRef = if let Some(s) = stdout.as_string() {
222+
use wasm_builtins::make_stdout_object;
223+
let stdout: PyObjectRef = if let Some(s) = stdout.as_string() {
227224
match s.as_str() {
228-
"console" => vm.ctx.new_method(wasm_builtins::builtin_print_console),
225+
"console" => make_stdout_object(vm, wasm_builtins::sys_stdout_write_console),
229226
_ => return Err(error()),
230227
}
231228
} else if stdout.is_function() {
232-
let func = js_sys::Function::from(stdout);
233-
JS_PRINT_FUNC.with(|thread_func| thread_func.replace(Some(func.clone())));
234-
vm.ctx
235-
.new_method(move |vm: &VirtualMachine, args: PyFuncArgs| -> PyResult {
236-
JS_PRINT_FUNC.with(|func| {
237-
func.borrow()
238-
.as_ref()
239-
.unwrap()
240-
.call1(
241-
&JsValue::UNDEFINED,
242-
&wasm_builtins::format_print_args(vm, args)?.into(),
243-
)
244-
.map_err(|err| convert::js_py_typeerror(vm, err))?;
245-
Ok(vm.get_none())
246-
})
247-
})
229+
let func_handle = JsHandle::new(stdout);
230+
make_stdout_object(vm, move |data, vm| {
231+
let func = js_sys::Function::from(func_handle.get());
232+
func.call1(&JsValue::UNDEFINED, &data.into())
233+
.map_err(|err| convert::js_py_typeerror(vm, err))?;
234+
Ok(())
235+
})
248236
} else if stdout.is_null() {
249-
fn noop(vm: &VirtualMachine, _args: PyFuncArgs) -> PyResult {
250-
Ok(vm.get_none())
251-
}
252-
vm.ctx.new_method(noop)
237+
make_stdout_object(vm, |_, _| Ok(()))
253238
} else if stdout.is_undefined() {
254-
vm.ctx.new_method(wasm_builtins::builtin_print_console)
239+
make_stdout_object(vm, wasm_builtins::sys_stdout_write_console)
255240
} else {
256241
return Err(error());
257242
};
258-
vm.set_attr(&vm.builtins, "print", print_fn).unwrap();
243+
vm.set_attr(&vm.sys_module, "stdout", stdout).unwrap();
259244
Ok(())
260245
})?
261246
}

‎wasm/lib/src/wasm_builtins.rs

Lines changed: 23 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -4,75 +4,37 @@
44
//! desktop.
55
//! Implements functions listed here: https://docs.python.org/3/library/builtins.html.
66
7-
use js_sys::{self, Array};
87
use web_sys::{self, console};
98

10-
use rustpython_vm::function::PyFuncArgs;
11-
use rustpython_vm::obj::{objstr, objtype};
12-
use rustpython_vm::pyobject::{IdProtocol, PyResult, TypeProtocol};
9+
use rustpython_vm::obj::objstr::PyStringRef;
10+
use rustpython_vm::pyobject::{PyObjectRef, PyResult};
1311
use rustpython_vm::VirtualMachine;
1412

1513
pub(crate) fn window() -> web_sys::Window {
1614
web_sys::window().expect("Window to be available")
1715
}
1816

19-
pub fn format_print_args(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult<String> {
20-
// Handle 'sep' kwarg:
21-
let sep_arg = args
22-
.get_optional_kwarg("sep")
23-
.filter(|obj| !obj.is(&vm.get_none()));
24-
if let Some(ref obj) = sep_arg {
25-
if !objtype::isinstance(obj, &vm.ctx.str_type()) {
26-
return Err(vm.new_type_error(format!(
27-
"sep must be None or a string, not {}",
28-
obj.class().name
29-
)));
30-
}
31-
}
32-
let sep_str = sep_arg.as_ref().map(|obj| objstr::borrow_value(obj));
33-
34-
// Handle 'end' kwarg:
35-
let end_arg = args
36-
.get_optional_kwarg("end")
37-
.filter(|obj| !obj.is(&vm.get_none()));
38-
if let Some(ref obj) = end_arg {
39-
if !objtype::isinstance(obj, &vm.ctx.str_type()) {
40-
return Err(vm.new_type_error(format!(
41-
"end must be None or a string, not {}",
42-
obj.class().name
43-
)));
44-
}
45-
}
46-
let end_str = end_arg.as_ref().map(|obj| objstr::borrow_value(obj));
47-
48-
// No need to handle 'flush' kwarg, irrelevant when writing to String
49-
50-
let mut output = String::new();
51-
let mut first = true;
52-
for a in args.args {
53-
if first {
54-
first = false;
55-
} else if let Some(ref sep_str) = sep_str {
56-
output.push_str(sep_str);
57-
} else {
58-
output.push(' ');
59-
}
60-
output.push_str(&vm.to_pystr(&a)?);
61-
}
62-
63-
if let Some(end_str) = end_str {
64-
output.push_str(end_str.as_ref())
65-
} else {
66-
output.push('\n');
67-
}
68-
Ok(output)
17+
pub fn sys_stdout_write_console(data: &str, _vm: &VirtualMachine) -> PyResult<()> {
18+
console::log_1(&data.into());
19+
Ok(())
6920
}
7021

71-
pub fn builtin_print_console(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
72-
let arr = Array::new();
73-
for arg in args.args {
74-
arr.push(&vm.to_pystr(&arg)?.into());
75-
}
76-
console::log(&arr);
77-
Ok(vm.get_none())
22+
pub fn make_stdout_object(
23+
vm: &VirtualMachine,
24+
write_f: impl Fn(&str, &VirtualMachine) -> PyResult<()> + Send + Sync + 'static,
25+
) -> PyObjectRef {
26+
let ctx = &vm.ctx;
27+
let write_method = ctx.new_method(
28+
move |_self: PyObjectRef, data: PyStringRef, vm: &VirtualMachine| -> PyResult<()> {
29+
write_f(data.as_str(), vm)
30+
},
31+
);
32+
let flush_method = ctx.new_method(|_self: PyObjectRef| {});
33+
// there's not really any point to storing this class so that there's a consistent type object,
34+
// we just want a half-decent repr() output
35+
let cls = py_class!(ctx, "JSStdout", vm.ctx.object(), {
36+
"write" => write_method,
37+
"flush" => flush_method,
38+
});
39+
ctx.new_base_object(cls, None)
7840
}

0 commit comments

Comments
 (0)