ax_percpu/
value.rs

1use core::{marker::PhantomData, ptr::NonNull};
2
3use cpu_local::{CpuPin, ExclusiveCpu};
4
5use crate::PerCpuArea;
6
7/// Provider generated for one concrete symbol in the per-CPU template.
8///
9/// # Safety
10///
11/// Every returned pointer must address the declared `T` in the selected live
12/// area. Primitive providers must use the matching atomic representation.
13#[doc(hidden)]
14pub unsafe trait PerCpuSymbol<T> {
15    fn offset() -> usize;
16    fn current_ptr(pin: &CpuPin<'_>) -> NonNull<T>;
17    fn remote_ptr(area: PerCpuArea) -> NonNull<T>;
18}
19
20/// Marker implemented only by macro-generated object symbols.
21///
22/// # Safety
23///
24/// The provider's storage must contain a live `T` in every initialized area.
25#[doc(hidden)]
26pub unsafe trait PerCpuObjectSymbol<T>: PerCpuSymbol<T> {}
27
28/// Marker implemented only by macro-generated atomic scalar symbols.
29///
30/// # Safety
31///
32/// The provider's storage must use the atomic representation paired with `T`.
33#[doc(hidden)]
34pub unsafe trait PerCpuPrimitiveSymbol<T>: PerCpuSymbol<T> {}
35
36type PerCpuMarker<T, S> = fn() -> (T, S);
37
38/// Typed descriptor for one symbol replicated in every runtime CPU area.
39pub struct PerCpu<T, S> {
40    _marker: PhantomData<PerCpuMarker<T, S>>,
41}
42
43impl<T, S> PerCpu<T, S>
44where
45    S: PerCpuSymbol<T>,
46{
47    /// Creates the zero-sized descriptor for a macro-generated symbol.
48    #[doc(hidden)]
49    pub const fn new() -> Self {
50        Self {
51            _marker: PhantomData,
52        }
53    }
54
55    /// Returns this symbol's byte offset in one area.
56    pub fn offset(&self) -> usize {
57        S::offset()
58    }
59
60    /// Returns a typed pointer whose address is stable for `pin`.
61    pub fn current_ptr(&self, pin: &CpuPin<'_>) -> NonNull<T> {
62        S::current_ptr(pin)
63    }
64
65    /// Returns a typed pointer in an explicitly selected remote area.
66    ///
67    /// The caller remains responsible for synchronization before dereference.
68    pub fn remote_ptr(&self, area: PerCpuArea) -> NonNull<T> {
69        S::remote_ptr(area)
70    }
71}
72
73impl<T, S> PerCpu<T, S>
74where
75    S: PerCpuObjectSymbol<T>,
76{
77    /// Mutates the current CPU's object without allowing its borrow to escape.
78    pub fn with_current_mut<R>(
79        &self,
80        exclusive: &ExclusiveCpu<'_>,
81        operation: impl for<'value> FnOnce(&'value mut T) -> R,
82    ) -> R {
83        // SAFETY: ExclusiveCpu proves local and remote alias exclusion for the
84        // closure, while its area fixes the generated address.
85        let mut pointer =
86            unsafe { NonNull::new_unchecked((exclusive.area().base() + S::offset()) as *mut T) };
87        operation(unsafe { pointer.as_mut() })
88    }
89}
90
91impl<T, S> PerCpu<T, S>
92where
93    T: Sync,
94    S: PerCpuObjectSymbol<T>,
95{
96    /// Borrows a shared current-CPU object for one non-escaping callback.
97    pub fn with_current<R>(
98        &self,
99        pin: &CpuPin<'_>,
100        operation: impl for<'value> FnOnce(&'value T) -> R,
101    ) -> R {
102        // SAFETY: T: Sync permits shared observation and the pin fixes address.
103        operation(unsafe { S::current_ptr(pin).as_ref() })
104    }
105}
106
107mod primitive {
108    use core::{
109        ptr::NonNull,
110        sync::atomic::{
111            AtomicBool, AtomicU8, AtomicU16, AtomicU32, AtomicU64, AtomicUsize, Ordering,
112        },
113    };
114
115    pub trait Sealed: Copy {
116        unsafe fn load(pointer: NonNull<Self>) -> Self;
117        unsafe fn store(pointer: NonNull<Self>, value: Self);
118    }
119
120    macro_rules! impl_atomic_primitive {
121        ($value:ty, $atomic:ty) => {
122            impl Sealed for $value {
123                unsafe fn load(pointer: NonNull<Self>) -> Self {
124                    unsafe { pointer.cast::<$atomic>().as_ref() }.load(Ordering::Relaxed)
125                }
126
127                unsafe fn store(pointer: NonNull<Self>, value: Self) {
128                    unsafe { pointer.cast::<$atomic>().as_ref() }.store(value, Ordering::Relaxed);
129                }
130            }
131        };
132    }
133
134    impl_atomic_primitive!(bool, AtomicBool);
135    impl_atomic_primitive!(u8, AtomicU8);
136    impl_atomic_primitive!(u16, AtomicU16);
137    impl_atomic_primitive!(u32, AtomicU32);
138    impl_atomic_primitive!(u64, AtomicU64);
139    impl_atomic_primitive!(usize, AtomicUsize);
140}
141
142impl<T, S> PerCpu<T, S>
143where
144    T: primitive::Sealed,
145    S: PerCpuPrimitiveSymbol<T>,
146{
147    /// Loads the current CPU's atomic scalar with relaxed ordering.
148    pub fn read_current(&self, pin: &CpuPin<'_>) -> T {
149        unsafe { T::load(S::current_ptr(pin)) }
150    }
151
152    /// Stores the current CPU's atomic scalar with relaxed ordering.
153    pub fn write_current(&self, pin: &CpuPin<'_>, value: T) {
154        unsafe { T::store(S::current_ptr(pin), value) }
155    }
156}
ax_percpu/value.rs

ax_percpu/
value.rs
