Fix RefCell borrowing errors

RustPython · windelbouwman · Feb 23, 2019 · Jan 30, 2019 · Jan 30, 2019 · Jan 30, 2019
commit 68517676541b3d6a6ee3d1987f80ba86b24e81c2
diff --git a/wasm/lib/src/convert.rs b/wasm/lib/src/convert.rs
@@ -1,7 +1,7 @@
 use js_sys::{Array, Object, Reflect};
 use rustpython_vm::pyobject::{self, PyFuncArgs, PyObjectRef, PyResult};
 use rustpython_vm::VirtualMachine;
-use vm_class::{StoredVirtualMachine, WASMVirtualMachine};
+use vm_class::{AccessibleVM, WASMVirtualMachine};
 use wasm_bindgen::{closure::Closure, prelude::*, JsCast};
 
 pub fn py_str_err(vm: &mut VirtualMachine, py_err: &PyObjectRef) -> String {
@@ -16,39 +16,36 @@ pub fn py_to_js(
 ) -> 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(
+            let closure =
                 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()),
-                                ));
-                            }
+                    let acc_vm = AccessibleVM::from(wasm_vm.clone());
+                    let vm = acc_vm
+                        .upgrade()
+                        .expect("acc. VM to be invalid when WASM vm is valid");
+                    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())));
                         }
-                        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())),
-                                ));
-                            }
+                    }
+                    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())),
+                            ));
                         }
-                        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 result = vm.invoke(py_obj.clone(), py_func_args);
+                    pyresult_to_jsresult(vm, result, Some(wasm_vm.clone()))
+                };
+            let closure = Closure::wrap(Box::new(closure)
+                as Box<dyn Fn(Option<Array>, Option<Object>) -> Result<JsValue, JsValue>>);
             let func = closure.as_ref().clone();
 
             // TODO: Come up with a way of managing closure handles

diff --git a/wasm/lib/src/vm_class.rs b/wasm/lib/src/vm_class.rs
@@ -1,9 +1,13 @@
 use convert;
 use js_sys::TypeError;
-use rustpython_vm::{compile, pyobject::PyObjectRef, VirtualMachine};
+use rustpython_vm::{
+    compile,
+    pyobject::{PyObjectRef, PyRef},
+    VirtualMachine,
+};
 use std::cell::RefCell;
 use std::collections::HashMap;
-use std::rc::Rc;
+use std::rc::{Rc, Weak};
 use wasm_bindgen::prelude::*;
 
 pub(crate) struct StoredVirtualMachine {
@@ -22,7 +26,8 @@ impl StoredVirtualMachine {
 
 // It's fine that it's thread local, since WASM doesn't even have threads yet
 thread_local! {
-    static STORED_VMS: Rc<RefCell<HashMap<String, StoredVirtualMachine>>> = Rc::default();
+    static STORED_VMS: PyRef<HashMap<String, PyRef<StoredVirtualMachine>>> = Rc::default();
+    static ACTIVE_VMS: PyRef<HashMap<String, *mut VirtualMachine>> = Rc::default();
 }
 
 #[wasm_bindgen(js_name = vmStore)]
@@ -34,7 +39,10 @@ impl VMStore {
         STORED_VMS.with(|cell| {
             let mut vms = cell.borrow_mut();
             if !vms.contains_key(&id) {
-                vms.insert(id.clone(), StoredVirtualMachine::new());
+                vms.insert(
+                    id.clone(),
+                    Rc::new(RefCell::new(StoredVirtualMachine::new())),
+                );
             }
         });
         WASMVirtualMachine { id }
@@ -53,7 +61,18 @@ impl VMStore {
 
     pub fn destroy(id: String) {
         STORED_VMS.with(|cell| {
-            cell.borrow_mut().remove(&id);
+            use std::collections::hash_map::Entry;
+            match cell.borrow_mut().entry(id) {
+                Entry::Occupied(o) => {
+                    let (_k, stored_vm) = o.remove_entry();
+                    // for f in stored_vm.drop_handlers.iter() {
+                    //     f();
+                    // }
+                    // deallocate the VM
+                    drop(stored_vm);
+                }
+                Entry::Vacant(_v) => {}
+            }
         });
     }
 
@@ -62,6 +81,47 @@ impl VMStore {
     }
 }
 
+pub(crate) struct AccessibleVM {
+    weak: Weak<RefCell<StoredVirtualMachine>>,
+    id: String,
+}
+
+impl AccessibleVM {
+    pub fn from_id(id: String) -> AccessibleVM {
+        let weak = STORED_VMS
+            .with(|cell| Rc::downgrade(cell.borrow().get(&id).expect("WASM VM to be valid")));
+        AccessibleVM { weak, id }
+    }
+
+    pub fn upgrade(&self) -> Option<&mut VirtualMachine> {
+        let vm_cell = self.weak.upgrade()?;
+        match vm_cell.try_borrow_mut() {
+            Ok(mut vm) => {
+                ACTIVE_VMS.with(|cell| {
+                    cell.borrow_mut().insert(self.id.clone(), &mut vm.vm);
+                });
+            }
+            Err(_) => {}
+        };
+        Some(ACTIVE_VMS.with(|cell| {
+            let vms = cell.borrow();
+            let ptr = vms.get(&self.id).expect("id to be in ACTIVE_VMS");
+            unsafe { &mut **ptr }
+        }))
+    }
+}
+
+impl From<WASMVirtualMachine> for AccessibleVM {
+    fn from(vm: WASMVirtualMachine) -> AccessibleVM {
+        AccessibleVM::from_id(vm.id)
+    }
+}
+impl From<&WASMVirtualMachine> for AccessibleVM {
+    fn from(vm: &WASMVirtualMachine) -> AccessibleVM {
+        AccessibleVM::from_id(vm.id.clone())
+    }
+}
+
 #[wasm_bindgen(js_name = VirtualMachine)]
 #[derive(Clone)]
 pub struct WASMVirtualMachine {
@@ -76,8 +136,8 @@ impl WASMVirtualMachine {
     {
         STORED_VMS.with(|cell| {
             let mut vms = cell.borrow_mut();
-            let stored_vm = vms.get_mut(&self.id).unwrap();
-            f(stored_vm)
+            let mut stored_vm = vms.get_mut(&self.id).unwrap().borrow_mut();
+            f(&mut stored_vm)
         })
     }