Add subclass type flags (#9327)

crates/capi/src/object/pytype.rs

@@ -325,6 +325,7 @@ pub extern "C" fn PyType_Freeze(_ty: *mut PyTypeObject) -> c_int {
 #[cfg(test)]
 mod tests {
     use pyo3::IntoPyObjectExt;
+    use pyo3::exceptions::PyValueError;
     use pyo3::ffi;
     use pyo3::prelude::*;
     use pyo3::types::{PyDict, PyInt, PyList, PyString, PyType, PyTypeMethods};
@@ -542,4 +543,22 @@ mod tests {
             assert_eq!(type_obj.qualname().unwrap(), "somemodule.MySpecClass");
         });
     }
+
+    #[test]
+    fn fast_subclass_check_using_flags() {
+        Python::attach(|py| {
+            assert!(unsafe {
+                ffi::PyType_HasFeature(
+                    py.get_type::<PyDict>().as_type_ptr(),
+                    ffi::Py_TPFLAGS_DICT_SUBCLASS,
+                ) != 0
+            });
+            assert!(unsafe {
+                ffi::PyType_HasFeature(
+                    py.get_type::<PyValueError>().as_type_ptr(),
+                    ffi::Py_TPFLAGS_BASE_EXC_SUBCLASS,
+                ) != 0
+            });
+        })
+    }
 }

crates/vm/src/builtins/bytes.rs

@@ -235,7 +235,7 @@ impl PyBytes {
 
 #[pyclass(
     itemsize = 1,
-    flags(BASETYPE, _MATCH_SELF),
+    flags(BASETYPE, _MATCH_SELF, BYTES_SUBCLASS),
     with(
         Py,
         PyRef,

crates/vm/src/builtins/dict.rs

@@ -804,7 +804,7 @@ impl PyDict {
         AsMapping,
         Representable
     ),
-    flags(BASETYPE, MAPPING, _MATCH_SELF)
+    flags(BASETYPE, MAPPING, _MATCH_SELF, DICT_SUBCLASS)
 )]
 impl PyDict {}
 

crates/vm/src/builtins/int.rs

@@ -531,7 +531,7 @@ struct RoundArgs {
 
 #[pyclass(
     itemsize = 4,
-    flags(BASETYPE, _MATCH_SELF),
+    flags(BASETYPE, _MATCH_SELF, LONG_SUBCLASS),
     with(PyRef, Comparable, Hashable, Constructor, AsNumber, Representable)
 )]
 impl Py<PyInt> {

crates/vm/src/builtins/list.rs

@@ -417,7 +417,7 @@ impl PyList {
         AsSequence,
         Representable
     ),
-    flags(BASETYPE, SEQUENCE, _MATCH_SELF)
+    flags(BASETYPE, SEQUENCE, _MATCH_SELF, LIST_SUBCLASS)
 )]
 impl Py<PyList> {
     #[pymethod]

crates/vm/src/builtins/str.rs

@@ -897,7 +897,7 @@ impl PyStr {
 }
 
 #[pyclass(
-    flags(BASETYPE, _MATCH_SELF),
+    flags(BASETYPE, _MATCH_SELF, UNICODE_SUBCLASS),
     with(
         AsMapping,
         AsNumber,

crates/vm/src/builtins/tuple.rs

@@ -543,7 +543,7 @@ impl PyTuple {
 
 #[pyclass(
     itemsize = core::mem::size_of::<crate::PyObjectRef>(),
-    flags(BASETYPE, SEQUENCE, _MATCH_SELF),
+    flags(BASETYPE, SEQUENCE, _MATCH_SELF, TUPLE_SUBCLASS),
     with(AsMapping, AsNumber, AsSequence, Hashable, Comparable, Iterable, Constructor, Representable)
 )]
 impl Py<PyTuple> {

crates/vm/src/builtins/type.rs

@@ -908,6 +908,58 @@ impl PyType {
         }
     }
 
+    fn inherit_subclass_flags(slots: &mut PyTypeSlots, bases: &[PyRef<Self>]) {
+        let mut inherited = PyTypeFlags::empty();
+        for base in bases {
+            inherited |= base.slots.flags.load() & PyTypeFlags::SUBCLASS_FLAGS;
+        }
+        slots.flags |= inherited;
+    }
+
+    fn compute_subclass_flags(&self, ctx: &Context) -> PyTypeFlags {
+        let is_subclass = |base: &Py<Self>| self.mro.read().iter().any(|cls| cls.is(base));
+        let mut flags = PyTypeFlags::empty();
+        if is_subclass(ctx.types.int_type) {
+            flags.insert(PyTypeFlags::LONG_SUBCLASS);
+        }
+        if is_subclass(ctx.types.list_type) {
+            flags.insert(PyTypeFlags::LIST_SUBCLASS);
+        }
+        if is_subclass(ctx.types.tuple_type) {
+            flags.insert(PyTypeFlags::TUPLE_SUBCLASS);
+        }
+        if is_subclass(ctx.types.bytes_type) {
+            flags.insert(PyTypeFlags::BYTES_SUBCLASS);
+        }
+        if is_subclass(ctx.types.str_type) {
+            flags.insert(PyTypeFlags::UNICODE_SUBCLASS);
+        }
+        if is_subclass(ctx.types.dict_type) {
+            flags.insert(PyTypeFlags::DICT_SUBCLASS);
+        }
+        if is_subclass(ctx.exceptions.base_exception_type) {
+            flags.insert(PyTypeFlags::BASE_EXC_SUBCLASS);
+        }
+        if is_subclass(ctx.types.type_type) {
+            flags.insert(PyTypeFlags::TYPE_SUBCLASS);
+        }
+        flags
+    }
+
+    fn update_subclass_flags_recursive(&self, ctx: &Context) {
+        self.slots.flags.replace_masked(
+            PyTypeFlags::SUBCLASS_FLAGS,
+            self.compute_subclass_flags(ctx),
+        );
+        for weak_ref in self.subclasses.read().iter() {
+            if let Some(subclass) = weak_ref.upgrade()
+                && let Some(subclass) = subclass.downcast_ref::<Self>()
+            {
+                subclass.update_subclass_flags_recursive(ctx);
+            }
+        }
+    }
+
     pub fn has_patma_collection_flag(&self, flag: u8) -> bool {
         debug_assert!(flag == PyTypeFlags::SEQUENCE || flag == PyTypeFlags::MAPPING);
         self.slots.flags.has_feature(flag)
@@ -1041,6 +1093,8 @@ impl PyType {
             slots.flags |= WEAKREF_FLAGS;
         }
 
+        Self::inherit_subclass_flags(&mut slots, bases.as_slice());
+
         // Inherit SEQUENCE and MAPPING flags from base classes
         Self::inherit_patma_flags(&mut slots, bases.as_slice());
 
@@ -1113,6 +1167,7 @@ impl PyType {
         // Inherit SEQUENCE and MAPPING flags from base class
         // For static types, we only have a single base
         Self::inherit_patma_flags(&mut slots, core::slice::from_ref(&base));
+        Self::inherit_subclass_flags(&mut slots, core::slice::from_ref(&base));
 
         if slots.basicsize == 0 {
             slots.basicsize = base.slots.basicsize;
@@ -1792,7 +1847,7 @@ impl PyType {
         AsNumber,
         Representable
     ),
-    flags(BASETYPE, HAS_DICT, HAS_WEAKREF)
+    flags(BASETYPE, HAS_DICT, HAS_WEAKREF, TYPE_SUBCLASS)
 )]
 impl PyType {}
 
@@ -2201,6 +2256,7 @@ impl Constructor for PyType {
         if custom_mro {
             mro_internal(&typ, vm)?;
             typ.init_slots(&vm.ctx);
+            typ.update_subclass_flags_recursive(&vm.ctx);
         }
 
         if let Some(ref slots) = heaptype_slots {
@@ -2741,6 +2797,7 @@ impl Py<PyType> {
             if core::ptr::eq(&**zelf.bases.read(), &*bases) {
                 remove_as_subclass(old_bases.as_slice(), &mut retired);
                 add_as_subclass(bases.as_slice());
+                zelf.update_subclass_flags_recursive(&vm.ctx);
                 zelf.update_all_slots(&vm.ctx);
             }
 

crates/vm/src/exceptions.rs

@@ -782,7 +782,7 @@ impl PyBaseException {
 
 #[pyclass(
     with(Py, PyRef, Constructor, Initializer, Representable),
-    flags(BASETYPE, HAS_DICT)
+    flags(BASETYPE, HAS_DICT, BASE_EXC_SUBCLASS)
 )]
 impl PyBaseException {}
 

crates/vm/src/types/slot.rs

@@ -266,6 +266,14 @@ bitflagset::bitflagset! {
         // (bool, int, float, str, bytes, bytearray, list, tuple, dict, set, frozenset).
         // This is not a stable API.
         const _MATCH_SELF = 22;
+        const LONG_SUBCLASS = 24;
+        const LIST_SUBCLASS = 25;
+        const TUPLE_SUBCLASS = 26;
+        const BYTES_SUBCLASS = 27;
+        const UNICODE_SUBCLASS = 28;
+        const DICT_SUBCLASS = 29;
+        const BASE_EXC_SUBCLASS = 30;
+        const TYPE_SUBCLASS = 31;
         const HAS_DICT = 40;
         const HAS_WEAKREF = 41;
         #[cfg(debug_assertions)]
@@ -288,6 +296,17 @@ impl PyTypeFlags {
     );
 
     pub const COLLECTION: Self = Self::from_slice(&[Self::SEQUENCE, Self::MAPPING]);
+
+    pub const SUBCLASS_FLAGS: Self = Self::from_slice(&[
+        Self::LONG_SUBCLASS,
+        Self::LIST_SUBCLASS,
+        Self::TUPLE_SUBCLASS,
+        Self::BYTES_SUBCLASS,
+        Self::UNICODE_SUBCLASS,
+        Self::DICT_SUBCLASS,
+        Self::BASE_EXC_SUBCLASS,
+        Self::TYPE_SUBCLASS,
+    ]);
 }
 
 impl AtomicPyTypeFlags {