Merge pull request RustPython#4513 from xiaozhiyan/optimize-pypydict-to_attributes-to-reuse-pyrefexact

Optimize Py<PyDict>::to_attributes() to reuse PyRefExact

Author: youknowone

vm/src/builtins/complex.rs

@@ -128,7 +128,7 @@ impl Constructor for PyComplex {
                 let val = if cls.is(vm.ctx.types.complex_type) && imag_missing {
                     match val.downcast_exact::<PyComplex>(vm) {
                         Ok(c) => {
-                            return Ok(c.into());
+                            return Ok(c.into_pyref().into());
                         }
                         Err(val) => val,
                     }

vm/src/builtins/dict.rs

@@ -1,6 +1,6 @@
 use super::{
     set::PySetInner, IterStatus, PositionIterInternal, PyBaseExceptionRef, PyGenericAlias,
-    PyMappingProxy, PySet, PyStrRef, PyTupleRef, PyType, PyTypeRef,
+    PyMappingProxy, PySet, PyStr, PyTupleRef, PyType, PyTypeRef,
 };
 use crate::{
     atomic_func,
@@ -11,7 +11,6 @@ use crate::{
     },
     class::{PyClassDef, PyClassImpl},
     common::ascii,
-    convert::ToPyObject,
     dictdatatype::{self, DictKey},
     function::{
         ArgIterable, FuncArgs, KwArgs, OptionalArg, PyArithmeticValue::*, PyComparisonValue,
@@ -24,7 +23,8 @@ use crate::{
         IterNextIterable, Iterable, PyComparisonOp, Unconstructible, Unhashable,
     },
     vm::VirtualMachine,
-    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,
+    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyRefExact, PyResult,
+    TryFromObject,
 };
 use once_cell::sync::Lazy;
 use rustpython_common::lock::PyMutex;
@@ -65,8 +65,9 @@ impl PyDict {
 
     // Used in update and ior.
     pub(crate) fn merge_object(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
-        let other = match other.downcast_exact(vm) {
-            Ok(dict_other) => return self.merge_dict(dict_other, vm),
+        let casted: Result<PyRefExact<PyDict>, _> = other.downcast_exact(vm);
+        let other = match casted {
+            Ok(dict_other) => return self.merge_dict(dict_other.into_pyref(), vm),
             Err(other) => other,
         };
         let dict = &self.entries;
@@ -227,7 +228,7 @@ impl PyDict {
                 for key in iterable.iter(vm)? {
                     pydict.setitem(key?, value.clone(), vm)?;
                 }
-                Ok(pydict.to_pyobject(vm))
+                Ok(pydict.into_pyref().into())
             }
             Err(pyobj) => {
                 for key in iterable.iter(vm)? {
@@ -533,9 +534,8 @@ impl Py<PyDict> {
     pub fn to_attributes(&self, vm: &VirtualMachine) -> PyAttributes {
         let mut attrs = PyAttributes::default();
         for (key, value) in self {
-            // TODO: use PyRefExact for interning
-            let key: PyStrRef = key.downcast().expect("dict has non-string keys");
-            attrs.insert(vm.ctx.intern_str(key.as_str()), value);
+            let key: PyRefExact<PyStr> = key.downcast_exact(vm).expect("dict has non-string keys");
+            attrs.insert(vm.ctx.intern_str(key), value);
         }
         attrs
     }

vm/src/builtins/int.rs

@@ -231,7 +231,7 @@ impl Constructor for PyInt {
                 let val = if cls.is(vm.ctx.types.int_type) {
                     match val.downcast_exact::<PyInt>(vm) {
                         Ok(i) => {
-                            return Ok(i.to_pyobject(vm));
+                            return Ok(i.into_pyref().into());
                         }
                         Err(val) => val,
                     }

vm/src/builtins/set.rs

@@ -354,7 +354,7 @@ impl PySetInner {
             self.merge_set(any_set, vm)
         // check Dict
         } else if let Ok(dict) = iterable.to_owned().downcast_exact::<PyDict>(vm) {
-            self.merge_dict(dict, vm)
+            self.merge_dict(dict.into_pyref(), vm)
         } else {
             // add iterable that is not AnySet or Dict
             for item in iterable.try_into_value::<ArgIterable>(vm)?.iter(vm)? {
@@ -799,7 +799,7 @@ impl Constructor for PyFrozenSet {
         let elements = if let OptionalArg::Present(iterable) = iterable {
             let iterable = if cls.is(vm.ctx.types.frozenset_type) {
                 match iterable.downcast_exact::<Self>(vm) {
-                    Ok(fs) => return Ok(fs.into()),
+                    Ok(fs) => return Ok(fs.into_pyref().into()),
                     Err(iterable) => iterable,
                 }
             } else {

vm/src/builtins/tuple.rs

@@ -95,7 +95,7 @@ impl Constructor for PyTuple {
         let elements = if let OptionalArg::Present(iterable) = iterable {
             let iterable = if cls.is(vm.ctx.types.tuple_type) {
                 match iterable.downcast_exact::<Self>(vm) {
-                    Ok(tuple) => return Ok(tuple.into()),
+                    Ok(tuple) => return Ok(tuple.into_pyref().into()),
                     Err(iterable) => iterable,
                 }
             } else {

vm/src/bytesinner.rs

@@ -79,9 +79,7 @@ impl ByteInnerNewOptions {
                 // construct an exact bytes from an exact bytes do not clone
                 let obj = if cls.is(PyBytes::class(vm)) {
                     match obj.downcast_exact::<PyBytes>(vm) {
-                        Ok(b) => {
-                            return Ok(b);
-                        }
+                        Ok(b) => return Ok(b.into_pyref()),
                         Err(obj) => obj,
                     }
                 } else {

vm/src/object/core.rs

@@ -12,7 +12,7 @@
 //! but not to do to remember it is a PyRef object.
 
 use super::{
-    ext::{AsObject, PyResult},
+    ext::{AsObject, PyRefExact, PyResult},
     payload::PyObjectPayload,
     PyAtomicRef,
 };
@@ -564,15 +564,15 @@ impl PyObjectRef {
     pub fn downcast_exact<T: PyObjectPayload + crate::PyPayload>(
         self,
         vm: &VirtualMachine,
-    ) -> Result<PyRef<T>, Self> {
+    ) -> Result<PyRefExact<T>, Self> {
         if self.class().is(T::class(vm)) {
             // TODO: is this always true?
             assert!(
                 self.payload_is::<T>(),
                 "obj.__class__ is T::class() but payload is not T"
             );
             // SAFETY: just asserted that payload_is::<T>()
-            Ok(unsafe { PyRef::from_obj_unchecked(self) })
+            Ok(unsafe { PyRefExact::new_unchecked(PyRef::from_obj_unchecked(self)) })
         } else {
             Err(self)
         }

vm/src/object/ext.rs

@@ -107,6 +107,7 @@ impl<T: PyPayload> std::borrow::ToOwned for PyExact<T> {
     }
 }
 
+/// PyRef but guaranteed not to be a subtype instance
 #[derive(Debug)]
 #[repr(transparent)]
 pub struct PyRefExact<T: PyObjectPayload> {

vm/src/stdlib/builtins.rs

@@ -872,7 +872,11 @@ mod builtins {
         // Use downcast_exact to keep ref to old object on error.
         let metaclass = kwargs
             .pop_kwarg("metaclass")
-            .map(|metaclass| metaclass.downcast_exact::<PyType>(vm))
+            .map(|metaclass| {
+                metaclass
+                    .downcast_exact::<PyType>(vm)
+                    .map(|m| m.into_pyref())
+            })
             .unwrap_or_else(|| Ok(vm.ctx.types.type_type.to_owned()));
 
         let (metaclass, meta_name) = match metaclass {