Make PyObjectRef and PyObjectWeakRef it's own type.

RustPython · BenLewis-Seequent · Feb 13, 2019 · Feb 16, 2019 · Feb 17, 2019 · Feb 18, 2019
commit 8dd58108485224bce8eeb672ad77468add7510f5
diff --git a/vm/src/builtins.rs b/vm/src/builtins.rs
@@ -16,8 +16,8 @@ use super::obj::objstr;
 use super::obj::objtype;
 
 use super::pyobject::{
-    AttributeProtocol, IdProtocol, PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectRef,
-    PyResult, Scope, TypeProtocol,
+    AttributeProtocol, IdProtocol, PyContext, PyFuncArgs, PyObjectRef,
+    PyResult, TypeProtocol,
 };
 use super::stdlib::io::io_open;
 
@@ -265,17 +265,7 @@ fn make_scope(vm: &mut VirtualMachine, locals: Option<&PyObjectRef>) -> PyObject
     };
 
     // TODO: handle optional globals
-    // Construct new scope:
-    let scope_inner = Scope {
-        locals,
-        parent: None,
-    };
-
-    PyObject {
-        payload: PyObjectPayload::Scope { scope: scope_inner },
-        typ: None,
-    }
-    .into_ref()
+    vm.ctx.new_scope_with_locals(None, locals)
 }
 
 fn builtin_format(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {

diff --git a/vm/src/obj/objstr.rs b/vm/src/obj/objstr.rs
@@ -1,6 +1,6 @@
 use super::super::format::{FormatParseError, FormatPart, FormatString};
 use super::super::pyobject::{
-    PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectRef, PyResult, TypeProtocol,
+    PyContext, PyFuncArgs, PyObjectPayload, PyObjectRef, PyResult, TypeProtocol,
 };
 use super::super::vm::VirtualMachine;
 use super::objint;
@@ -1126,8 +1126,8 @@ pub fn subscript(vm: &mut VirtualMachine, value: &str, b: PyObjectRef) -> PyResu
 
 // help get optional string indices
 fn get_slice(
-    start: Option<&std::rc::Rc<std::cell::RefCell<PyObject>>>,
-    end: Option<&std::rc::Rc<std::cell::RefCell<PyObject>>>,
+    start: Option<&PyObjectRef>,
+    end: Option<&PyObjectRef>,
     len: usize,
 ) -> Result<(usize, usize), String> {
     let start_idx = match start {

diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs
@@ -38,6 +38,7 @@ use std::cell::RefCell;
 use std::collections::HashMap;
 use std::fmt;
 use std::rc::{Rc, Weak};
+use std::ops::Deref;
 
 /* Python objects and references.
 
@@ -54,20 +55,48 @@ Basically reference counting, but then done by rust.
  */
 
 /*
-The PyRef type implements
+The PyObjectRef type implements
 https://doc.rust-lang.org/std/cell/index.html#introducing-mutability-inside-of-something-immutable
 */
-pub type PyRef<T> = Rc<RefCell<T>>;
 
 /// The `PyObjectRef` is one of the most used types. It is a reference to a
 /// python object. A single python object can have multiple references, and
 /// this reference counting is accounted for by this type. Use the `.clone()`
 /// method to create a new reference and increment the amount of references
 /// to the python object by 1.
-pub type PyObjectRef = PyRef<PyObject>;
+#[derive(Debug, Clone)]
+pub struct PyObjectRef {
+    rc: Rc<RefCell<PyObject>>
+}
+
+impl PyObjectRef {
+    pub fn downgrade(this: &Self) -> PyObjectWeakRef {
+        PyObjectWeakRef { weak: Rc::downgrade(&this.rc) }
+    }
+
+    pub fn strong_count(this: &Self) -> usize {
+        Rc::strong_count(&this.rc)
+    }
+}
 
-/// Same as PyObjectRef, except for being a weak reference.
-pub type PyObjectWeakRef = Weak<RefCell<PyObject>>;
+impl Deref for PyObjectRef {
+    type Target = RefCell<PyObject>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.rc
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct PyObjectWeakRef {
+    weak: Weak<RefCell<PyObject>>
+}
+
+impl PyObjectWeakRef {
+    pub fn upgrade(&self) -> Option<PyObjectRef> {
+        self.weak.upgrade().map(|x| PyObjectRef { rc:x })
+    }
+}
 
 /// Use this type for function which return a python object or and exception.
 /// Both the python object and the python exception are `PyObjectRef` types
@@ -164,11 +193,7 @@ pub struct Scope {
 }
 
 fn _nothing() -> PyObjectRef {
-    PyObject {
-        payload: PyObjectPayload::None,
-        typ: None,
-    }
-    .into_ref()
+    PyObject::new_no_type(PyObjectPayload::None)
 }
 
 pub fn create_type(
@@ -535,12 +560,12 @@ impl PyContext {
 
     pub fn new_scope(&self, parent: Option<PyObjectRef>) -> PyObjectRef {
         let locals = self.new_dict();
+        self.new_scope_with_locals(parent, locals)
+    }
+
+    pub fn new_scope_with_locals(&self, parent: Option<PyObjectRef>, locals: PyObjectRef) -> PyObjectRef {
         let scope = Scope { locals, parent };
-        PyObject {
-            payload: PyObjectPayload::Scope { scope },
-            typ: None,
-        }
-        .into_ref()
+        PyObject::new_no_type(PyObjectPayload::Scope { scope })
     }
 
     pub fn new_module(&self, name: &str, scope: PyObjectRef) -> PyObjectRef {
@@ -697,6 +722,7 @@ impl PyContext {
 pub struct PyObject {
     pub payload: PyObjectPayload,
     pub typ: Option<PyObjectRef>,
+    pub weak_refs: Vec<PyObjectRef>
     // pub dict: HashMap<String, PyObjectRef>, // __dict__ member
 }
 
@@ -1015,6 +1041,7 @@ pub enum PyObjectPayload {
     },
     WeakRef {
         referent: PyObjectWeakRef,
+        callback: Option<PyObjectRef>,
     },
     Instance {
         dict: RefCell<PyAttributes>,
@@ -1072,13 +1099,24 @@ impl PyObject {
             payload,
             typ: Some(typ),
             // dict: HashMap::new(),  // dict,
+            weak_refs: Vec::new()
+        }
+        .into_ref()
+    }
+
+    fn new_no_type(payload: PyObjectPayload) -> PyObjectRef {
+        PyObject {
+            payload,
+            typ: None,
+            // dict: HashMap::new(),  // dict,
+            weak_refs: Vec::new()
         }
         .into_ref()
     }
 
     // Move this object into a reference object, transferring ownership.
     pub fn into_ref(self) -> PyObjectRef {
-        Rc::new(RefCell::new(self))
+        PyObjectRef { rc: Rc::new(RefCell::new(self)) }
     }
 }
 

diff --git a/vm/src/stdlib/weakref.rs b/vm/src/stdlib/weakref.rs
@@ -11,7 +11,6 @@ use super::super::pyobject::{
     TypeProtocol,
 };
 use super::super::VirtualMachine;
-use std::rc::Rc;
 
 pub fn mk_module(ctx: &PyContext) -> PyObjectRef {
     let py_mod = ctx.new_module("_weakref", ctx.new_scope(None));
@@ -25,10 +24,11 @@ pub fn mk_module(ctx: &PyContext) -> PyObjectRef {
 
 fn ref_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     // TODO: check first argument for subclass of `ref`.
-    arg_check!(vm, args, required = [(cls, None), (referent, None)]);
-    let referent = Rc::downgrade(referent);
+    arg_check!(vm, args, required = [(cls, None), (referent, None)],
+        optional = [(callback, None)]);
+    let referent = PyObjectRef::downgrade(referent);
     Ok(PyObject::new(
-        PyObjectPayload::WeakRef { referent },
+        PyObjectPayload::WeakRef { referent, callback: callback.cloned() },
         cls.clone(),
     ))
 }
@@ -47,7 +47,7 @@ fn ref_call(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 }
 
 fn get_value(obj: &PyObjectRef) -> PyObjectWeakRef {
-    if let PyObjectPayload::WeakRef { referent } = &obj.borrow().payload {
+    if let PyObjectPayload::WeakRef { referent, .. } = &obj.borrow().payload {
         referent.clone()
     } else {
         panic!("Inner error getting weak ref {:?}", obj);

diff --git a/vm/src/sysmodule.rs b/vm/src/sysmodule.rs
@@ -1,6 +1,5 @@
 use obj::objtype;
 use pyobject::{PyContext, PyFuncArgs, PyObjectRef, PyResult, TypeProtocol};
-use std::rc::Rc;
 use std::{env, mem};
 use vm::VirtualMachine;
 
@@ -24,7 +23,7 @@ fn getframe(vm: &mut VirtualMachine, _args: PyFuncArgs) -> PyResult {
 
 fn sys_getrefcount(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(object, None)]);
-    let size = Rc::strong_count(&object);
+    let size = PyObjectRef::strong_count(object);
     Ok(vm.ctx.new_int(size))
 }