Make StaticCell less complex to use

common/src/static_cell.rs

@@ -1,16 +1,66 @@
 #[cfg(not(feature = "threading"))]
 mod non_threading {
     use crate::lock::OnceCell;
-    pub type StaticCell<T> = std::thread::LocalKey<OnceCell<&'static T>>;
+    use std::thread::LocalKey;
+
+    pub struct StaticCell<T: 'static> {
+        inner: &'static LocalKey<OnceCell<&'static T>>,
+    }
+
+    fn leak<T>(x: T) -> &'static T {
+        Box::leak(Box::new(x))
+    }
+
+    impl<T> StaticCell<T> {
+        #[doc(hidden)]
+        pub const fn _from_localkey(inner: &'static LocalKey<OnceCell<&'static T>>) -> Self {
+            Self { inner }
+        }
+
+        pub fn get(&'static self) -> Option<&'static T> {
+            self.inner.with(|x| x.get().copied())
+        }
+
+        pub fn set(&'static self, value: T) -> Result<(), T> {
+            // thread-safe because it's a unsync::OnceCell
+            self.inner.with(|x| {
+                if x.get().is_some() {
+                    Err(value)
+                } else {
+                    // will never fail
+                    let _ = x.set(leak(value));
+                    Ok(())
+                }
+            })
+        }
+
+        pub fn get_or_init<F>(&'static self, f: F) -> &'static T
+        where
+            F: FnOnce() -> T,
+        {
+            self.inner.with(|x| *x.get_or_init(|| leak(f())))
+        }
+
+        pub fn get_or_try_init<F, E>(&'static self, f: F) -> Result<&'static T, E>
+        where
+            F: FnOnce() -> Result<T, E>,
+        {
+            self.inner
+                .with(|x| x.get_or_try_init(|| f().map(leak)).map(|&x| x))
+        }
+    }
 
     #[macro_export]
-    macro_rules! static_cells {
-        // process multiple declarations
-        ($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty;) => (
-            std::thread_local! {
-                $(#[$attr])* $vis static $name: $crate::lock::OnceCell<&'static $t> = $crate::lock::OnceCell::new();
-            }
-        );
+    macro_rules! static_cell {
+        ($($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty;)+) => {
+            $($(#[$attr])*
+            $vis static $name: $crate::static_cell::StaticCell<$t> = {
+                ::std::thread_local! {
+                     $vis static $name: $crate::lock::OnceCell<&'static $t> = $crate::lock::OnceCell::new();
+                }
+                $crate::static_cell::StaticCell::_from_localkey(&$name)
+            };)+
+        };
     }
 }
 #[cfg(not(feature = "threading"))]
@@ -20,57 +70,47 @@ pub use non_threading::*;
 mod threading {
     use crate::lock::OnceCell;
 
-    pub struct StaticKey<T> {
-        inner: T,
+    pub struct StaticCell<T: 'static> {
+        inner: OnceCell<T>,
     }
 
-    impl<T> StaticKey<T> {
-        pub const fn new(inner: T) -> Self {
+    impl<T> StaticCell<T> {
+        #[doc(hidden)]
+        pub const fn _from_oncecell(inner: OnceCell<T>) -> Self {
             Self { inner }
         }
 
-        pub fn with<F, R>(&self, f: F) -> R
+        pub fn get(&'static self) -> Option<&'static T> {
+            self.inner.get()
+        }
+
+        pub fn set(&'static self, value: T) -> Result<(), T> {
+            self.inner.set(value)
+        }
+
+        pub fn get_or_init<F>(&'static self, f: F) -> &'static T
         where
-            F: FnOnce(&T) -> R,
+            F: FnOnce() -> T,
         {
-            f(&self.inner)
+            self.inner.get_or_init(f)
+        }
+
+        pub fn get_or_try_init<F, E>(&'static self, f: F) -> Result<&'static T, E>
+        where
+            F: FnOnce() -> Result<T, E>,
+        {
+            self.inner.get_or_try_init(f)
         }
     }
 
-    pub type StaticCell<T> = StaticKey<OnceCell<&'static T>>;
-
     #[macro_export]
-    macro_rules! static_cells {
-        // process multiple declarations
-        ($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty;) => (
-            $(#[$attr])* $vis static $name: $crate::static_cell::StaticKey<$crate::lock::OnceCell<&'static $t>> = $crate::static_cell::StaticKey::new($crate::lock::OnceCell::new());
-        );
+    macro_rules! static_cell {
+        ($($(#[$attr:meta])* $vis:vis static $name:ident: $t:ty;)+) => {
+            $($(#[$attr])*
+            $vis static $name: $crate::static_cell::StaticCell<$t> =
+                $crate::static_cell::StaticCell::_from_oncecell($crate::lock::OnceCell::new());)+
+        };
     }
 }
 #[cfg(feature = "threading")]
 pub use threading::*;
-
-pub fn get<T>(cell: &'static StaticCell<T>) -> Option<&'static T> {
-    cell.with(|cell| cell.get().copied())
-}
-
-pub fn init_expect<T>(cell: &'static StaticCell<T>, value: T, msg: &'static str) -> &'static T {
-    cell.with(|cell| {
-        let static_ref = Box::leak(Box::new(value)) as &_;
-        cell.set(static_ref)
-            .unwrap_or_else(|_| panic!("double initializing '{}'", msg));
-        static_ref
-    })
-}
-
-pub fn get_or_init<T, F>(cell: &'static StaticCell<T>, f: F) -> &'static T
-where
-    F: FnOnce() -> T,
-{
-    cell.with(|cell| {
-        *cell.get_or_init(|| {
-            let value = f();
-            Box::leak(Box::new(value))
-        })
-    })
-}

derive/src/pyclass.rs

@@ -197,8 +197,7 @@ fn generate_class_def(
 
         impl ::rustpython_vm::pyobject::StaticType for #ident {
             fn static_cell() -> &'static ::rustpython_common::static_cell::StaticCell<::rustpython_vm::builtins::PyTypeRef> {
-                use ::rustpython_common::static_cells;
-                static_cells! {
+                ::rustpython_common::static_cell! {
                     static CELL: ::rustpython_vm::builtins::PyTypeRef;
                 }
                 &CELL

vm/src/pyobject.rs

@@ -111,63 +111,59 @@ impl PyContext {
     pub const INT_CACHE_POOL_MIN: i32 = -5;
     pub const INT_CACHE_POOL_MAX: i32 = 256;
 
+    fn init() -> Self {
+        flame_guard!("init PyContext");
+        let types = TypeZoo::init();
+        let exceptions = exceptions::ExceptionZoo::init();
+
+        fn create_object<T: PyObjectPayload + PyValue>(payload: T, cls: &PyTypeRef) -> PyRef<T> {
+            PyRef::new_ref(payload, cls.clone(), None)
+        }
+
+        let none = create_object(PyNone, PyNone::static_type());
+        let ellipsis = create_object(PyEllipsis, PyEllipsis::static_type());
+        let not_implemented = create_object(PyNotImplemented, PyNotImplemented::static_type());
+
+        let int_cache_pool = (Self::INT_CACHE_POOL_MIN..=Self::INT_CACHE_POOL_MAX)
+            .map(|v| PyRef::new_ref(PyInt::from(BigInt::from(v)), types.int_type.clone(), None))
+            .collect();
+
+        let true_value = create_object(PyInt::from(1), &types.bool_type);
+        let false_value = create_object(PyInt::from(0), &types.bool_type);
+
+        let empty_tuple = create_object(
+            PyTuple::_new(Vec::new().into_boxed_slice()),
+            &types.tuple_type,
+        );
+
+        let tp_new_wrapper = create_object(
+            PyNativeFuncDef::from(pytype::tp_new_wrapper.into_func()).into_function(),
+            &types.builtin_function_or_method_type,
+        )
+        .into_object();
+
+        let context = PyContext {
+            true_value,
+            false_value,
+            not_implemented,
+            none,
+            empty_tuple,
+            ellipsis,
+
+            types,
+            exceptions,
+            int_cache_pool,
+            tp_new_wrapper,
+        };
+        TypeZoo::extend(&context);
+        exceptions::ExceptionZoo::extend(&context);
+        context
+    }
     pub fn new() -> Self {
-        use rustpython_common::static_cells;
-        static_cells! {
+        rustpython_common::static_cell! {
             static CONTEXT: PyContext;
         }
-        static_cell::get_or_init(&CONTEXT, || {
-            flame_guard!("init PyContext");
-            let types = TypeZoo::init();
-            let exceptions = exceptions::ExceptionZoo::init();
-
-            fn create_object<T: PyObjectPayload + PyValue>(
-                payload: T,
-                cls: &PyTypeRef,
-            ) -> PyRef<T> {
-                PyRef::new_ref(payload, cls.clone(), None)
-            }
-
-            let none = create_object(PyNone, PyNone::static_type());
-            let ellipsis = create_object(PyEllipsis, PyEllipsis::static_type());
-            let not_implemented = create_object(PyNotImplemented, PyNotImplemented::static_type());
-
-            let int_cache_pool = (Self::INT_CACHE_POOL_MIN..=Self::INT_CACHE_POOL_MAX)
-                .map(|v| PyRef::new_ref(PyInt::from(BigInt::from(v)), types.int_type.clone(), None))
-                .collect();
-
-            let true_value = create_object(PyInt::from(1), &types.bool_type);
-            let false_value = create_object(PyInt::from(0), &types.bool_type);
-
-            let empty_tuple = create_object(
-                PyTuple::_new(Vec::new().into_boxed_slice()),
-                &types.tuple_type,
-            );
-
-            let tp_new_wrapper = create_object(
-                PyNativeFuncDef::from(pytype::tp_new_wrapper.into_func()).into_function(),
-                &types.builtin_function_or_method_type,
-            )
-            .into_object();
-
-            let context = PyContext {
-                true_value,
-                false_value,
-                not_implemented,
-                none,
-                empty_tuple,
-                ellipsis,
-
-                types,
-                exceptions,
-                int_cache_pool,
-                tp_new_wrapper,
-            };
-            TypeZoo::extend(&context);
-            exceptions::ExceptionZoo::extend(&context);
-            context
-        })
-        .clone()
+        CONTEXT.get_or_init(Self::init).clone()
     }
 
     pub fn none(&self) -> PyObjectRef {
@@ -1242,23 +1238,25 @@ pub trait StaticType {
         crate::builtins::object::PyBaseObject::static_type()
     }
     fn static_type() -> &'static PyTypeRef {
-        static_cell::get(Self::static_cell()).unwrap_or_else(|| unsafe {
-            // SAFETY: object must be initialized by init_* method.
-            // So this is actually not safe as itself.
-            // But easy to find out when it happened in debug build
-            std::hint::unreachable_unchecked()
-        })
+        Self::static_cell()
+            .get()
+            .expect("static type has not been initialized")
     }
-
     fn init_manually(typ: PyTypeRef) -> &'static PyTypeRef {
-        static_cell::init_expect(Self::static_cell(), typ, "init_manually")
+        let cell = Self::static_cell();
+        cell.set(typ)
+            .unwrap_or_else(|_| panic!("double initialization from init_manually"));
+        cell.get().unwrap()
     }
     fn init_bare_type() -> &'static PyTypeRef
     where
         Self: PyClassImpl,
     {
         let typ = Self::create_bare_type();
-        static_cell::init_expect(Self::static_cell(), typ, Self::NAME)
+        let cell = Self::static_cell();
+        cell.set(typ)
+            .unwrap_or_else(|_| panic!("double initialization of {}", Self::NAME));
+        cell.get().unwrap()
     }
     fn create_bare_type() -> PyTypeRef
     where
@@ -1313,12 +1311,13 @@ pub trait PyClassImpl: PyClassDef {
     where
         Self: StaticType,
     {
-        static_cell::get_or_init(Self::static_cell(), || {
-            let typ = Self::create_bare_type();
-            Self::extend_class(ctx, &typ);
-            typ
-        })
-        .clone()
+        Self::static_cell()
+            .get_or_init(|| {
+                let typ = Self::create_bare_type();
+                Self::extend_class(ctx, &typ);
+                typ
+            })
+            .clone()
     }
 
     fn extend_slots(slots: &mut PyTypeSlots);