generic-static-cache 0.2.1

#![cfg_attr(docsrs, feature(doc_cfg))]
#![allow(named_asm_labels)]
#![cfg_attr(not(feature = "std"), no_std)]

//! # generic_static_cache
//!
//! 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 speed up the lookup by allocating the storage using inline
//! assembly: Accessing a generic static provided by this crate is instant,
//! whereas using a hashmap takes more than 10×instant.
//!
//! ## ⚠ Caveats ⚠
//!
//! This crate isn't as well-tested as is should be.
//!
//! Supported targets are **x86-64**, **aarch64**, **arm** and **x86**;
//! on other targets, this crate falls back to a hashmap.
//!
//! ## no_std
//! On supported platforms, `global` is always available.
//!
//! With the `alloc` feature, `generic_static!` and `non_zeroable_global` become available.
//!
//! With the `std` feature, everything also becomes available on unsupported platforms.

pub use bytemuck;

/// Access a zero-initialized global static instance of `T`.
///
/// 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);
/// assert_eq!(*generic_static_cache::global::<i32>(), 0);
/// ```
#[cfg(any(
    feature = "std",
    target_arch = "x86_64",
    target_arch = "aarch64",
    target_arch = "arm",
    target_arch = "x86"
))]
#[inline(always)]
pub fn global<T: bytemuck::Zeroable + Sync + 'static>() -> &'static T {
    global_impl::<false, T>()
}

/// Like [`global`], but the symbol identifying the storage will not be exported.
///
/// This means different binaries linked together may not share the same instance
/// of this global variable, but the binary may be smaller (if many such globals
/// are exposed).
///
/// Globals defined with `global` and `local_global` are separate from each other.
#[cfg(any(
    feature = "std",
    target_arch = "x86_64",
    target_arch = "aarch64",
    target_arch = "arm",
    target_arch = "x86"
))]
#[inline(always)]
pub fn local_global<T: bytemuck::Zeroable + Sync + 'static>() -> &'static T {
    global_impl::<true, T>()
}

#[cfg(any(
    feature = "std",
    target_arch = "x86_64",
    target_arch = "aarch64",
    target_arch = "arm",
    target_arch = "x86"
))]
#[inline(always)]
fn global_impl<const LOCAL: bool, 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 zeroed-initialized writable memory (bss).
            // rustc may duplicate the content of this function,
            // so using .ifndef to deduplicate them within the binary.
            core::arch::asm!(
                ".ifnotdef global_{local}_{id}",
                ".if {local}",
                ".local global_{local}_{id}",
                ".endif",
                ".comm global_{local}_{id}, {size}, {align}",
                ".endif",
                local = const if LOCAL {1} else {0},
                id = sym global::<T>,
                size = const core::mem::size_of::<T>(),
                align = const core::mem::align_of::<T>(),
                options(nomem)
            );
            // Now load the symbol's address
            let addr: usize;
            #[cfg(target_arch = "x86_64")]
            {
                core::arch::asm!(
                    "lea {addr}, [rip+global_{local}_{id}]",
                    addr = out(reg) addr,
                    local = const if LOCAL {1} else {0},
                    id = sym global::<T>,
                    options(pure, nomem)
                );
            }
            #[cfg(target_arch = "aarch64")]
            {
                core::arch::asm!(
                    "adrp {addr}, global_{local}_{id}",
                    "add {addr}, {addr}, :lo12:global_{local}_{id}",
                    addr = out(reg) addr,
                    local = const if LOCAL {1} else {0},
                    id = sym global::<T>,
                    options(pure, nomem)
                );
            }
            #[cfg(target_arch = "arm")]
            {
                core::arch::asm!(
                    "ldr {addr}, =global_{local}_{id}",
                    addr = out(reg) addr,
                    local = const if LOCAL {1} else {0},
                    id = sym global::<T>,
                    options(pure, nomem)
                );
            }
            #[cfg(target_arch = "x86")]
            {
                core::arch::asm!(
                    // AT&T syntax allows using multiple symbols in the memory operand
                    "call 2f",
                    "2: pop {addr}",
                    "lea global_{local}_{id}-2b({addr}), {addr}",
                    addr = out(reg) addr,
                    local = const if LOCAL {1} else {0},
                    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 core::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(
                    alloc::boxed::Box::into_raw(alloc::boxed::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 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));
/// ```
#[cfg(all(
    feature = "alloc",
    any(
        feature = "std",
        target_arch = "x86_64",
        target_arch = "aarch64",
        target_arch = "arm",
        target_arch = "x86"
    )
))]
pub mod non_zeroable_global {
    extern crate alloc;

    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 = alloc::boxed::Box::into_raw(alloc::boxed::Box::new(data));
        match global::<Heap<T>>().0.compare_exchange(
            core::ptr::null_mut(),
            boxed,
            Ordering::SeqCst,
            Ordering::SeqCst,
        ) {
            Ok(_) => Ok(()),
            Err(_) => {
                unsafe {
                    drop(alloc::boxed::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`]`+'static` type,
/// and the initializer expression must start with a `&`.
/// Outer type variables may be used and the type hint is optional.
///
/// The initializing expression doesn't need to be `const`.
/// If this is executed for the first time in multiple threads simultaneously,
/// the initializing expression may get executed multiple times.
///
/// # Example
/// ```
/// # use core::sync::atomic::{AtomicU32, Ordering::Relaxed};
/// # use generic_static_cache::generic_static;
/// fn numeric_type_id<T>() -> u32 {
///     static NEXT: AtomicU32 = AtomicU32::new(0);
///     generic_static!{
///         static ID: &u32 = &NEXT.fetch_add(1, Relaxed);
///     }
///     *ID
/// }
/// assert_eq!(numeric_type_id::<bool>(), 0);
/// assert_eq!(numeric_type_id::<String>(), 1);
/// assert_eq!(numeric_type_id::<i32>(), 2);
/// assert_eq!(numeric_type_id::<bool>(), 0);
/// ```
#[cfg(all(
    feature = "alloc",
    any(
        feature = "std",
        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)? = {
            #[cfg(any(
                target_arch = "x86_64",
                target_arch = "aarch64",
                target_arch = "arm",
                target_arch = "x86"
            ))] {
                extern crate alloc;

                let init = ||$init;
                fn assert_sync_static<T: Sync + 'static>(_: &impl FnOnce() -> T) {}
                assert_sync_static(&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 = alloc::boxed::Box::into_raw(alloc::boxed::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(alloc::boxed::Box::from_raw(boxed));
                        }
                    }
                    unsafe { &*boxed }
                } else {
                    unsafe { &*data }
                }
            }
            #[cfg(not(any(
                target_arch = "x86_64",
                target_arch = "aarch64",
                target_arch = "arm",
                target_arch = "x86"
            )))] {
                #[cfg(not(feature = "std"))]
                compile_error!("Unsupported platform, enable feature \"std\" to enable fallback");

                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(
                        alloc::boxed::Box::into_raw(alloc::boxed::Box::new($init)) as *const u8,
                    ));
                unsafe { &*(value.0 as *const _) }
            }
        };
    };
}

#[cfg(feature = "std")]
pub use with_std::*;

#[cfg(feature = "std")]
mod with_std {
    use core::any::TypeId;
    use core::hash::{BuildHasher, Hasher};
    use std::collections::HashMap;

    /// Fast map keyed by `TypeId`.
    pub type TypeIdMap<T> = HashMap<TypeId, T, NoOpTypeIdBuildHasher>;

    /// Hasher for [`TypeIdMap`].
    #[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]) {
            // Slow, bad quality fallback to not break applications in case std implementation changes.
            self.0 = bytes.iter().fold(self.0, |hash, b| {
                hash.rotate_left(8).wrapping_add(*b as u64)
            });
        }

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

#[cfg(test)]
mod test {
    use std::{
        any::TypeId,
        hash::{Hash, Hasher},
    };

    #[test]
    fn test_local_global() {
        use crate::local_global;
        use core::sync::atomic::{AtomicI32, AtomicI64, Ordering};

        let a = local_global::<AtomicI32>();
        let b = local_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!(*local_global::<i64>(), 0);

        core::hint::black_box(local_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);
    }

    #[test]
    fn test_macro_types() {
        fn generic<T: Sync + 'static>(t: T) {
            generic_static! {
                static _FOO = &t;
            }
        }
        generic(0);
        generic(true);
    }

    #[test]
    fn type_id_hash() {
        TypeId::of::<()>().hash(&mut {
            struct H;
            impl Hasher for H {
                fn finish(&self) -> u64 {
                    0
                }

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

                fn write_u64(&mut self, _: u64) {}
            }
            H
        });
    }
}