From d0b4751ab2f84de0825b18d4e8724ff02eed0a13 Mon Sep 17 00:00:00 2001
From: coolreader18 <33094578+coolreader18@users.noreply.github.com>
Date: Sun, 17 Feb 2019 08:08:22 -0600
Subject: [PATCH] Remove the VM pointer from the map when it's dropped

---
 wasm/lib/src/convert.rs       |  2 +-
 wasm/lib/src/vm_class.rs      | 41 +++++++++++++++++++++++++++++++----
 wasm/lib/src/wasm_builtins.rs |  2 +-
 3 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/wasm/lib/src/convert.rs b/wasm/lib/src/convert.rs
index d8510db4f..379fc86c1 100644
--- a/wasm/lib/src/convert.rs
+++ b/wasm/lib/src/convert.rs
@@ -20,7 +20,7 @@ pub fn py_to_js(vm: &mut VirtualMachine, py_obj: PyObjectRef) -> JsValue {
                     let py_obj = py_obj.clone();
                     wasm_vm.assert_valid()?;
                     let acc_vm = AccessibleVM::from(wasm_vm.clone());
-                    let vm = acc_vm
+                    let vm = &mut acc_vm
                         .upgrade()
                         .expect("acc. VM to be invalid when WASM vm is valid");
                     let mut py_func_args = rustpython_vm::pyobject::PyFuncArgs::default();
diff --git a/wasm/lib/src/vm_class.rs b/wasm/lib/src/vm_class.rs
index c9c210451..319f116b4 100644
--- a/wasm/lib/src/vm_class.rs
+++ b/wasm/lib/src/vm_class.rs
@@ -106,20 +106,26 @@ impl AccessibleVM {
         )
     }
 
-    pub fn upgrade(&self) -> Option<&mut VirtualMachine> {
+    pub fn upgrade(&self) -> Option<AccessibleVMPtr> {
         let vm_cell = self.weak.upgrade()?;
-        match vm_cell.try_borrow_mut() {
+        let top_level = match vm_cell.try_borrow_mut() {
             Ok(mut vm) => {
                 ACTIVE_VMS.with(|cell| {
                     cell.borrow_mut().insert(self.id.clone(), &mut vm.vm);
                 });
+                true
             }
-            Err(_) => {}
+            Err(_) => false,
         };
         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 }
+            let vm = unsafe { &mut **ptr };
+            AccessibleVMPtr {
+                id: self.id.clone(),
+                top_level,
+                inner: vm,
+            }
         }))
     }
 }
@@ -135,6 +141,33 @@ impl From<&WASMVirtualMachine> for AccessibleVM {
     }
 }
 
+pub(crate) struct AccessibleVMPtr<'a> {
+    id: String,
+    top_level: bool,
+    inner: &'a mut VirtualMachine,
+}
+
+impl std::ops::Deref for AccessibleVMPtr<'_> {
+    type Target = VirtualMachine;
+    fn deref(&self) -> &VirtualMachine {
+        &self.inner
+    }
+}
+impl std::ops::DerefMut for AccessibleVMPtr<'_> {
+    fn deref_mut(&mut self) -> &mut VirtualMachine {
+        &mut self.inner
+    }
+}
+
+impl Drop for AccessibleVMPtr<'_> {
+    fn drop(&mut self) {
+        if self.top_level {
+            // remove the (now invalid) pointer from the map
+            ACTIVE_VMS.with(|cell| cell.borrow_mut().remove(&self.id));
+        }
+    }
+}
+
 #[wasm_bindgen(js_name = VirtualMachine)]
 #[derive(Clone)]
 pub struct WASMVirtualMachine {
diff --git a/wasm/lib/src/wasm_builtins.rs b/wasm/lib/src/wasm_builtins.rs
index b6764d60e..2f1d843ff 100644
--- a/wasm/lib/src/wasm_builtins.rs
+++ b/wasm/lib/src/wasm_builtins.rs
@@ -164,7 +164,7 @@ fn builtin_fetch(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         })
         .and_then(|prom| JsFuture::from(prom))
         .then(move |val| {
-            let vm = acc_vm
+            let vm = &mut acc_vm
                 .upgrade()
                 .expect("that the VM *not* be destroyed while promise is being resolved");
             match val {