generic_static_cache/

lib.rs

#![cfg_attr(docsrs, feature(doc_cfg))]
#![allow(named_asm_labels)]
#![cfg_attr(
    any(
        target_arch = "x86_64",
        target_arch = "aarch64",
        target_arch = "arm",
        target_arch = "x86"
    ),
    no_std
)]
// TODO: Properly test data for same type from different compilation units
// TODO: More platforms
// TODO: Benches

//! Quoting the [Rust Reference](https://doc.rust-lang.org/reference/items/static-items.html#statics--generics):
//!
//! > A static item defined in a generic scope (for example in a blanket or default implementation)
//! > will result in exactly one static item being defined, as if the static definition was pulled
//! > out of the current scope into the module. There will not be one item per monomorphization.
//!
//! One way to work around this is to use a `HashMap<TypeId,Data>`. This is a simple & usually the best solution.
//! If lookup performance is important, you can skip hashing the `TypeId` for minor gains as it
//! [already contains](https://github.com/rust-lang/rust/blob/eeff92ad32c2627876112ccfe812e19d38494087/library/core/src/any.rs#L645)
//! a good-quality hash.
//!
//! This crate aims to further fully remove the lookup by allocating the storage using inline
//! assembly.
//!
//! # ⚠ Caveats ⚠
//!
//! THIS PLACE IS NOT A PLACE OF HONOR.  
//! NO HIGHLY ESTEEMED DEED IS COMMEMORATED HERE.  
//! NOTHING VALUED IS HERE.  
//! WHAT IS HERE WAS DANGEROUS AND REPULSIVE TO US.  
//! THE DANGER IS IN A PARTICULAR LOCATION.  
//! THE DANGER IS STILL PRESENT, IN YOUR TIME, AS IT WAS IN OURS.  
//!
//! Different compilation units may access different instances of the data
//! (at least without `share-generics`).
//!
//! Supported targets are **x86-64**, **aarch64**, **arm** and **x86**;
//! on other targets, this crate falls back to a hashmap.
//!
//! The linker may not use a smaller align for the `data` section than that of
//! the types used with `global`.
//!
//! This crate isn't as well-tested as is should be.
//!
//! # Usage
//! The `generic_static` macro defines a static inside of a function,
//! with every instantiation of the function having it's own instance of the static:
//!
//! ```rust
//! # use core::sync::atomic::{AtomicI32, Ordering};
//! # use generic_static_cache::generic_static;
//! fn get_and_inc<T>() -> i32 {
//!     generic_static!{
//!         static COUNTER: &AtomicI32 = &AtomicI32::new(1);
//!     }
//!     COUNTER.fetch_add(1, Ordering::Relaxed)
//! }
//! assert_eq!(get_and_inc::<bool>(), 1);
//! assert_eq!(get_and_inc::<bool>(), 2);
//! assert_eq!(get_and_inc::<String>(), 1);
//! assert_eq!(get_and_inc::<bool>(), 3);
//! ```
//!
//! Underlying this is the `global::<T>()` function,
//! which allocates a shared global static for each type it's used with.

extern crate alloc;

pub use alloc::boxed::Box;
pub use bytemuck;

/// Access a global instance of `T`, zero-initialized at program start.
///
/// For types that cannot be zero-initialized, use [`non_zeroable_global`].
///
/// # Example
/// ```
/// # use core::sync::atomic::{AtomicI32, Ordering::Relaxed};
/// # use generic_static_cache::bytemuck;
/// struct MyType(AtomicI32);
/// unsafe impl bytemuck::Zeroable for MyType {}
///
/// assert_eq!(generic_static_cache::global::<MyType>().0.load(Relaxed), 0);
/// generic_static_cache::global::<MyType>().0.store(1, Relaxed);
/// assert_eq!(generic_static_cache::global::<MyType>().0.load(Relaxed), 1);
/// ```
pub fn global<T: bytemuck::Zeroable + Sync + 'static>() -> &'static T {
    assert!(core::mem::align_of::<T>() <= 4 * 1024);
    #[cfg(any(
        target_arch = "x86_64",
        target_arch = "aarch64",
        target_arch = "arm",
        target_arch = "x86"
    ))]
    {
        unsafe {
            // Tested both with position-independent code and with -C relocation-model=static

            // Reserve space in writable memory
            core::arch::asm!(
                ".pushsection .data",
                ".balign {align}",
                "global_static_{id}:",
                ".skip {size}",
                ".popsection",
                id = sym global::<T>,
                align = const core::mem::align_of::<T>(),
                size = const core::mem::size_of::<T>(),
                options(nomem)
            );
            let addr: usize;
            #[cfg(target_arch = "x86_64")]
            {
                core::arch::asm!(
                    "lea {addr}, [rip+global_static_{id}]",
                    addr = out(reg) addr,
                    id = sym global::<T>,
                    options(pure, nomem)
                );
            }
            #[cfg(target_arch = "aarch64")]
            {
                core::arch::asm!(
                    "adrp {addr}, global_static_{id}",
                    "add {addr}, {addr}, :lo12:global_static_{id}",
                    addr = out(reg) addr,
                    id = sym global::<T>,
                    options(pure, nomem)
                );
            }
            #[cfg(target_arch = "arm")]
            {
                core::arch::asm!(
                    "ldr {addr}, =global_static_{id}",
                    addr = out(reg) addr,
                    id = sym global::<T>,
                    options(pure, nomem)
                );
            }
            #[cfg(target_arch = "x86")]
            {
                core::arch::asm!(
                    // AT&T syntax as intel syntax seems to throw an error
                    // when using two symbols in the memory operand
                    "call 1f",
                    "1: pop {addr}",
                    "lea global_static_{id}-1b({addr}), {addr}",
                    addr = out(reg) addr,
                    id = sym global::<T>,
                    options(pure, nomem, att_syntax)
                );
            }
            &*(addr as *const _)
        }
    }
    #[cfg(not(any(
        target_arch = "x86_64",
        target_arch = "aarch64",
        target_arch = "arm",
        target_arch = "x86"
    )))]
    {
        pub(crate) struct SyncWrapper(*const u8);
        unsafe impl Sync for SyncWrapper {}
        unsafe impl Send for SyncWrapper {}
        use std::any::TypeId;
        static MAP: std::sync::RwLock<TypeIdMap<SyncWrapper>> =
            std::sync::RwLock::new(TypeIdMap::with_hasher(NoOpTypeIdBuildHasher));
        {
            let guard = MAP.read().unwrap();
            if let Some(value) = guard.get(&TypeId::of::<T>()) {
                return unsafe { &*(value.0 as *const T) };
            }
        }
        let mut guard = MAP.write().unwrap();
        let value = guard
            .entry(TypeId::of::<T>())
            .or_insert(SyncWrapper(
                Box::into_raw(Box::new(<T as bytemuck::Zeroable>::zeroed())) as *const u8,
            ));
        unsafe { &*(value.0 as *const T) }
    }
}

/// Access a generic global that can contain non-zeroable types.
/// This is done via an indirection to the heap, so use
/// [`global`] if this is not needed.
///
/// This doesn't interfere with data accessed with [`global`].
///
/// # Example
/// ```
/// # use generic_static_cache::non_zeroable_global;
/// #[derive(PartialEq, Debug)]
/// struct MyType;
///
/// assert_eq!(non_zeroable_global::get::<MyType>(), None);
/// let _ = non_zeroable_global::init(MyType);
/// assert_eq!(non_zeroable_global::get::<MyType>(), Some(&MyType));
// TODO: Add fallback
#[cfg(any(
    target_arch = "x86_64",
    target_arch = "aarch64",
    target_arch = "arm",
    target_arch = "x86"
))]
pub mod non_zeroable_global {
    use super::*;

    // TODO: Remove one layer of indirection on unsupported platforms

    /// Wrapper to prevent interfering with the user's `direct` calls
    struct Heap<T>(core::sync::atomic::AtomicPtr<T>);
    unsafe impl<T> bytemuck::Zeroable for Heap<T> {}
    unsafe impl<T: Sync> Sync for Heap<T> {}

    #[derive(Debug)]
    pub struct AlreadyInitialized;

    /// Initialize the global initiable storage of type `Type`.
    ///
    /// If called multiple times, only the first call will succeed.
    ///
    /// [`global`]: crate::global
    pub fn init<T: Sync + 'static>(data: T) -> Result<(), AlreadyInitialized> {
        use core::sync::atomic::Ordering;
        let boxed = Box::into_raw(Box::new(data));
        match global::<Heap<T>>().0.compare_exchange(
            core::ptr::null_mut(),
            boxed,
            Ordering::SeqCst,
            Ordering::SeqCst,
        ) {
            Ok(_) => Ok(()),
            Err(_) => {
                unsafe {
                    drop(Box::from_raw(boxed));
                }
                Err(AlreadyInitialized)
            }
        }
    }

    /// Access the global initiable storage of type `Type`.
    ///
    /// [`global`]: crate::global
    pub fn get<T: Sync + 'static>() -> Option<&'static T> {
        use core::sync::atomic::Ordering;
        let data = global::<Heap<T>>().0.load(Ordering::SeqCst);
        if data.is_null() {
            None
        } else {
            Some(unsafe { &*data })
        }
    }

    /// Initialize & access the global initiable storage of type `Type`.
    ///
    /// If this is called multiple times simultaneously,
    /// the `cons` argument of multiple invocations may be called,
    /// but only one result will be used.
    ///
    /// [`global`]: crate::global
    pub fn get_or_init<T: Sync + 'static>(cons: impl Fn() -> T) -> &'static T {
        use core::sync::atomic::Ordering;
        let data = global::<Heap<T>>().0.load(Ordering::SeqCst);
        if data.is_null() {
            let _ = init::<_>(cons());
            get::<_>().unwrap()
        } else {
            unsafe { &*data }
        }
    }
}

/// Declare a static variable that is not shared across different monomorphizations
/// of the containing functions. Its type must be a shared reference to a [`Sync`] type.
///
/// If this is executed for the first time in multiple threads simultaneously,
/// the initializing expression may get executed multiple times.
///
/// # Example
/// ```
/// # use std::sync::Mutex;
/// # use generic_static_cache::generic_static;
/// fn generic_function<T>() {
///     generic_static!{
///         static NAME: &Mutex<String> = &Mutex::new("Ferris".to_string());
///     }
/// # /*
///     …
/// # */
/// }
/// ```
#[cfg(any(
    target_arch = "x86_64",
    target_arch = "aarch64",
    target_arch = "arm",
    target_arch = "x86"
))]
#[macro_export]
macro_rules! generic_static {
    {static $ident:ident $(: &$type:ty)? = &$init:expr;} => {
        #[allow(non_snake_case)]
        let $ident $(: &'static $type)? = {
            let init = ||$init;
            fn assert_sync<T: Sync>(_: &impl FnOnce() -> T) {}
            assert_sync(&init);

            // Use empty closure to create a new type to use as a unique key,
            // use reference to initializer to infer type of static data
            fn make<Key: 'static, Value: Sync + 'static>(_: Key, _: &impl FnOnce()->Value)
            -> &'static ::core::sync::atomic::AtomicPtr<Value> {
                struct Holder<T, D> {
                    _marker: core::marker::PhantomData<T>,
                    value: core::sync::atomic::AtomicPtr<D>
                }
                unsafe impl<T, D> $crate::bytemuck::Zeroable for Holder<T,D>{}
                unsafe impl<T, D> Sync for Holder<T,D>{}
                &$crate::global::<Holder<Key, Value>>().value
            }
            let ptr = make(||(), &init);

            let data = ptr.load(::core::sync::atomic::Ordering::SeqCst);
            if data.is_null() {
                // Need to call initializer
                // This can be called multiple times if executed for the first time
                // in multiple threads simultaneously!
                let boxed = $crate::Box::into_raw($crate::Box::new(init()));
                if ptr
                    .compare_exchange(
                        ::core::ptr::null_mut(),
                        boxed,
                        ::core::sync::atomic::Ordering::SeqCst,
                        ::core::sync::atomic::Ordering::SeqCst,
                    )
                    .is_err()
                {
                    // Was simultaneously initialized by another thread
                    unsafe {
                        drop($crate::Box::from_raw(boxed));
                    }
                }
                unsafe { &*boxed }
            } else {
                unsafe { &*data }
            }
        };
    };
}

/// Declare a static variable that is not shared across different monomorphizations
/// of the containing functions. Its type must be a shared reference to a [`Sync`] type.
///
/// If this is executed for the first time in multiple threads simultaneously,
/// the initializing expression may get executed multiple times.
///
/// # Example
/// ```
/// # use std::sync::Mutex;
/// # use generic_static_cache::generic_static;
/// fn generic_function<T>() {
///     generic_static!{
///         static NAME: &Mutex<String> = &Mutex::new("Ferris".to_string());
///     }
/// # /*
///     …
/// # */
/// }
/// ```
#[cfg(not(any(
    target_arch = "x86_64",
    target_arch = "aarch64",
    target_arch = "arm",
    target_arch = "x86"
)))]
#[macro_export]
macro_rules! generic_static {
    // TODO: Use a RwLock<*mut u8> to remove the restrictions on type inference
    // and usage of outer types.
    {static $ident:ident $(: &$type:ty)? = &$init:expr;} => {
        #[allow(non_snake_case)]
        let $ident $(: &'static $type)? = 'block: {
            struct SyncWrapper(*const u8);
            unsafe impl Sync for SyncWrapper {}
            unsafe impl Send for SyncWrapper {}
            fn id<T: 'static>(_: T) -> core::any::TypeId {
                core::any::TypeId::of::<T>()
            }
            // Each closure expression defines a closure type [expr.closure.intro].
            // As closures can depend on generic types from a surrounding items,
            // closures in different instantiations of outer generics must have different types
            // and therefore must be different closure expressions.
            // The types defined by closure expressions are unique [expr.closure.unique-type].
            let id = id(||());
            static MAP: ::std::sync::RwLock<$crate::TypeIdMap<SyncWrapper>> =
                ::std::sync::RwLock::new($crate::TypeIdMap::with_hasher(
                    $crate::NoOpTypeIdBuildHasher,
                ));
            {
                let guard = MAP.read().unwrap();
                if let Some(value) = guard.get(&id) {
                    break 'block unsafe { &*(value.0 as *const _) };
                }
            }
            let mut guard = MAP.write().unwrap();
            let value = guard
                .entry(id)
                .or_insert(SyncWrapper(
                    Box::into_raw(Box::new($init)) as *const u8,
                ));
            unsafe { &*(value.0 as *const _) }
        };

    };
}

#[cfg(not(any(
    target_arch = "x86_64",
    target_arch = "aarch64",
    target_arch = "arm",
    target_arch = "x86"
)))]
pub use fallback::*;

#[cfg(not(any(
    target_arch = "x86_64",
    target_arch = "aarch64",
    target_arch = "arm",
    target_arch = "x86"
)))]
mod fallback {
    use core::any::TypeId;
    use core::hash::{BuildHasher, Hasher};
    use std::collections::HashMap;

    /// Fast type map as a fallback.
    pub type TypeIdMap<T> = HashMap<TypeId, T, NoOpTypeIdBuildHasher>;

    #[doc(hidden)]
    #[derive(Default)]
    pub struct NoOpTypeIdBuildHasher;

    impl BuildHasher for NoOpTypeIdBuildHasher {
        type Hasher = NoOpTypeIdHasher;

        fn build_hasher(&self) -> Self::Hasher {
            NoOpTypeIdHasher(0)
        }
    }

    #[doc(hidden)]
    #[derive(Default)]
    pub struct NoOpTypeIdHasher(u64);

    impl Hasher for NoOpTypeIdHasher {
        fn finish(&self) -> u64 {
            self.0
        }

        fn write(&mut self, _bytes: &[u8]) {
            unimplemented!()
        }

        fn write_u64(&mut self, i: u64) {
            self.0 = i
        }
    }
}

#[cfg(test)]
mod test {
    #[test]
    fn test_global() {
        use crate::global;
        use core::sync::atomic::{AtomicI32, AtomicI64, Ordering};

        let a = global::<AtomicI32>();
        let b = global::<AtomicI64>();
        assert_eq!(a.load(Ordering::Relaxed), 0);
        a.store(69, Ordering::Relaxed);
        assert_eq!(a.load(Ordering::Relaxed), 69);
        assert_eq!(b.load(Ordering::Relaxed), 0);
        assert_eq!(*global::<i64>(), 0);

        core::hint::black_box(global::<AtomicI64>());
    }

    #[test]
    fn test_macro() {
        use core::sync::atomic::{AtomicI32, Ordering};
        #[allow(clippy::extra_unused_type_parameters)]
        fn get_and_inc<T: 'static>() -> i32 {
            generic_static!(
                static BLUB: &AtomicI32 = &AtomicI32::new(1);
            );
            let value = BLUB.load(Ordering::Relaxed);
            BLUB.fetch_add(1, Ordering::Relaxed);
            value
        }
        assert_eq!(get_and_inc::<bool>(), 1);
        assert_eq!(get_and_inc::<bool>(), 2);
        assert_eq!(get_and_inc::<i32>(), 1);
        assert_eq!(get_and_inc::<bool>(), 3);

        generic_static!(
            static FOO_1: &AtomicI32 = &AtomicI32::new(0);
        );
        generic_static!(
            static FOO_2: &AtomicI32 = &AtomicI32::new(69);
        );
        assert_eq!(FOO_1.load(Ordering::Relaxed), 0);
        assert_eq!(FOO_2.load(Ordering::Relaxed), 69);
        FOO_1.store(1, Ordering::Relaxed);
        FOO_2.store(2, Ordering::Relaxed);
        assert_eq!(FOO_1.load(Ordering::Relaxed), 1);
        assert_eq!(FOO_2.load(Ordering::Relaxed), 2);
    }

    #[cfg(any(
        target_arch = "x86_64",
        target_arch = "aarch64",
        target_arch = "arm",
        target_arch = "x86"
    ))]
    #[test]
    fn test_macro_type_inference() {
        generic_static! {
            static _FOO = &();
        }
    }
}