Merge pull request RustPython#551 from coolreader18/store-rcs-in-stored-vm

cthulahoops · web-flow · commit 3730ca2b43f3 · 2019-03-09T06:36:01.000Z
[WASM] Keep a PyObjectWeakRef to the PyObjectRef being moved to a closure
diff --git a/wasm/lib/src/convert.rs b/wasm/lib/src/convert.rs
@@ -22,13 +22,15 @@ pub fn py_to_js(vm: &mut VirtualMachine, py_obj: PyObjectRef) -> JsValue {
             let wasm_vm = WASMVirtualMachine {
                 id: wasm_id.clone(),
             };
-            let mut py_obj = Some(py_obj);
+            let weak_py_obj = wasm_vm.push_held_rc(py_obj).unwrap();
+
             let closure =
                 move |args: Option<Array>, kwargs: Option<Object>| -> Result<JsValue, JsValue> {
                     let py_obj = match wasm_vm.assert_valid() {
-                        Ok(_) => py_obj.clone().expect("py_obj to be valid if VM is valid"),
+                        Ok(_) => weak_py_obj
+                            .upgrade()
+                            .expect("weak_py_obj to be valid if VM is valid"),
                         Err(err) => {
-                            py_obj = None;
                             return Err(err);
                         }
                     };
@@ -57,7 +59,8 @@ pub fn py_to_js(vm: &mut VirtualMachine, py_obj: PyObjectRef) -> JsValue {
                 as Box<dyn FnMut(Option<Array>, Option<Object>) -> Result<JsValue, JsValue>>);
             let func = closure.as_ref().clone();
 
-            // TODO: Come up with a way of managing closure handles
+            // stores pretty much nothing, it's fine to leak this because if it gets dropped
+            // the error message is worse
             closure.forget();
 
             return func;
diff --git a/wasm/lib/src/vm_class.rs b/wasm/lib/src/vm_class.rs
@@ -13,9 +13,16 @@ use std::collections::HashMap;
 use std::rc::{Rc, Weak};
 use wasm_bindgen::prelude::*;
 
+pub trait HeldRcInner {}
+
+impl<T> HeldRcInner for T {}
+
 pub(crate) struct StoredVirtualMachine {
     pub vm: VirtualMachine,
     pub scope: ScopeRef,
+    /// you can put a Rc in here, keep it as a Weak, and it'll be held only for
+    /// as long as the StoredVM is alive
+    held_rcs: Vec<Rc<dyn HeldRcInner>>,
 }
 
 impl StoredVirtualMachine {
@@ -27,7 +34,11 @@ impl StoredVirtualMachine {
             setup_browser_module(&mut vm);
         }
         vm.wasm_id = Some(id);
-        StoredVirtualMachine { vm, scope }
+        StoredVirtualMachine {
+            vm,
+            scope,
+            held_rcs: vec![],
+        }
     }
 }
 
@@ -211,6 +222,17 @@ impl WASMVirtualMachine {
         STORED_VMS.with(|cell| cell.borrow().contains_key(&self.id))
     }
 
+    pub(crate) fn push_held_rc<T: HeldRcInner + 'static>(
+        &self,
+        rc: Rc<T>,
+    ) -> Result<Weak<T>, JsValue> {
+        self.with(|stored_vm| {
+            let weak = Rc::downgrade(&rc);
+            stored_vm.held_rcs.push(rc);
+            weak
+        })
+    }
+
     pub fn assert_valid(&self) -> Result<(), JsValue> {
         if self.valid() {
             Ok(())
@@ -234,6 +256,7 @@ impl WASMVirtualMachine {
             move |StoredVirtualMachine {
                       ref mut vm,
                       ref mut scope,
+                      ..
                   }| {
                 let value = convert::js_to_py(vm, value);
                 scope.locals.set_item(&vm.ctx, &name, value);
@@ -247,6 +270,7 @@ impl WASMVirtualMachine {
             move |StoredVirtualMachine {
                       ref mut vm,
                       ref mut scope,
+                      ..
                   }| {
                 let print_fn: Box<Fn(&mut VirtualMachine, PyFuncArgs) -> PyResult> =
                     if let Some(selector) = stdout.as_string() {
@@ -315,6 +339,7 @@ impl WASMVirtualMachine {
             |StoredVirtualMachine {
                  ref mut vm,
                  ref mut scope,
+                 ..
              }| {
                 source.push('\n');
                 let code =
