refactor replace nonnull with pointerslot

Author: qingshi163

derive/src/pyclass.rs

@@ -571,7 +571,7 @@ where
                 }
             } else if POINTER_SLOTS.contains(&slot_name.as_str()) {
                 quote_spanned! { span =>
-                    slots.#slot_ident.store(Some(NonNull::from(Self::#ident())));
+                    slots.#slot_ident.store(Some(PointerSlot::from(Self::#ident())));
                 }
             } else {
                 quote_spanned! { span =>

vm/src/builtins/type.rs

@@ -20,7 +20,7 @@ use crate::{
 };
 use indexmap::{map::Entry, IndexMap};
 use itertools::Itertools;
-use std::{borrow::Borrow, collections::HashSet, fmt, ops::Deref, pin::Pin};
+use std::{borrow::Borrow, collections::HashSet, fmt, ops::Deref, pin::Pin, ptr::NonNull};
 
 /// type(object_or_name, bases, dict)
 /// type(object) -> the object's type
@@ -41,6 +41,39 @@ pub struct HeapTypeExt {
     pub number_methods: PyNumberMethods,
 }
 
+pub struct PointerSlot<T>(NonNull<T>);
+
+impl<T> Clone for PointerSlot<T> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
+impl<T> Copy for PointerSlot<T> {}
+
+impl<T> From<&'static T> for PointerSlot<T> {
+    fn from(x: &'static T) -> Self {
+        Self(NonNull::from(x))
+    }
+}
+
+impl<T> AsRef<T> for PointerSlot<T> {
+    fn as_ref(&self) -> &T {
+        unsafe { self.0.as_ref() }
+    }
+}
+
+impl<T> PointerSlot<T> {
+    pub unsafe fn from_heaptype<F>(typ: &PyType, f: F) -> Option<Self>
+    where
+        F: FnOnce(&HeapTypeExt) -> &T,
+    {
+        typ.heaptype_ext
+            .as_ref()
+            .map(|ext| Self(NonNull::from(f(ext))))
+    }
+}
+
 pub type PyTypeRef = PyRef<PyType>;
 
 cfg_if::cfg_if! {

vm/src/protocol/number.rs

@@ -1,10 +1,10 @@
-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},
+    builtins::{
+        int, type_::PointerSlot, PyByteArray, PyBytes, PyComplex, PyFloat, PyInt, PyIntRef, PyStr,
+    },
     function::ArgBytesLike,
     stdlib::warnings,
     AsObject, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromBorrowedObject,
@@ -109,7 +109,7 @@ impl PyNumberMethods {
 pub struct PyNumber<'a> {
     pub obj: &'a PyObject,
     // some fast path do not need methods, so we do lazy initialize
-    methods: OnceCell<NonNull<PyNumberMethods>>,
+    methods: OnceCell<PointerSlot<PyNumberMethods>>,
 }
 
 impl<'a> From<&'a PyObject> for PyNumber<'a> {
@@ -124,13 +124,16 @@ impl<'a> From<&'a PyObject> for PyNumber<'a> {
 impl PyNumber<'_> {
     pub fn methods(&self) -> &PyNumberMethods {
         static GLOBAL_NOT_IMPLEMENTED: PyNumberMethods = PyNumberMethods::NOT_IMPLEMENTED;
-        let as_number = self.methods.get_or_init(|| {
-            Self::find_methods(self.obj).unwrap_or_else(|| NonNull::from(&GLOBAL_NOT_IMPLEMENTED))
-        });
-        unsafe { as_number.as_ref() }
+
+        self.methods
+            .get_or_init(|| {
+                Self::find_methods(self.obj)
+                    .unwrap_or_else(|| PointerSlot::from(&GLOBAL_NOT_IMPLEMENTED))
+            })
+            .as_ref()
     }
 
-    fn find_methods(obj: &PyObject) -> Option<NonNull<PyNumberMethods>> {
+    fn find_methods(obj: &PyObject) -> Option<PointerSlot<PyNumberMethods>> {
         obj.class().mro_find_map(|x| x.slots.as_number.load())
     }
 

vm/src/types/slot.rs

@@ -1,6 +1,6 @@
 use crate::common::{hash::PyHash, lock::PyRwLock};
 use crate::{
-    builtins::{PyFloat, PyInt, PyStrInterned, PyStrRef, PyType, PyTypeRef},
+    builtins::{type_::PointerSlot, PyFloat, PyInt, PyStrInterned, PyStrRef, PyType, PyTypeRef},
     bytecode::ComparisonOperator,
     convert::ToPyResult,
     function::Either,
@@ -15,7 +15,6 @@ 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.
@@ -31,7 +30,7 @@ pub struct PyTypeSlots {
     // Methods to implement standard operations
 
     // Method suites for standard classes
-    pub as_number: AtomicCell<Option<NonNull<PyNumberMethods>>>,
+    pub as_number: AtomicCell<Option<PointerSlot<PyNumberMethods>>>,
     pub as_sequence: AtomicCell<Option<AsSequenceFunc>>,
     pub as_mapping: AtomicCell<Option<AsMappingFunc>>,
 
@@ -356,11 +355,9 @@ impl PyType {
 
         macro_rules! update_pointer_slot {
             ($name:ident, $pointed:ident) => {{
-                self.slots.$name.store(
-                    self.heaptype_ext
-                        .as_ref()
-                        .map(|ext| NonNull::from(&ext.$pointed)),
-                );
+                self.slots
+                    .$name
+                    .store(unsafe { PointerSlot::from_heaptype(self, |ext| &ext.$pointed) });
             }};
         }
         match name.as_str() {