From 0d6c77cd6795fd8915cd2fcd73992d6267c1d095 Mon Sep 17 00:00:00 2001
From: Jeong YunWon <jeong@youknowone.org>
Date: Mon, 14 Sep 2020 16:26:30 +0900
Subject: [PATCH] PyObjectWeak as a struct

---
 vm/src/obj/objweakref.rs |  2 +-
 vm/src/pyobjectrc.rs     | 96 +++++++++++++++++++++++++++++++++++++---
 wasm/lib/src/convert.rs  |  3 +-
 wasm/lib/src/vm_class.rs |  9 ++--
 4 files changed, 95 insertions(+), 15 deletions(-)

diff --git a/vm/src/obj/objweakref.rs b/vm/src/obj/objweakref.rs
index 02c60f7d9..3e89ea48a 100644
--- a/vm/src/obj/objweakref.rs
+++ b/vm/src/obj/objweakref.rs
@@ -26,7 +26,7 @@ impl PyWeak {
     }
 
     pub fn upgrade(&self) -> Option<PyObjectRef> {
-        PyObjectRc::upgrade_weak(&self.referent)
+        self.referent.upgrade()
     }
 }
 
diff --git a/vm/src/pyobjectrc.rs b/vm/src/pyobjectrc.rs
index 01bc60e16..05a15069d 100644
--- a/vm/src/pyobjectrc.rs
+++ b/vm/src/pyobjectrc.rs
@@ -12,7 +12,12 @@ where
     inner: PyRc<PyObject<T>>,
 }
 
-pub type PyObjectWeak<T = dyn PyObjectPayload> = PyWeak<PyObject<T>>;
+pub struct PyObjectWeak<T = dyn PyObjectPayload>
+where
+    T: ?Sized + PyObjectPayload,
+{
+    inner: PyWeak<PyObject<T>>,
+}
 
 pub trait AsPyObjectRef {
     fn _as_ref(self) -> PyRc<PyObject<dyn PyObjectPayload>>;
@@ -80,11 +85,28 @@ where
     }
 
     pub fn downgrade(this: &Self) -> PyObjectWeak<T> {
-        PyRc::downgrade(&this.inner)
+        PyObjectWeak {
+            inner: PyRc::downgrade(&this.inner),
+        }
     }
+}
 
-    pub fn upgrade_weak(weak: &PyObjectWeak<T>) -> Option<Self> {
-        weak.upgrade().map(|inner| PyObjectRc { inner })
+impl<T: ?Sized + PyObjectPayload> IdProtocol for PyObjectRc<T>
+where
+    PyRc<PyObject<T>>: IdProtocol + AsPyObjectRef,
+{
+    fn get_id(&self) -> usize {
+        self.inner.get_id()
+    }
+}
+
+impl<T> PyObjectWeak<T>
+where
+    T: ?Sized + PyObjectPayload,
+    PyRc<PyObject<T>>: AsPyObjectRef,
+{
+    pub fn upgrade(&self) -> Option<PyObjectRc<T>> {
+        self.inner.upgrade().map(|inner| PyObjectRc { inner })
     }
 }
 
@@ -103,6 +125,11 @@ where
 {
 }
 
+#[cfg(feature = "threading")]
+unsafe impl<T> Send for PyObjectWeak<T> where T: ?Sized + PyObjectPayload {}
+#[cfg(feature = "threading")]
+unsafe impl<T> Sync for PyObjectWeak<T> where T: ?Sized + PyObjectPayload {}
+
 impl<T> Deref for PyObjectRc<T>
 where
     T: ?Sized + PyObjectPayload,
@@ -171,6 +198,16 @@ where
     }
 }
 
+impl<T> borrow::BorrowMut<T> for PyObjectRc<T>
+where
+    T: ?Sized + PyObjectPayload,
+    PyRc<PyObject<T>>: AsPyObjectRef + borrow::BorrowMut<T>,
+{
+    fn borrow_mut(&mut self) -> &mut T {
+        self.inner.borrow_mut()
+    }
+}
+
 impl<T> AsRef<T> for PyObjectRc<T>
 where
     T: ?Sized + PyObjectPayload,
@@ -181,8 +218,53 @@ where
     }
 }
 
-impl IdProtocol for PyObjectRc {
-    fn get_id(&self) -> usize {
-        self.inner.get_id()
+impl<T> Clone for PyObjectWeak<T>
+where
+    T: ?Sized + PyObjectPayload,
+{
+    fn clone(&self) -> Self {
+        PyObjectWeak {
+            inner: self.inner.clone(),
+        }
+    }
+}
+
+impl<T> fmt::Debug for PyObjectWeak<T>
+where
+    T: ?Sized + PyObjectPayload,
+    PyObject<T>: fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.inner.fmt(f)
+    }
+}
+
+impl<T> borrow::Borrow<T> for PyObjectWeak<T>
+where
+    T: ?Sized + PyObjectPayload,
+    PyWeak<PyObject<T>>: borrow::Borrow<T>,
+{
+    fn borrow(&self) -> &T {
+        self.inner.borrow()
+    }
+}
+
+impl<T> borrow::BorrowMut<T> for PyObjectWeak<T>
+where
+    T: ?Sized + PyObjectPayload,
+    PyWeak<PyObject<T>>: borrow::BorrowMut<T>,
+{
+    fn borrow_mut(&mut self) -> &mut T {
+        self.inner.borrow_mut()
+    }
+}
+
+impl<T> AsRef<T> for PyObjectWeak<T>
+where
+    T: ?Sized + PyObjectPayload,
+    PyWeak<PyObject<T>>: AsRef<T>,
+{
+    fn as_ref(&self) -> &T {
+        self.inner.as_ref()
     }
 }
diff --git a/wasm/lib/src/convert.rs b/wasm/lib/src/convert.rs
index be97c9ac5..a95e366c0 100644
--- a/wasm/lib/src/convert.rs
+++ b/wasm/lib/src/convert.rs
@@ -76,7 +76,8 @@ pub fn py_to_js(vm: &VirtualMachine, py_obj: PyObjectRef) -> JsValue {
             let closure =
                 move |args: Option<Array>, kwargs: Option<Object>| -> Result<JsValue, JsValue> {
                     let py_obj = match wasm_vm.assert_valid() {
-                        Ok(_) => PyObjectRef::upgrade_weak(&weak_py_obj)
+                        Ok(_) => weak_py_obj
+                            .upgrade()
                             .expect("weak_py_obj to be valid if VM is valid"),
                         Err(err) => {
                             return Err(err);
diff --git a/wasm/lib/src/vm_class.rs b/wasm/lib/src/vm_class.rs
index 5bcc3e1b0..29ba5cc16 100644
--- a/wasm/lib/src/vm_class.rs
+++ b/wasm/lib/src/vm_class.rs
@@ -6,8 +6,8 @@ use js_sys::{Object, TypeError};
 use wasm_bindgen::prelude::*;
 
 use rustpython_compiler::compile;
-use rustpython_vm::common::rc::{PyRc, PyWeak};
-use rustpython_vm::pyobject::{ItemProtocol, PyObject, PyObjectPayload, PyObjectRef, PyValue};
+use rustpython_vm::common::rc::PyRc;
+use rustpython_vm::pyobject::{ItemProtocol, PyObjectRef, PyObjectWeak, PyValue};
 use rustpython_vm::scope::{NameProtocol, Scope};
 use rustpython_vm::{InitParameter, PySettings, VirtualMachine};
 
@@ -174,10 +174,7 @@ impl WASMVirtualMachine {
         STORED_VMS.with(|cell| cell.borrow().contains_key(&self.id))
     }
 
-    pub(crate) fn push_held_rc(
-        &self,
-        obj: PyObjectRef,
-    ) -> Result<PyWeak<PyObject<dyn PyObjectPayload>>, JsValue> {
+    pub(crate) fn push_held_rc(&self, obj: PyObjectRef) -> Result<PyObjectWeak, JsValue> {
         self.with(|stored_vm| {
             let weak = PyObjectRef::downgrade(&obj);
             stored_vm.held_objects.borrow_mut().push(obj);