Optimize PyType Creation (#4209)

stdlib/src/ssl.rs

@@ -228,7 +228,7 @@ mod _ssl {
     /// SSL/TLS connection terminated abruptly.
     #[pyattr(name = "SSLEOFError", once)]
     fn ssl_eof_error(vm: &VirtualMachine) -> PyTypeRef {
-        PyType::new_simple_ref("ssl.SSLEOFError", &ssl_error(vm)).unwrap()
+        PyType::new_simple_ref("ssl.SSLEOFError", &ssl_error(vm), &vm.ctx).unwrap()
     }
 
     type OpensslVersionInfo = (u8, u8, u8, u8, u8);

vm/src/builtins/type.rs

@@ -121,13 +121,18 @@ impl PyPayload for PyType {
 }
 
 impl PyType {
-    pub fn new_simple_ref(name: &str, base: &PyTypeRef) -> Result<PyRef<Self>, String> {
+    pub fn new_simple_ref(
+        name: &str,
+        base: &PyTypeRef,
+        ctx: &Context,
+    ) -> Result<PyRef<Self>, String> {
         Self::new_ref(
             name,
             vec![base.clone()],
             Default::default(),
             Default::default(),
             Self::static_type().to_owned(),
+            ctx,
         )
     }
     pub fn new_ref(
@@ -136,6 +141,7 @@ impl PyType {
         attrs: PyAttributes,
         slots: PyTypeSlots,
         metaclass: PyRef<Self>,
+        ctx: &Context,
     ) -> Result<PyRef<Self>, String> {
         Self::new_verbose_ref(
             name,
@@ -145,8 +151,11 @@ impl PyType {
             slots,
             HeapTypeExt::default(),
             metaclass,
+            ctx,
         )
     }
+
+    #[allow(clippy::too_many_arguments)]
     fn new_verbose_ref(
         name: &str,
         base: PyRef<Self>,
@@ -155,6 +164,7 @@ impl PyType {
         mut slots: PyTypeSlots,
         heaptype_ext: HeapTypeExt,
         metaclass: PyRef<Self>,
+        ctx: &Context,
     ) -> Result<PyRef<Self>, String> {
         // Check for duplicates in bases.
         let mut unique_bases = HashSet::new();
@@ -177,17 +187,22 @@ impl PyType {
         *slots.name.get_mut() = Some(String::from(name));
 
         #[allow(clippy::mutable_key_type)]
-        let mut attr_name_set = HashSet::new();
+        let mut slot_name_set = HashSet::new();
 
         for cls in mro.iter() {
             for &name in cls.attributes.read().keys() {
-                if name.as_str() != "__new__" {
-                    attr_name_set.insert(name);
+                if name != identifier!(ctx, __new__)
+                    && name.as_str().starts_with("__")
+                    && name.as_str().ends_with("__")
+                {
+                    slot_name_set.insert(name);
                 }
             }
         }
         for &name in attrs.keys() {
-            attr_name_set.insert(name);
+            if name.as_str().starts_with("__") && name.as_str().ends_with("__") {
+                slot_name_set.insert(name);
+            }
         }
 
         let new_type = PyRef::new_ref(
@@ -204,11 +219,53 @@ impl PyType {
             None,
         );
 
-        for attr_name in attr_name_set {
-            if attr_name.as_str().starts_with("__") && attr_name.as_str().ends_with("__") {
-                new_type.update_slot(attr_name, true);
-            }
+        for attr_name in slot_name_set {
+            new_type.update_slot::<true>(attr_name, ctx);
         }
+
+        let weakref_type = super::PyWeak::static_type();
+        for base in &new_type.bases {
+            base.subclasses.write().push(
+                new_type
+                    .as_object()
+                    .downgrade_with_weakref_typ_opt(None, weakref_type.to_owned())
+                    .unwrap(),
+            );
+        }
+
+        Ok(new_type)
+    }
+
+    pub fn new_bare_ref(
+        name: &str,
+        base: PyRef<Self>,
+        attrs: PyAttributes,
+        mut slots: PyTypeSlots,
+        metaclass: PyRef<Self>,
+    ) -> Result<PyRef<Self>, String> {
+        if base.slots.flags.has_feature(PyTypeFlags::HAS_DICT) {
+            slots.flags |= PyTypeFlags::HAS_DICT
+        }
+
+        *slots.name.get_mut() = Some(String::from(name));
+
+        let bases = vec![base.clone()];
+        let mro = base.iter_mro().cloned().collect();
+
+        let new_type = PyRef::new_ref(
+            PyType {
+                base: Some(base),
+                bases,
+                mro,
+                subclasses: PyRwLock::default(),
+                attributes: PyRwLock::new(attrs),
+                slots,
+                heaptype_ext: None,
+            },
+            metaclass,
+            None,
+        );
+
         let weakref_type = super::PyWeak::static_type();
         for base in &new_type.bases {
             base.subclasses.write().push(
@@ -693,6 +750,7 @@ impl PyType {
             slots,
             heaptype_ext,
             metatype,
+            &vm.ctx,
         )
         .map_err(|e| vm.new_type_error(e))?;
 
@@ -923,7 +981,11 @@ impl SetAttr for PyType {
             }
         }
         if attr_name.as_str().starts_with("__") && attr_name.as_str().ends_with("__") {
-            zelf.update_slot(attr_name, assign);
+            if assign {
+                zelf.update_slot::<true>(attr_name, &vm.ctx);
+            } else {
+                zelf.update_slot::<false>(attr_name, &vm.ctx);
+            }
         }
         Ok(())
     }
@@ -1177,6 +1239,7 @@ mod tests {
             PyAttributes::default(),
             Default::default(),
             type_type.clone(),
+            &context,
         )
         .unwrap();
         let b = PyType::new_ref(
@@ -1185,6 +1248,7 @@ mod tests {
             PyAttributes::default(),
             Default::default(),
             type_type,
+            &context,
         )
         .unwrap();
 

vm/src/class.rs

@@ -43,9 +43,9 @@ pub trait StaticType {
     where
         Self: PyClassImpl,
     {
-        PyType::new_ref(
+        PyType::new_bare_ref(
             Self::NAME,
-            vec![Self::static_baseclass().to_owned()],
+            Self::static_baseclass().to_owned(),
             Default::default(),
             Self::make_slots(),
             Self::static_metaclass().to_owned(),

vm/src/stdlib/io.rs

@@ -3669,6 +3669,7 @@ mod _io {
             Default::default(),
             Default::default(),
             ctx.types.type_type.to_owned(),
+            ctx,
         )
         .unwrap()
     }

vm/src/types/slot.rs

@@ -356,13 +356,13 @@ fn del_wrapper(zelf: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
 }
 
 impl PyType {
-    pub(crate) fn update_slot(&self, name: &'static PyStrInterned, add: bool) {
+    pub(crate) fn update_slot<const ADD: bool>(&self, name: &'static PyStrInterned, ctx: &Context) {
         debug_assert!(name.as_str().starts_with("__"));
         debug_assert!(name.as_str().ends_with("__"));
 
         macro_rules! toggle_slot {
             ($name:ident, $func:expr) => {{
-                self.slots.$name.store(if add { Some($func) } else { None });
+                self.slots.$name.store(if ADD { Some($func) } else { None });
             }};
         }
 
@@ -382,7 +382,7 @@ impl PyType {
 
         macro_rules! toggle_ext_func {
             ($n1:ident, $n2:ident, $func:expr) => {{
-                self.heaptype_ext.as_ref().unwrap().$n1.$n2.store(if add {
+                self.heaptype_ext.as_ref().unwrap().$n1.$n2.store(if ADD {
                     Some($func)
                 } else {
                     None
@@ -390,87 +390,80 @@ impl PyType {
             }};
         }
 
-        match name.as_str() {
-            "__len__" => {
+        match name {
+            _ if name == identifier!(ctx, __len__) => {
                 update_slot!(as_mapping, slot_as_mapping);
                 toggle_ext_func!(sequence_methods, length, |seq, vm| len_wrapper(seq.obj, vm));
                 update_pointer_slot!(as_sequence, sequence_methods);
             }
-            "__getitem__" => {
+            _ if name == identifier!(ctx, __getitem__) => {
                 update_slot!(as_mapping, slot_as_mapping);
                 toggle_ext_func!(sequence_methods, item, |seq, i, vm| getitem_wrapper(
                     seq.obj, i, vm
                 ));
                 update_pointer_slot!(as_sequence, sequence_methods);
             }
-            "__setitem__" | "__delitem__" => {
+            _ if name == identifier!(ctx, __setitem__) || name == identifier!(ctx, __delitem__) => {
                 update_slot!(as_mapping, slot_as_mapping);
                 toggle_ext_func!(sequence_methods, ass_item, |seq, i, value, vm| {
                     setitem_wrapper(seq.obj, i, value, vm)
                 });
                 update_pointer_slot!(as_sequence, sequence_methods);
             }
-            "__hash__" => {
+            _ if name == identifier!(ctx, __hash__) => {
                 toggle_slot!(hash, hash_wrapper);
             }
-            "__call__" => {
+            _ if name == identifier!(ctx, __call__) => {
                 toggle_slot!(call, call_wrapper);
             }
-            "__getattr__" | "__getattribute__" => {
+            _ if name == identifier!(ctx, __getattr__)
+                || name == identifier!(ctx, __getattribute__) =>
+            {
                 update_slot!(getattro, getattro_wrapper);
             }
-            "__setattr__" | "__delattr__" => {
+            _ if name == identifier!(ctx, __setattr__) || name == identifier!(ctx, __delattr__) => {
                 update_slot!(setattro, setattro_wrapper);
             }
-            "__eq__" | "__ne__" | "__le__" | "__lt__" | "__ge__" | "__gt__" => {
+            _ if name == identifier!(ctx, __eq__)
+                || name == identifier!(ctx, __ne__)
+                || name == identifier!(ctx, __le__)
+                || name == identifier!(ctx, __lt__)
+                || name == identifier!(ctx, __ge__)
+                || name == identifier!(ctx, __gt__) =>
+            {
                 update_slot!(richcompare, richcompare_wrapper);
             }
-            "__iter__" => {
+            _ if name == identifier!(ctx, __iter__) => {
                 toggle_slot!(iter, iter_wrapper);
             }
-            "__next__" => {
+            _ if name == identifier!(ctx, __next__) => {
                 toggle_slot!(iternext, iternext_wrapper);
             }
-            "__get__" => {
+            _ if name == identifier!(ctx, __get__) => {
                 toggle_slot!(descr_get, descr_get_wrapper);
             }
-            "__set__" | "__delete__" => {
+            _ if name == identifier!(ctx, __set__) || name == identifier!(ctx, __delete__) => {
                 update_slot!(descr_set, descr_set_wrapper);
             }
-            "__init__" => {
+            _ if name == identifier!(ctx, __init__) => {
                 toggle_slot!(init, init_wrapper);
             }
-            "__new__" => {
+            _ if name == identifier!(ctx, __new__) => {
                 toggle_slot!(new, new_wrapper);
             }
-            "__del__" => {
+            _ if name == identifier!(ctx, __del__) => {
                 toggle_slot!(del, del_wrapper);
             }
-            "__int__" => {
-                self.heaptype_ext
-                    .as_ref()
-                    .unwrap()
-                    .number_methods
-                    .int
-                    .store(Some(int_wrapper));
+            _ if name == identifier!(ctx, __int__) => {
+                toggle_ext_func!(number_methods, int, int_wrapper);
                 update_pointer_slot!(as_number, number_methods);
             }
-            "__index__" => {
-                self.heaptype_ext
-                    .as_ref()
-                    .unwrap()
-                    .number_methods
-                    .index
-                    .store(Some(index_wrapper));
+            _ if name == identifier!(ctx, __index__) => {
+                toggle_ext_func!(number_methods, index, index_wrapper);
                 update_pointer_slot!(as_number, number_methods);
             }
-            "__float__" => {
-                self.heaptype_ext
-                    .as_ref()
-                    .unwrap()
-                    .number_methods
-                    .float
-                    .store(Some(float_wrapper));
+            _ if name == identifier!(ctx, __float__) => {
+                toggle_ext_func!(number_methods, float, float_wrapper);
                 update_pointer_slot!(as_number, number_methods);
             }
             _ => {}

vm/src/vm/context.rs

@@ -422,6 +422,7 @@ impl Context {
             attrs,
             slots,
             self.types.type_type.to_owned(),
+            self,
         )
         .unwrap()
     }
@@ -446,6 +447,7 @@ impl Context {
             attrs,
             PyBaseException::make_slots(),
             self.types.type_type.to_owned(),
+            self,
         )
         .unwrap()
     }

wasm/lib/src/js_module.rs

@@ -619,7 +619,7 @@ mod _js {
     fn js_error(vm: &VirtualMachine) -> PyTypeRef {
         let ctx = &vm.ctx;
         let js_error = PyRef::leak(
-            PyType::new_simple_ref("JSError", &vm.ctx.exceptions.exception_type.to_owned())
+            PyType::new_simple_ref("JSError", &vm.ctx.exceptions.exception_type.to_owned(), ctx)
                 .unwrap(),
         );
         extend_class!(ctx, js_error, {