Allow closures to be passed from Python to JS if using a WASM VM

windelbouwman · Feb 23, 2019 · Jan 30, 2019 · Jan 30, 2019 · Jan 30, 2019 · Jan 30, 2019
commit 9b6516e1756147dcb5abcc427100ec03f8e426f7
diff --git a/vm/src/builtins.rs b/vm/src/builtins.rs
@@ -18,6 +18,8 @@ use super::pyobject::{
     AttributeProtocol, IdProtocol, PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectRef,
     PyResult, Scope, TypeProtocol,
 };
+
+#[cfg(not(target_arch = "wasm32"))]
 use super::stdlib::io::io_open;
 
 use super::vm::VirtualMachine;
@@ -769,6 +771,7 @@ pub fn make_module(ctx: &PyContext) -> PyObjectRef {
     ctx.set_attr(&py_mod, "min", ctx.new_rustfunc(builtin_min));
     ctx.set_attr(&py_mod, "object", ctx.object());
     ctx.set_attr(&py_mod, "oct", ctx.new_rustfunc(builtin_oct));
+    #[cfg(not(target_arch = "wasm32"))]
     ctx.set_attr(&py_mod, "open", ctx.new_rustfunc(io_open));
     ctx.set_attr(&py_mod, "ord", ctx.new_rustfunc(builtin_ord));
     ctx.set_attr(&py_mod, "next", ctx.new_rustfunc(builtin_next));

diff --git a/vm/src/stdlib/mod.rs b/vm/src/stdlib/mod.rs
@@ -1,9 +1,7 @@
 mod ast;
-pub mod io;
 mod json;
 mod keyword;
 mod math;
-mod os;
 mod pystruct;
 mod random;
 mod re;
@@ -14,18 +12,21 @@ mod types;
 mod weakref;
 use std::collections::HashMap;
 
+#[cfg(not(target_arch = "wasm32"))]
+pub mod io;
+#[cfg(not(target_arch = "wasm32"))]
+mod os;
+
 use super::pyobject::{PyContext, PyObjectRef};
 
 pub type StdlibInitFunc = fn(&PyContext) -> PyObjectRef;
 
 pub fn get_module_inits() -> HashMap<String, StdlibInitFunc> {
     let mut modules = HashMap::new();
     modules.insert("ast".to_string(), ast::mk_module as StdlibInitFunc);
-    modules.insert("io".to_string(), io::mk_module as StdlibInitFunc);
     modules.insert("json".to_string(), json::mk_module as StdlibInitFunc);
     modules.insert("keyword".to_string(), keyword::mk_module as StdlibInitFunc);
     modules.insert("math".to_string(), math::mk_module as StdlibInitFunc);
-    modules.insert("os".to_string(), os::mk_module as StdlibInitFunc);
     modules.insert("re".to_string(), re::mk_module as StdlibInitFunc);
     modules.insert("random".to_string(), random::mk_module as StdlibInitFunc);
     modules.insert("string".to_string(), string::mk_module as StdlibInitFunc);
@@ -37,5 +38,13 @@ pub fn get_module_inits() -> HashMap<String, StdlibInitFunc> {
     );
     modules.insert("types".to_string(), types::mk_module as StdlibInitFunc);
     modules.insert("_weakref".to_string(), weakref::mk_module as StdlibInitFunc);
+
+    // disable some modules on WASM
+    #[cfg(not(target_arch = "wasm32"))]
+    {
+        modules.insert("io".to_string(), io::mk_module as StdlibInitFunc);
+        modules.insert("os".to_string(), os::mk_module as StdlibInitFunc);
+    }
+
     modules
 }
diff --git a/wasm/lib/src/convert.rs b/wasm/lib/src/convert.rs
@@ -1,14 +1,62 @@
 use js_sys::{Array, Object, Reflect};
 use rustpython_vm::pyobject::{self, PyFuncArgs, PyObjectRef, PyResult};
 use rustpython_vm::VirtualMachine;
-use wasm_bindgen::{prelude::*, JsCast};
+use vm_class::{StoredVirtualMachine, WASMVirtualMachine};
+use wasm_bindgen::{closure::Closure, prelude::*, JsCast};
 
 pub fn py_str_err(vm: &mut VirtualMachine, py_err: &PyObjectRef) -> String {
     vm.to_pystr(&py_err)
         .unwrap_or_else(|_| "Error, and error getting error message".into())
 }
 
-pub fn py_to_js(vm: &mut VirtualMachine, py_obj: PyObjectRef) -> JsValue {
+pub fn py_to_js(
+    vm: &mut VirtualMachine,
+    py_obj: PyObjectRef,
+    wasm_vm: Option<WASMVirtualMachine>,
+) -> JsValue {
+    if let Some(wasm_vm) = wasm_vm {
+        if rustpython_vm::obj::objtype::isinstance(&py_obj, &vm.ctx.function_type()) {
+            let closure = Closure::wrap(Box::new(
+                move |args: Option<Array>, kwargs: Option<Object>| -> Result<JsValue, JsValue> {
+                    let py_obj = py_obj.clone();
+                    let wasm_vm = wasm_vm.clone();
+                    wasm_vm.assert_valid()?;
+                    let wasm_vm_clone = wasm_vm.clone();
+                    wasm_vm.with_unchecked(move |StoredVirtualMachine { ref mut vm, .. }| {
+                        let vm: &mut VirtualMachine = vm;
+                        let mut py_func_args = rustpython_vm::pyobject::PyFuncArgs::default();
+                        if let Some(ref args) = args {
+                            for arg in args.values() {
+                                py_func_args.args.push(js_to_py(
+                                    vm,
+                                    arg?,
+                                    Some(wasm_vm_clone.clone()),
+                                ));
+                            }
+                        }
+                        if let Some(ref kwargs) = kwargs {
+                            for pair in object_entries(kwargs) {
+                                let (key, val) = pair?;
+                                py_func_args.kwargs.push((
+                                    js_sys::JsString::from(key).into(),
+                                    js_to_py(vm, val, Some(wasm_vm_clone.clone())),
+                                ));
+                            }
+                        }
+                        let result = vm.invoke(py_obj.clone(), py_func_args);
+                        pyresult_to_jsresult(vm, result, Some(wasm_vm_clone.clone()))
+                    })
+                },
+            )
+                as Box<Fn(Option<Array>, Option<Object>) -> Result<JsValue, JsValue>>);
+            let func = closure.as_ref().clone();
+
+            // TODO: Come up with a way of managing closure handles
+            closure.forget();
+
+            return func;
+        }
+    }
     let dumps = rustpython_vm::import::import(
         vm,
         std::path::PathBuf::default(),
@@ -25,14 +73,46 @@ pub fn py_to_js(vm: &mut VirtualMachine, py_obj: PyObjectRef) -> JsValue {
     }
 }
 
-pub fn js_to_py(vm: &mut VirtualMachine, js_val: JsValue) -> PyObjectRef {
+pub fn object_entries(obj: &Object) -> impl Iterator<Item = Result<(JsValue, JsValue), JsValue>> {
+    Object::entries(obj).values().into_iter().map(|pair| {
+        pair.map(|pair| {
+            let key = Reflect::get(&pair, &"0".into()).unwrap();
+            let val = Reflect::get(&pair, &"1".into()).unwrap();
+            (key, val)
+        })
+    })
+}
+
+pub fn pyresult_to_jsresult(
+    vm: &mut VirtualMachine,
+    result: PyResult,
+    wasm_vm: Option<WASMVirtualMachine>,
+) -> Result<JsValue, JsValue> {
+    result
+        .map(|value| py_to_js(vm, value, wasm_vm))
+        .map_err(|err| py_str_err(vm, &err).into())
+}
+
+pub fn js_to_py(
+    vm: &mut VirtualMachine,
+    js_val: JsValue,
+    // Accept a WASM VM because if js_val is a function it has to be able to convert
+    // its arguments to JS, and those arguments might include a closure
+    wasm_vm: Option<WASMVirtualMachine>,
+) -> PyObjectRef {
     if js_val.is_object() {
         if Array::is_array(&js_val) {
             let js_arr: Array = js_val.into();
             let elems = js_arr
                 .values()
                 .into_iter()
-                .map(|val| js_to_py(vm, val.expect("Iteration over array failed")))
+                .map(|val| {
+                    js_to_py(
+                        vm,
+                        val.expect("Iteration over array failed"),
+                        wasm_vm.clone(),
+                    )
+                })
                 .collect();
             vm.ctx.new_list(elems)
         } else {
@@ -41,7 +121,7 @@ pub fn js_to_py(vm: &mut VirtualMachine, js_val: JsValue) -> PyObjectRef {
                 let pair = pair.expect("Iteration over object failed");
                 let key = Reflect::get(&pair, &"0".into()).unwrap();
                 let val = Reflect::get(&pair, &"1".into()).unwrap();
-                let py_val = js_to_py(vm, val);
+                let py_val = js_to_py(vm, val, wasm_vm.clone());
                 vm.ctx
                     .set_item(&dict, &String::from(js_sys::JsString::from(key)), py_val);
             }
@@ -54,16 +134,16 @@ pub fn js_to_py(vm: &mut VirtualMachine, js_val: JsValue) -> PyObjectRef {
                 let func = func.clone();
                 let this = Object::new();
                 for (k, v) in args.kwargs {
-                    Reflect::set(&this, &k.into(), &py_to_js(vm, v))
+                    Reflect::set(&this, &k.into(), &py_to_js(vm, v, wasm_vm.clone()))
                         .expect("Couldn't set this property");
                 }
                 let js_args = Array::new();
                 for v in args.args {
-                    js_args.push(&py_to_js(vm, v));
+                    js_args.push(&py_to_js(vm, v, wasm_vm.clone()));
                 }
                 func.apply(&this, &js_args)
-                    .map(|val| js_to_py(vm, val))
-                    .map_err(|err| js_to_py(vm, err))
+                    .map(|val| js_to_py(vm, val, wasm_vm.clone()))
+                    .map_err(|err| js_to_py(vm, err, wasm_vm.clone()))
             },
         )
     } else if let Some(err) = js_val.dyn_ref::<js_sys::Error>() {

diff --git a/wasm/lib/src/lib.rs b/wasm/lib/src/lib.rs
@@ -99,7 +99,7 @@ pub fn eval_py(source: &str, options: Option<Object>) -> Result<JsValue, JsValue
                             &JsValue::UNDEFINED,
                             &wasm_builtins::format_print_args(vm, args)?.into(),
                         )
-                        .map_err(|err| convert::js_to_py(vm, err))?;
+                        .map_err(|err| convert::js_to_py(vm, err, None))?;
                         Ok(vm.get_none())
                     },
                 )
@@ -121,11 +121,9 @@ pub fn eval_py(source: &str, options: Option<Object>) -> Result<JsValue, JsValue
     let injections = vm.new_dict();
 
     if let Some(js_vars) = js_vars.clone() {
-        for pair in Object::entries(&js_vars).values() {
-            let pair = pair?;
-            let key = Reflect::get(&pair, &"0".into()).unwrap();
-            let val = Reflect::get(&pair, &"1".into()).unwrap();
-            let py_val = convert::js_to_py(&mut vm, val);
+        for pair in convert::object_entries(&js_vars) {
+            let (key, val) = pair?;
+            let py_val = convert::js_to_py(&mut vm, val, None);
             vm.ctx.set_item(
                 &injections,
                 &String::from(js_sys::JsString::from(key)),
@@ -136,9 +134,8 @@ pub fn eval_py(source: &str, options: Option<Object>) -> Result<JsValue, JsValue
 
     vm.ctx.set_item(&mut vars, "js_vars", injections);
 
-    eval(&mut vm, source, vars)
-        .map(|value| convert::py_to_js(&mut vm, value))
-        .map_err(|err| convert::py_str_err(&mut vm, &err).into())
+    let result = eval(&mut vm, source, vars);
+    convert::pyresult_to_jsresult(&mut vm, result, None)
 }
 
 #[wasm_bindgen(typescript_custom_section)]

diff --git a/wasm/lib/src/vm_class.rs b/wasm/lib/src/vm_class.rs
@@ -3,9 +3,10 @@ use js_sys::TypeError;
 use rustpython_vm::{compile, pyobject::PyObjectRef, VirtualMachine};
 use std::cell::RefCell;
 use std::collections::HashMap;
+use std::rc::Rc;
 use wasm_bindgen::prelude::*;
 
-struct StoredVirtualMachine {
+pub(crate) struct StoredVirtualMachine {
     pub vm: VirtualMachine,
     pub scope: PyObjectRef,
 }
@@ -21,7 +22,7 @@ impl StoredVirtualMachine {
 
 // It's fine that it's thread local, since WASM doesn't even have threads yet
 thread_local! {
-    static STORED_VMS: RefCell<HashMap<String, StoredVirtualMachine>> = RefCell::default();
+    static STORED_VMS: Rc<RefCell<HashMap<String, StoredVirtualMachine>>> = Rc::default();
 }
 
 #[wasm_bindgen(js_name = vmStore)]
@@ -62,17 +63,37 @@ impl VMStore {
 }
 
 #[wasm_bindgen(js_name = VirtualMachine)]
+#[derive(Clone)]
 pub struct WASMVirtualMachine {
     id: String,
 }
 
 #[wasm_bindgen(js_class = VirtualMachine)]
 impl WASMVirtualMachine {
+    pub(crate) fn with_unchecked<F, R>(&self, f: F) -> R
+    where
+        F: FnOnce(&mut StoredVirtualMachine) -> R,
+    {
+        STORED_VMS.with(|cell| {
+            let mut vms = cell.borrow_mut();
+            let stored_vm = vms.get_mut(&self.id).unwrap();
+            f(stored_vm)
+        })
+    }
+
+    pub(crate) fn with<F, R>(&self, f: F) -> Result<R, JsValue>
+    where
+        F: FnOnce(&mut StoredVirtualMachine) -> R,
+    {
+        self.assert_valid()?;
+        Ok(self.with_unchecked(f))
+    }
+
     pub fn valid(&self) -> bool {
         STORED_VMS.with(|cell| cell.borrow().contains_key(&self.id))
     }
 
-    fn assert_valid(&self) -> Result<(), JsValue> {
+    pub fn assert_valid(&self) -> Result<(), JsValue> {
         if self.valid() {
             Ok(())
         } else {
@@ -91,34 +112,32 @@ impl WASMVirtualMachine {
 
     #[wasm_bindgen(js_name = addToScope)]
     pub fn add_to_scope(&self, name: String, value: JsValue) -> Result<(), JsValue> {
-        self.assert_valid()?;
-        STORED_VMS.with(|cell| {
-            let mut vms = cell.borrow_mut();
-            let StoredVirtualMachine {
-                ref mut vm,
-                ref mut scope,
-            } = vms.get_mut(&self.id).unwrap();
-            let value = convert::js_to_py(vm, value);
-            vm.ctx.set_item(scope, &name, value);
-            Ok(())
-        })
+        self.with(
+            move |StoredVirtualMachine {
+                      ref mut vm,
+                      ref mut scope,
+                  }| {
+                let value = convert::js_to_py(vm, value, Some(self.clone()));
+                vm.ctx.set_item(scope, &name, value);
+            },
+        )
     }
 
     pub fn run(&self, mut source: String) -> Result<JsValue, JsValue> {
         self.assert_valid()?;
-        STORED_VMS.with(|cell| {
-            let mut vms = cell.borrow_mut();
-            let StoredVirtualMachine {
-                ref mut vm,
-                ref mut scope,
-            } = vms.get_mut(&self.id).unwrap();
-            source.push('\n');
-            let code = compile::compile(vm, &source, compile::Mode::Single, None)
-                .map_err(|err| convert::py_str_err(vm, &err))?;
-            let result = vm
-                .run_code_obj(code, scope.clone())
-                .map_err(|err| convert::py_str_err(vm, &err))?;
-            Ok(convert::py_to_js(vm, result))
-        })
+        self.with_unchecked(
+            |StoredVirtualMachine {
+                 ref mut vm,
+                 ref mut scope,
+             }| {
+                source.push('\n');
+                let code = compile::compile(vm, &source, &compile::Mode::Single, None)
+                    .map_err(|err| convert::py_str_err(vm, &err))?;
+                let result = vm
+                    .run_code_obj(code, scope.clone())
+                    .map_err(|err| convert::py_str_err(vm, &err))?;
+                Ok(convert::py_to_js(vm, result, Some(self.clone())))
+            },
+        )
     }
 }
diff --git a/wasm/lib/src/wasm_builtins.rs b/wasm/lib/src/wasm_builtins.rs
@@ -47,7 +47,7 @@ pub fn format_print_args(vm: &mut VirtualMachine, args: PyFuncArgs) -> Result<St
 
 pub fn builtin_print_html(vm: &mut VirtualMachine, args: PyFuncArgs, selector: &str) -> PyResult {
     let output = format_print_args(vm, args)?;
-    print_to_html(&output, selector).map_err(|err| convert::js_to_py(vm, err))?;
+    print_to_html(&output, selector).map_err(|err| convert::js_to_py(vm, err, None))?;
     Ok(vm.get_none())
 }