Put JS_FUNC in thread local storage

palaviv · palaviv · commit 7101a53879b8 · 2020-04-29T12:05:16.000+03:00
diff --git a/wasm/lib/src/convert.rs b/wasm/lib/src/convert.rs
@@ -1,3 +1,5 @@
+use std::cell::RefCell;
+
 use js_sys::{Array, ArrayBuffer, Object, Promise, Reflect, SyntaxError, Uint8Array};
 use wasm_bindgen::{closure::Closure, prelude::*, JsCast};
 
@@ -13,6 +15,9 @@ use rustpython_vm::{exceptions, py_serde};
 use crate::browser_module;
 use crate::vm_class::{stored_vm_from_wasm, WASMVirtualMachine};
 
+// Currently WASM do not support multithreading. We should change this once it is enabled.
+thread_local!(static JS_FUNC: RefCell<Option<js_sys::Function>> = RefCell::new(None));
+
 #[wasm_bindgen(inline_js = r"
 export class PyError extends Error {
     constructor(info) {
@@ -149,6 +154,25 @@ pub fn pyresult_to_jsresult(vm: &VirtualMachine, result: PyResult) -> Result<JsV
         .map_err(|err| py_err_to_js_err(vm, &err))
 }
 
+fn js_func_to_py(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
+    JS_FUNC.with(|func| {
+        let this = Object::new();
+        for (k, v) in args.kwargs {
+            Reflect::set(&this, &k.into(), &py_to_js(vm, v)).expect("property to be settable");
+        }
+        let js_args = Array::new();
+        for v in args.args {
+            js_args.push(&py_to_js(vm, v));
+        }
+        func.borrow()
+            .as_ref()
+            .unwrap()
+            .apply(&this, &js_args)
+            .map(|val| js_to_py(vm, val))
+            .map_err(|err| js_err_to_py_err(vm, &err))
+    })
+}
+
 pub fn js_to_py(vm: &VirtualMachine, js_val: JsValue) -> PyObjectRef {
     if js_val.is_object() {
         if let Some(promise) = js_val.dyn_ref::<Promise>() {
@@ -191,22 +215,8 @@ pub fn js_to_py(vm: &VirtualMachine, js_val: JsValue) -> PyObjectRef {
         }
     } else if js_val.is_function() {
         let func = js_sys::Function::from(js_val);
-        vm.ctx
-            .new_method(move |vm: &VirtualMachine, args: PyFuncArgs| -> PyResult {
-                let func = func.clone();
-                let this = Object::new();
-                for (k, v) in args.kwargs {
-                    Reflect::set(&this, &k.into(), &py_to_js(vm, v))
-                        .expect("property to be settable");
-                }
-                let js_args = Array::new();
-                for v in args.args {
-                    js_args.push(&py_to_js(vm, v));
-                }
-                func.apply(&this, &js_args)
-                    .map(|val| js_to_py(vm, val))
-                    .map_err(|err| js_err_to_py_err(vm, &err))
-            })
+        JS_FUNC.with(|thread_func| thread_func.replace(Some(func.clone())));
+        vm.ctx.new_method(js_func_to_py)
     } else if let Some(err) = js_val.dyn_ref::<js_sys::Error>() {
         js_err_to_py_err(vm, err).into_object()
     } else if js_val.is_undefined() {
diff --git a/wasm/lib/src/vm_class.rs b/wasm/lib/src/vm_class.rs
@@ -19,6 +19,9 @@ use crate::js_module;
 use crate::wasm_builtins;
 use rustpython_compiler::mode::Mode;
 
+// Currently WASM do not support multithreading. We should change this once it is enabled.
+thread_local!(static JS_FUNC: RefCell<Option<js_sys::Function>> = RefCell::new(None));
+
 pub(crate) struct StoredVirtualMachine {
     pub vm: VirtualMachine,
     pub scope: RefCell<Scope>,
@@ -151,6 +154,20 @@ pub struct WASMVirtualMachine {
     pub(crate) id: String,
 }
 
+fn stdout_js_func(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
+    JS_FUNC.with(|func| {
+        func.borrow()
+            .as_ref()
+            .unwrap()
+            .call1(
+                &JsValue::UNDEFINED,
+                &wasm_builtins::format_print_args(vm, args)?.into(),
+            )
+            .map_err(|err| convert::js_py_typeerror(vm, err))?;
+        Ok(vm.get_none())
+    })
+}
+
 #[wasm_bindgen(js_class = VirtualMachine)]
 impl WASMVirtualMachine {
     pub(crate) fn with_unchecked<F, R>(&self, f: F) -> R
@@ -229,15 +246,8 @@ impl WASMVirtualMachine {
                 }
             } else if stdout.is_function() {
                 let func = js_sys::Function::from(stdout);
-                vm.ctx
-                    .new_method(move |vm: &VirtualMachine, args: PyFuncArgs| -> PyResult {
-                        func.call1(
-                            &JsValue::UNDEFINED,
-                            &wasm_builtins::format_print_args(vm, args)?.into(),
-                        )
-                        .map_err(|err| convert::js_py_typeerror(vm, err))?;
-                        Ok(vm.get_none())
-                    })
+                JS_FUNC.with(|thread_func| thread_func.replace(Some(func.clone())));
+                vm.ctx.new_method(stdout_js_func)
             } else if stdout.is_null() {
                 fn noop(vm: &VirtualMachine, _args: PyFuncArgs) -> PyResult {
                     Ok(vm.get_none())
