proto heaptypeext

Author: qingshi163

derive/src/pyclass.rs

@@ -564,10 +564,15 @@ where
         let slot_name = slot_ident.to_string();
         let tokens = {
             const NON_ATOMIC_SLOTS: &[&str] = &["as_buffer"];
+            const POINTER_SLOTS: &[&str] = &["as_number"];
             if NON_ATOMIC_SLOTS.contains(&slot_name.as_str()) {
                 quote_spanned! { span =>
                     slots.#slot_ident = Some(Self::#ident as _);
                 }
+            } else if POINTER_SLOTS.contains(&slot_name.as_str()) {
+                quote_spanned! { span =>
+                    slots.#slot_ident.store(Some(NonNull::from(Self::#ident())));
+                }
             } else {
                 quote_spanned! { span =>
                     slots.#slot_ident.store(Some(Self::#ident as _));

vm/src/builtins/float.rs

@@ -60,11 +60,11 @@ impl From<f64> for PyFloat {
 
 impl PyObject {
     pub fn try_float_opt(&self, vm: &VirtualMachine) -> PyResult<Option<PyRef<PyFloat>>> {
-        PyNumber::new(self, vm).float_opt(vm)
+        PyNumber::from(self).float_opt(vm)
     }
 
     pub fn try_float(&self, vm: &VirtualMachine) -> PyResult<PyRef<PyFloat>> {
-        PyNumber::new(self, vm).float(vm)
+        PyNumber::from(self).float(vm)
     }
 }
 

vm/src/builtins/int.rs

@@ -266,7 +266,7 @@ impl Constructor for PyInt {
                     val
                 };
 
-                PyNumber::new(val.as_ref(), vm)
+                PyNumber::from(val.as_ref())
                     .int(vm)
                     .map(|x| x.as_bigint().clone())
             }

vm/src/builtins/type.rs

@@ -2,11 +2,6 @@ use super::{
     mappingproxy::PyMappingProxy, object, union_, PyClassMethod, PyDictRef, PyList, PyStaticMethod,
     PyStr, PyStrInterned, PyStrRef, PyTuple, PyTupleRef, PyWeak,
 };
-use crate::common::{
-    ascii,
-    borrow::BorrowedValue,
-    lock::{PyRwLock, PyRwLockReadGuard},
-};
 use crate::{
     builtins::PyBaseExceptionRef,
     class::{PyClassImpl, StaticType},
@@ -15,9 +10,17 @@ use crate::{
     types::{Callable, GetAttr, PyTypeFlags, PyTypeSlots, SetAttr},
     AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
 };
+use crate::{
+    common::{
+        ascii,
+        borrow::BorrowedValue,
+        lock::{PyRwLock, PyRwLockReadGuard},
+    },
+    protocol::PyNumberMethods,
+};
 use indexmap::{map::Entry, IndexMap};
 use itertools::Itertools;
-use std::{borrow::Borrow, collections::HashSet, fmt, ops::Deref};
+use std::{borrow::Borrow, collections::HashSet, fmt, ops::Deref, pin::Pin};
 
 /// type(object_or_name, bases, dict)
 /// type(object) -> the object's type
@@ -30,10 +33,23 @@ pub struct PyType {
     pub subclasses: PyRwLock<Vec<PyRef<PyWeak>>>,
     pub attributes: PyRwLock<PyAttributes>,
     pub slots: PyTypeSlots,
+    pub heaptype_ext: Option<Pin<Box<HeapTypeExt>>>,
+}
+
+#[derive(Default)]
+pub struct HeapTypeExt {
+    pub number_methods: PyNumberMethods,
 }
 
 pub type PyTypeRef = PyRef<PyType>;
 
+cfg_if::cfg_if! {
+    if #[cfg(feature = "threading")] {
+        unsafe impl Send for PyType {}
+        unsafe impl Sync for PyType {}
+    }
+}
+
 /// For attributes we do not use a dict, but an IndexMap, which is an Hash Table
 /// that maintains order and is compatible with the standard HashMap  This is probably
 /// faster and only supports strings as keys.
@@ -112,6 +128,7 @@ impl PyType {
                 subclasses: PyRwLock::default(),
                 attributes: PyRwLock::new(attrs),
                 slots,
+                heaptype_ext: Some(Pin::new(Box::new(HeapTypeExt::default()))),
             },
             metaclass,
             None,

vm/src/function/number.rs

@@ -82,7 +82,7 @@ impl Deref for ArgIntoFloat {
 impl TryFromObject for ArgIntoFloat {
     // Equivalent to PyFloat_AsDouble.
     fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
-        let value = PyNumber::new(obj.as_ref(), vm).float(vm)?.to_f64();
+        let value = PyNumber::from(obj.as_ref()).float(vm)?.to_f64();
         Ok(ArgIntoFloat { value })
     }
 }

vm/src/object/core.rs

@@ -1105,6 +1105,7 @@ pub(crate) fn init_type_hierarchy() -> (PyTypeRef, PyTypeRef, PyTypeRef) {
             subclasses: PyRwLock::default(),
             attributes: PyRwLock::new(Default::default()),
             slots: PyType::make_slots(),
+            heaptype_ext: None,
         };
         let object_payload = PyType {
             base: None,
@@ -1113,6 +1114,7 @@ pub(crate) fn init_type_hierarchy() -> (PyTypeRef, PyTypeRef, PyTypeRef) {
             subclasses: PyRwLock::default(),
             attributes: PyRwLock::new(Default::default()),
             slots: object::PyBaseObject::make_slots(),
+            heaptype_ext: None,
         };
         let type_type_ptr = Box::into_raw(Box::new(partially_init!(
             PyInner::<PyType> {
@@ -1173,6 +1175,7 @@ pub(crate) fn init_type_hierarchy() -> (PyTypeRef, PyTypeRef, PyTypeRef) {
         subclasses: PyRwLock::default(),
         attributes: PyRwLock::default(),
         slots: PyWeak::make_slots(),
+        heaptype_ext: None,
     };
     let weakref_type = PyRef::new_ref(weakref_type, type_type.clone(), None);
 

vm/src/protocol/number.rs

@@ -1,4 +1,7 @@
+use std::ptr::NonNull;
+
 use crossbeam_utils::atomic::AtomicCell;
+use once_cell::sync::OnceCell;
 
 use crate::{
     builtins::{int, PyByteArray, PyBytes, PyComplex, PyFloat, PyInt, PyIntRef, PyStr},
@@ -12,6 +15,7 @@ type UnaryFunc<R = PyObjectRef> = AtomicCell<Option<fn(&PyNumber, &VirtualMachin
 type BinaryFunc<R = PyObjectRef> =
     AtomicCell<Option<fn(&PyNumber, &PyObject, &VirtualMachine) -> PyResult<R>>>;
 
+#[derive(Default)]
 pub struct PyNumberMethods {
     /* Number implementations must check *both*
     arguments for proper type and implement the necessary conversions
@@ -124,36 +128,38 @@ impl PyNumberMethods {
 pub struct PyNumber<'a> {
     pub obj: &'a PyObject,
     // some fast path do not need methods, so we do lazy initialize
-    pub methods: Option<&'static PyNumberMethods>,
+    methods: OnceCell<NonNull<PyNumberMethods>>,
 }
 
-impl<'a> PyNumber<'a> {
-    pub fn new(obj: &'a PyObject, vm: &VirtualMachine) -> Self {
+impl<'a> From<&'a PyObject> for PyNumber<'a> {
+    fn from(obj: &'a PyObject) -> Self {
         Self {
             obj,
-            methods: Self::find_methods(obj, vm),
+            methods: OnceCell::new(),
         }
     }
 }
 
 impl PyNumber<'_> {
-    pub fn find_methods(obj: &PyObject, vm: &VirtualMachine) -> Option<&'static PyNumberMethods> {
-        let as_number = obj.class().mro_find_map(|x| x.slots.as_number.load());
-        as_number.map(|f| f(obj, vm))
+    pub fn methods(&self) -> &PyNumberMethods {
+        let as_number = self.methods.get_or_init(|| {
+            Self::find_methods(self.obj).unwrap_or(NonNull::from(&PyNumberMethods::NOT_IMPLEMENTED))
+        });
+        unsafe { as_number.as_ref() }
     }
 
-    pub fn methods(&self) -> &'static PyNumberMethods {
-        self.methods.unwrap_or(&PyNumberMethods::NOT_IMPLEMENTED)
+    fn find_methods<'a>(obj: &'a PyObject) -> Option<NonNull<PyNumberMethods>> {
+        obj.class().mro_find_map(|x| x.slots.as_number.load())
     }
 
     // PyNumber_Check
-    pub fn check(obj: &PyObject, vm: &VirtualMachine) -> bool {
-        Self::find_methods(obj, vm).map_or(false, |methods| {
-            methods.int.load().is_some()
-                || methods.index.load().is_some()
-                || methods.float.load().is_some()
-                || obj.payload_is::<PyComplex>()
-        })
+    pub fn check<'a>(obj: &'a PyObject, vm: &VirtualMachine) -> bool {
+        let num = PyNumber::from(obj);
+        let methods = num.methods();
+        methods.int.load().is_some()
+            || methods.index.load().is_some()
+            || methods.float.load().is_some()
+            || obj.payload_is::<PyComplex>()
     }
 
     // PyIndex_Check
@@ -207,7 +213,7 @@ impl PyNumber<'_> {
             //     vm,
             // )?;
             let ret = f.invoke((), vm)?;
-            PyNumber::new(ret.as_ref(), vm).index(vm).map_err(|_| {
+            PyNumber::from(ret.as_ref()).index(vm).map_err(|_| {
                 vm.new_type_error(format!(
                     "__trunc__ returned non-Integral (type {})",
                     ret.class()

vm/src/types/slot.rs

@@ -15,6 +15,7 @@ use crate::{
 };
 use crossbeam_utils::atomic::AtomicCell;
 use num_traits::{Signed, ToPrimitive};
+use std::ptr::NonNull;
 use std::{borrow::Borrow, cmp::Ordering};
 
 // The corresponding field in CPython is `tp_` prefixed.
@@ -30,7 +31,7 @@ pub struct PyTypeSlots {
     // Methods to implement standard operations
 
     // Method suites for standard classes
-    pub as_number: AtomicCell<Option<AsNumberFunc>>,
+    pub as_number: AtomicCell<Option<NonNull<PyNumberMethods>>>,
     pub as_sequence: AtomicCell<Option<AsSequenceFunc>>,
     pub as_mapping: AtomicCell<Option<AsMappingFunc>>,
 
@@ -139,7 +140,6 @@ impl Default for PyTypeFlags {
 
 pub(crate) type GenericMethod = fn(&PyObject, FuncArgs, &VirtualMachine) -> PyResult;
 pub(crate) type AsMappingFunc = fn(&PyObject, &VirtualMachine) -> &'static PyMappingMethods;
-pub(crate) type AsNumberFunc = fn(&PyObject, &VirtualMachine) -> &'static PyNumberMethods;
 pub(crate) type HashFunc = fn(&PyObject, &VirtualMachine) -> PyResult<PyHash>;
 // CallFunc = GenericMethod
 pub(crate) type GetattroFunc = fn(&PyObject, PyStrRef, &VirtualMachine) -> PyResult;
@@ -205,15 +205,6 @@ fn slot_as_sequence(zelf: &PyObject, vm: &VirtualMachine) -> &'static PySequence
     PySequenceMethods::generic(has_length, has_ass_item)
 }
 
-fn slot_as_number(zelf: &PyObject, vm: &VirtualMachine) -> &'static PyNumberMethods {
-    let (has_int, has_float, has_index) = (
-        zelf.class().has_attr(identifier!(vm, __int__)),
-        zelf.class().has_attr(identifier!(vm, __float__)),
-        zelf.class().has_attr(identifier!(vm, __index__)),
-    );
-    PyNumberMethods::generic(has_int, has_float, has_index)
-}
-
 fn hash_wrapper(zelf: &PyObject, vm: &VirtualMachine) -> PyResult<PyHash> {
     let hash_obj = vm.call_special_method(zelf.to_owned(), identifier!(vm, __hash__), ())?;
     match hash_obj.payload_if_subclass::<PyInt>(vm) {
@@ -374,7 +365,7 @@ impl PyType {
                 update_slot!(del, del_wrapper);
             }
             "__int__" | "__index__" | "__float__" => {
-                update_slot!(as_number, slot_as_number);
+                // update_slot!(as_number, slot_as_number);
             }
             _ => {}
         }
@@ -884,7 +875,7 @@ pub trait AsNumber: PyPayload {
 
     #[inline]
     #[pyslot]
-    fn as_number(_zelf: &PyObject, _vm: &VirtualMachine) -> &'static PyNumberMethods {
+    fn as_number() -> &'static PyNumberMethods {
         &Self::AS_NUMBER
     }
 

vm/src/vm/vm_ops.rs

@@ -10,13 +10,11 @@ use crate::{
 /// Collection of operators
 impl VirtualMachine {
     pub fn to_index_opt(&self, obj: PyObjectRef) -> Option<PyResult<PyIntRef>> {
-        PyNumber::new(obj.as_ref(), self)
-            .index_opt(self)
-            .transpose()
+        PyNumber::from(obj.as_ref()).index_opt(self).transpose()
     }
 
     pub fn to_index(&self, obj: &PyObject) -> PyResult<PyIntRef> {
-        PyNumber::new(obj, self).index(self)
+        PyNumber::from(obj).index(self)
     }
 
     #[inline]