percpu_macros/lib.rs
1//! Macros to define and access a per-CPU data structure.
2//!
3//! **DO NOT** use this crate directly. Use the [percpu] crate instead.
4//!
5//! [percpu]: https://docs.rs/percpu
6//!
7//! ## Implementation details of the `def_percpu` macro
8//!
9//! ### Core idea
10//!
11//! The core idea is to collect all per-CPU static variables to a single section (i.e., `.percpu`), then allocate a
12//! per-CPU data area, with the size equals to the size of the `.percpu` section, for each CPU (it can be done
13//! statically or dynamically), then copy the `.percpu` section to each per-CPU data area during initialization.
14//!
15//! The address of a per-CPU static variable on a given CPU can be calculated by adding the offset of the variable
16//! (relative to the section base) to the base address of the per-CPU data area on the CPU.
17//!
18//! ### How to access the per-CPU data
19//!
20//! To access a per-CPU static variable on a given CPU, three values are needed:
21//!
22//! - The base address of the per-CPU data area on the CPU,
23//! - which can be calculated by the base address of the whole per-CPU data area and the CPU ID,
24//! - and then stored in a register, like `TPIDR_EL1`/`TPIDR_EL2` on AArch64, or `gs` on x86_64.
25//! - The offset of the per-CPU static variable relative to the per-CPU data area base,
26//! - which can be calculated by assembly notations, like `offset symbol` on x86_64, or `#:abs_g0_nc:symbol` on
27//! AArch64, or `%hi(symbol)` and `%lo(symbol)` on RISC-V.
28//! - The size of the per-CPU static variable,
29//! - which we actually do not need to know, just give the right type to rust compiler.
30//!
31//! ### Generated code
32//!
33//! For each static variable `X` with type `T` that is defined with the `def_percpu` macro, the following items are
34//! generated:
35//!
36//! - A static variable `__PERCPU_X` with type `T` that stores the per-CPU data.
37//!
38//! This variable is placed in the `.percpu` section. All attributes of the original static variable, as well as the
39//! initialization expression, are preserved.
40//!
41//! This variable is never, and should never be, accessed directly. To access the per-CPU data, the offset of the
42//! variable is, and should be, used.
43//!
44//! - A zero-sized wrapper struct `X_WRAPPER` that is used to access the per-CPU data.
45//!
46//! Some methods are generated in this struct to access the per-CPU data. For primitive integer types, extra methods
47//! are generated to accelerate the access.
48//!
49//! - A static variable `X` of type `X_WRAPPER` that is used to access the per-CPU data.
50//!
51//! This variable is always generated with the same visibility and attributes as the original static variable.
52//!
53
54use proc_macro::TokenStream;
55use proc_macro2::Span;
56use quote::{format_ident, quote};
57use syn::{Error, ItemStatic};
58
59#[cfg_attr(feature = "sp-naive", path = "naive.rs")]
60mod arch;
61
62fn compiler_error(err: Error) -> TokenStream {
63 err.to_compile_error().into()
64}
65
66/// Defines a per-CPU static variable.
67///
68/// It should be used on a `static` variable definition.
69///
70/// See the documentation of the [percpu](https://docs.rs/percpu) crate for more details.
71#[proc_macro_attribute]
72pub fn def_percpu(attr: TokenStream, item: TokenStream) -> TokenStream {
73 if !attr.is_empty() {
74 return compiler_error(Error::new(
75 Span::call_site(),
76 "expect an empty attribute: `#[def_percpu]`",
77 ));
78 }
79
80 let ast = syn::parse_macro_input!(item as ItemStatic);
81
82 let attrs = &ast.attrs;
83 let vis = &ast.vis;
84 let name = &ast.ident;
85 let ty = &ast.ty;
86 let init_expr = &ast.expr;
87
88 let inner_symbol_name = &format_ident!("__PERCPU_{}", name);
89 let struct_name = &format_ident!("{}_WRAPPER", name);
90
91 let ty_str = quote!(#ty).to_string();
92 let is_primitive_int = ["bool", "u8", "u16", "u32", "u64", "usize"].contains(&ty_str.as_str());
93
94 let no_preempt_guard = if cfg!(feature = "preempt") {
95 quote! { let _guard = percpu::__priv::NoPreemptGuard::new(); }
96 } else {
97 quote! {}
98 };
99
100 // Do not generate `fn read_current()`, `fn write_current()`, etc for non primitive types.
101 let read_write_methods = if is_primitive_int {
102 let read_current_raw = arch::gen_read_current_raw(inner_symbol_name, ty);
103 let write_current_raw =
104 arch::gen_write_current_raw(inner_symbol_name, &format_ident!("val"), ty);
105
106 quote! {
107 /// Returns the value of the per-CPU static variable on the current CPU.
108 ///
109 /// # Safety
110 ///
111 /// Caller must ensure that preemption is disabled on the current CPU.
112 #[inline]
113 pub unsafe fn read_current_raw(&self) -> #ty {
114 #read_current_raw
115 }
116
117 /// Set the value of the per-CPU static variable on the current CPU.
118 ///
119 /// # Safety
120 ///
121 /// Caller must ensure that preemption is disabled on the current CPU.
122 #[inline]
123 pub unsafe fn write_current_raw(&self, val: #ty) {
124 #write_current_raw
125 }
126
127 /// Returns the value of the per-CPU static variable on the current CPU. Preemption will be disabled during
128 /// the call.
129 pub fn read_current(&self) -> #ty {
130 #no_preempt_guard
131 unsafe { self.read_current_raw() }
132 }
133
134 /// Set the value of the per-CPU static variable on the current CPU. Preemption will be disabled during the
135 /// call.
136 pub fn write_current(&self, val: #ty) {
137 #no_preempt_guard
138 unsafe { self.write_current_raw(val) }
139 }
140 }
141
142 // Todo: maybe add `(read|write)_remote(_raw)?` here?
143 } else {
144 quote! {}
145 };
146
147 let symbol_vma = arch::gen_symbol_vma(inner_symbol_name);
148 let offset = arch::gen_offset(inner_symbol_name);
149 let current_ptr = arch::gen_current_ptr(inner_symbol_name, ty);
150 quote! {
151 #[cfg_attr(not(target_os = "macos"), link_section = ".percpu")] // unimplemented on macos
152 #(#attrs)*
153 static mut #inner_symbol_name: #ty = #init_expr;
154
155 #[doc = concat!("Wrapper struct for the per-CPU data [`", stringify!(#name), "`]")]
156 #[allow(non_camel_case_types)]
157 #vis struct #struct_name {}
158
159 #(#attrs)*
160 #vis static #name: #struct_name = #struct_name {};
161
162 impl #struct_name {
163 /// Returns the virtual memory address of this per-CPU static
164 /// variable in the `.percpu` section.
165 ///
166 /// It's same across all CPUs, and also the same as `offset` if the
167 /// "non-zero-vma" feature is not enabled.
168 #[inline]
169 pub fn symbol_vma(&self) -> usize {
170 #symbol_vma
171 }
172
173 /// Returns the offset relative to the per-CPU data area base.
174 #[inline]
175 pub fn offset(&self) -> usize {
176 #offset
177 }
178
179 /// Returns the raw pointer of this per-CPU static variable on the current CPU.
180 ///
181 /// # Safety
182 ///
183 /// Caller must ensure that preemption is disabled on the current CPU.
184 #[inline]
185 pub unsafe fn current_ptr(&self) -> *const #ty {
186 #current_ptr
187 }
188
189 /// Returns the reference of the per-CPU static variable on the current CPU.
190 ///
191 /// # Safety
192 ///
193 /// Caller must ensure that preemption is disabled on the current CPU.
194 #[inline]
195 pub unsafe fn current_ref_raw(&self) -> &#ty {
196 &*self.current_ptr()
197 }
198
199 /// Returns the mutable reference of the per-CPU static variable on the current CPU.
200 ///
201 /// # Safety
202 ///
203 /// Caller must ensure that preemption is disabled on the current CPU.
204 #[inline]
205 #[allow(clippy::mut_from_ref)]
206 pub unsafe fn current_ref_mut_raw(&self) -> &mut #ty {
207 &mut *(self.current_ptr() as *mut #ty)
208 }
209
210 /// Manipulate the per-CPU data on the current CPU in the given closure.
211 /// Preemption will be disabled during the call.
212 pub fn with_current<F, T>(&self, f: F) -> T
213 where
214 F: FnOnce(&mut #ty) -> T,
215 {
216 #no_preempt_guard
217 f(unsafe { self.current_ref_mut_raw() })
218 }
219
220 /// Returns the raw pointer of this per-CPU static variable on the given CPU.
221 ///
222 /// # Safety
223 ///
224 /// Caller must ensure that
225 /// - the CPU ID is valid, and
226 /// - data races will not happen.
227 #[inline]
228 pub unsafe fn remote_ptr(&self, cpu_id: usize) -> *const #ty {
229 let base = percpu::percpu_area_base(cpu_id);
230 let offset = #offset;
231 (base + offset) as *const #ty
232 }
233
234 /// Returns the reference of the per-CPU static variable on the given CPU.
235 ///
236 /// # Safety
237 ///
238 /// Caller must ensure that
239 /// - the CPU ID is valid, and
240 /// - data races will not happen.
241 #[inline]
242 pub unsafe fn remote_ref_raw(&self, cpu_id: usize) -> &#ty {
243 &*self.remote_ptr(cpu_id)
244 }
245
246 /// Returns the mutable reference of the per-CPU static variable on the given CPU.
247 ///
248 /// # Safety
249 ///
250 /// Caller must ensure that
251 /// - the CPU ID is valid, and
252 /// - data races will not happen.
253 #[inline]
254 #[allow(clippy::mut_from_ref)]
255 pub unsafe fn remote_ref_mut_raw(&self, cpu_id: usize) -> &mut #ty {
256 &mut *(self.remote_ptr(cpu_id) as *mut #ty)
257 }
258
259 #read_write_methods
260 }
261 }
262 .into()
263}
264
265#[doc(hidden)]
266#[cfg(not(feature = "sp-naive"))]
267#[proc_macro]
268pub fn percpu_symbol_vma(item: TokenStream) -> TokenStream {
269 let symbol = &format_ident!("{}", item.to_string());
270 let offset = arch::gen_symbol_vma(symbol);
271 quote!({ #offset }).into()
272}