xref: /linux/rust/kernel/cpumask.rs (revision 0074281bb6316108e0cff094bd4db78ab3eee236)
1 // SPDX-License-Identifier: GPL-2.0
2 
3 //! CPU Mask abstractions.
4 //!
5 //! C header: [`include/linux/cpumask.h`](srctree/include/linux/cpumask.h)
6 
7 use crate::{
8     alloc::{AllocError, Flags},
9     cpu::CpuId,
10     prelude::*,
11     types::Opaque,
12 };
13 
14 #[cfg(CONFIG_CPUMASK_OFFSTACK)]
15 use core::ptr::{self, NonNull};
16 
17 use core::ops::{Deref, DerefMut};
18 
19 /// A CPU Mask.
20 ///
21 /// Rust abstraction for the C `struct cpumask`.
22 ///
23 /// # Invariants
24 ///
25 /// A [`Cpumask`] instance always corresponds to a valid C `struct cpumask`.
26 ///
27 /// The callers must ensure that the `struct cpumask` is valid for access and
28 /// remains valid for the lifetime of the returned reference.
29 ///
30 /// # Examples
31 ///
32 /// The following example demonstrates how to update a [`Cpumask`].
33 ///
34 /// ```
35 /// use kernel::bindings;
36 /// use kernel::cpu::CpuId;
37 /// use kernel::cpumask::Cpumask;
38 ///
39 /// fn set_clear_cpu(ptr: *mut bindings::cpumask, set_cpu: CpuId, clear_cpu: CpuId) {
40 ///     // SAFETY: The `ptr` is valid for writing and remains valid for the lifetime of the
41 ///     // returned reference.
42 ///     let mask = unsafe { Cpumask::as_mut_ref(ptr) };
43 ///
44 ///     mask.set(set_cpu);
45 ///     mask.clear(clear_cpu);
46 /// }
47 /// ```
48 #[repr(transparent)]
49 pub struct Cpumask(Opaque<bindings::cpumask>);
50 
51 impl Cpumask {
52     /// Creates a mutable reference to an existing `struct cpumask` pointer.
53     ///
54     /// # Safety
55     ///
56     /// The caller must ensure that `ptr` is valid for writing and remains valid for the lifetime
57     /// of the returned reference.
as_mut_ref<'a>(ptr: *mut bindings::cpumask) -> &'a mut Self58     pub unsafe fn as_mut_ref<'a>(ptr: *mut bindings::cpumask) -> &'a mut Self {
59         // SAFETY: Guaranteed by the safety requirements of the function.
60         //
61         // INVARIANT: The caller ensures that `ptr` is valid for writing and remains valid for the
62         // lifetime of the returned reference.
63         unsafe { &mut *ptr.cast() }
64     }
65 
66     /// Creates a reference to an existing `struct cpumask` pointer.
67     ///
68     /// # Safety
69     ///
70     /// The caller must ensure that `ptr` is valid for reading and remains valid for the lifetime
71     /// of the returned reference.
as_ref<'a>(ptr: *const bindings::cpumask) -> &'a Self72     pub unsafe fn as_ref<'a>(ptr: *const bindings::cpumask) -> &'a Self {
73         // SAFETY: Guaranteed by the safety requirements of the function.
74         //
75         // INVARIANT: The caller ensures that `ptr` is valid for reading and remains valid for the
76         // lifetime of the returned reference.
77         unsafe { &*ptr.cast() }
78     }
79 
80     /// Obtain the raw `struct cpumask` pointer.
as_raw(&self) -> *mut bindings::cpumask81     pub fn as_raw(&self) -> *mut bindings::cpumask {
82         let this: *const Self = self;
83         this.cast_mut().cast()
84     }
85 
86     /// Set `cpu` in the cpumask.
87     ///
88     /// ATTENTION: Contrary to C, this Rust `set()` method is non-atomic.
89     /// This mismatches kernel naming convention and corresponds to the C
90     /// function `__cpumask_set_cpu()`.
91     #[inline]
set(&mut self, cpu: CpuId)92     pub fn set(&mut self, cpu: CpuId) {
93         // SAFETY: By the type invariant, `self.as_raw` is a valid argument to `__cpumask_set_cpu`.
94         unsafe { bindings::__cpumask_set_cpu(u32::from(cpu), self.as_raw()) };
95     }
96 
97     /// Clear `cpu` in the cpumask.
98     ///
99     /// ATTENTION: Contrary to C, this Rust `clear()` method is non-atomic.
100     /// This mismatches kernel naming convention and corresponds to the C
101     /// function `__cpumask_clear_cpu()`.
102     #[inline]
clear(&mut self, cpu: CpuId)103     pub fn clear(&mut self, cpu: CpuId) {
104         // SAFETY: By the type invariant, `self.as_raw` is a valid argument to
105         // `__cpumask_clear_cpu`.
106         unsafe { bindings::__cpumask_clear_cpu(i32::from(cpu), self.as_raw()) };
107     }
108 
109     /// Test `cpu` in the cpumask.
110     ///
111     /// Equivalent to the kernel's `cpumask_test_cpu` API.
112     #[inline]
test(&self, cpu: CpuId) -> bool113     pub fn test(&self, cpu: CpuId) -> bool {
114         // SAFETY: By the type invariant, `self.as_raw` is a valid argument to `cpumask_test_cpu`.
115         unsafe { bindings::cpumask_test_cpu(i32::from(cpu), self.as_raw()) }
116     }
117 
118     /// Set all CPUs in the cpumask.
119     ///
120     /// Equivalent to the kernel's `cpumask_setall` API.
121     #[inline]
setall(&mut self)122     pub fn setall(&mut self) {
123         // SAFETY: By the type invariant, `self.as_raw` is a valid argument to `cpumask_setall`.
124         unsafe { bindings::cpumask_setall(self.as_raw()) };
125     }
126 
127     /// Checks if cpumask is empty.
128     ///
129     /// Equivalent to the kernel's `cpumask_empty` API.
130     #[inline]
empty(&self) -> bool131     pub fn empty(&self) -> bool {
132         // SAFETY: By the type invariant, `self.as_raw` is a valid argument to `cpumask_empty`.
133         unsafe { bindings::cpumask_empty(self.as_raw()) }
134     }
135 
136     /// Checks if cpumask is full.
137     ///
138     /// Equivalent to the kernel's `cpumask_full` API.
139     #[inline]
full(&self) -> bool140     pub fn full(&self) -> bool {
141         // SAFETY: By the type invariant, `self.as_raw` is a valid argument to `cpumask_full`.
142         unsafe { bindings::cpumask_full(self.as_raw()) }
143     }
144 
145     /// Get weight of the cpumask.
146     ///
147     /// Equivalent to the kernel's `cpumask_weight` API.
148     #[inline]
weight(&self) -> u32149     pub fn weight(&self) -> u32 {
150         // SAFETY: By the type invariant, `self.as_raw` is a valid argument to `cpumask_weight`.
151         unsafe { bindings::cpumask_weight(self.as_raw()) }
152     }
153 
154     /// Copy cpumask.
155     ///
156     /// Equivalent to the kernel's `cpumask_copy` API.
157     #[inline]
copy(&self, dstp: &mut Self)158     pub fn copy(&self, dstp: &mut Self) {
159         // SAFETY: By the type invariant, `Self::as_raw` is a valid argument to `cpumask_copy`.
160         unsafe { bindings::cpumask_copy(dstp.as_raw(), self.as_raw()) };
161     }
162 }
163 
164 /// A CPU Mask pointer.
165 ///
166 /// Rust abstraction for the C `struct cpumask_var_t`.
167 ///
168 /// # Invariants
169 ///
170 /// A [`CpumaskVar`] instance always corresponds to a valid C `struct cpumask_var_t`.
171 ///
172 /// The callers must ensure that the `struct cpumask_var_t` is valid for access and remains valid
173 /// for the lifetime of [`CpumaskVar`].
174 ///
175 /// # Examples
176 ///
177 /// The following example demonstrates how to create and update a [`CpumaskVar`].
178 ///
179 /// ```
180 /// use kernel::cpu::CpuId;
181 /// use kernel::cpumask::CpumaskVar;
182 ///
183 /// let mut mask = CpumaskVar::new_zero(GFP_KERNEL).unwrap();
184 ///
185 /// assert!(mask.empty());
186 /// let mut count = 0;
187 ///
188 /// let cpu2 = CpuId::from_u32(2);
189 /// if let Some(cpu) = cpu2 {
190 ///     mask.set(cpu);
191 ///     assert!(mask.test(cpu));
192 ///     count += 1;
193 /// }
194 ///
195 /// let cpu3 = CpuId::from_u32(3);
196 /// if let Some(cpu) = cpu3 {
197 ///     mask.set(cpu);
198 ///     assert!(mask.test(cpu));
199 ///     count += 1;
200 /// }
201 ///
202 /// assert_eq!(mask.weight(), count);
203 ///
204 /// let mask2 = CpumaskVar::try_clone(&mask).unwrap();
205 ///
206 /// if let Some(cpu) = cpu2 {
207 ///     assert!(mask2.test(cpu));
208 /// }
209 ///
210 /// if let Some(cpu) = cpu3 {
211 ///     assert!(mask2.test(cpu));
212 /// }
213 /// assert_eq!(mask2.weight(), count);
214 /// ```
215 pub struct CpumaskVar {
216     #[cfg(CONFIG_CPUMASK_OFFSTACK)]
217     ptr: NonNull<Cpumask>,
218     #[cfg(not(CONFIG_CPUMASK_OFFSTACK))]
219     mask: Cpumask,
220 }
221 
222 impl CpumaskVar {
223     /// Creates a zero-initialized instance of the [`CpumaskVar`].
new_zero(_flags: Flags) -> Result<Self, AllocError>224     pub fn new_zero(_flags: Flags) -> Result<Self, AllocError> {
225         Ok(Self {
226             #[cfg(CONFIG_CPUMASK_OFFSTACK)]
227             ptr: {
228                 let mut ptr: *mut bindings::cpumask = ptr::null_mut();
229 
230                 // SAFETY: It is safe to call this method as the reference to `ptr` is valid.
231                 //
232                 // INVARIANT: The associated memory is freed when the `CpumaskVar` goes out of
233                 // scope.
234                 unsafe { bindings::zalloc_cpumask_var(&mut ptr, _flags.as_raw()) };
235                 NonNull::new(ptr.cast()).ok_or(AllocError)?
236             },
237 
238             #[cfg(not(CONFIG_CPUMASK_OFFSTACK))]
239             mask: Cpumask(Opaque::zeroed()),
240         })
241     }
242 
243     /// Creates an instance of the [`CpumaskVar`].
244     ///
245     /// # Safety
246     ///
247     /// The caller must ensure that the returned [`CpumaskVar`] is properly initialized before
248     /// getting used.
new(_flags: Flags) -> Result<Self, AllocError>249     pub unsafe fn new(_flags: Flags) -> Result<Self, AllocError> {
250         Ok(Self {
251             #[cfg(CONFIG_CPUMASK_OFFSTACK)]
252             ptr: {
253                 let mut ptr: *mut bindings::cpumask = ptr::null_mut();
254 
255                 // SAFETY: It is safe to call this method as the reference to `ptr` is valid.
256                 //
257                 // INVARIANT: The associated memory is freed when the `CpumaskVar` goes out of
258                 // scope.
259                 unsafe { bindings::alloc_cpumask_var(&mut ptr, _flags.as_raw()) };
260                 NonNull::new(ptr.cast()).ok_or(AllocError)?
261             },
262             #[cfg(not(CONFIG_CPUMASK_OFFSTACK))]
263             mask: Cpumask(Opaque::uninit()),
264         })
265     }
266 
267     /// Creates a mutable reference to an existing `struct cpumask_var_t` pointer.
268     ///
269     /// # Safety
270     ///
271     /// The caller must ensure that `ptr` is valid for writing and remains valid for the lifetime
272     /// of the returned reference.
as_mut_ref<'a>(ptr: *mut bindings::cpumask_var_t) -> &'a mut Self273     pub unsafe fn as_mut_ref<'a>(ptr: *mut bindings::cpumask_var_t) -> &'a mut Self {
274         // SAFETY: Guaranteed by the safety requirements of the function.
275         //
276         // INVARIANT: The caller ensures that `ptr` is valid for writing and remains valid for the
277         // lifetime of the returned reference.
278         unsafe { &mut *ptr.cast() }
279     }
280 
281     /// Creates a reference to an existing `struct cpumask_var_t` pointer.
282     ///
283     /// # Safety
284     ///
285     /// The caller must ensure that `ptr` is valid for reading and remains valid for the lifetime
286     /// of the returned reference.
as_ref<'a>(ptr: *const bindings::cpumask_var_t) -> &'a Self287     pub unsafe fn as_ref<'a>(ptr: *const bindings::cpumask_var_t) -> &'a Self {
288         // SAFETY: Guaranteed by the safety requirements of the function.
289         //
290         // INVARIANT: The caller ensures that `ptr` is valid for reading and remains valid for the
291         // lifetime of the returned reference.
292         unsafe { &*ptr.cast() }
293     }
294 
295     /// Clones cpumask.
try_clone(cpumask: &Cpumask) -> Result<Self>296     pub fn try_clone(cpumask: &Cpumask) -> Result<Self> {
297         // SAFETY: The returned cpumask_var is initialized right after this call.
298         let mut cpumask_var = unsafe { Self::new(GFP_KERNEL) }?;
299 
300         cpumask.copy(&mut cpumask_var);
301         Ok(cpumask_var)
302     }
303 }
304 
305 // Make [`CpumaskVar`] behave like a pointer to [`Cpumask`].
306 impl Deref for CpumaskVar {
307     type Target = Cpumask;
308 
309     #[cfg(CONFIG_CPUMASK_OFFSTACK)]
deref(&self) -> &Self::Target310     fn deref(&self) -> &Self::Target {
311         // SAFETY: The caller owns CpumaskVar, so it is safe to deref the cpumask.
312         unsafe { &*self.ptr.as_ptr() }
313     }
314 
315     #[cfg(not(CONFIG_CPUMASK_OFFSTACK))]
deref(&self) -> &Self::Target316     fn deref(&self) -> &Self::Target {
317         &self.mask
318     }
319 }
320 
321 impl DerefMut for CpumaskVar {
322     #[cfg(CONFIG_CPUMASK_OFFSTACK)]
deref_mut(&mut self) -> &mut Cpumask323     fn deref_mut(&mut self) -> &mut Cpumask {
324         // SAFETY: The caller owns CpumaskVar, so it is safe to deref the cpumask.
325         unsafe { self.ptr.as_mut() }
326     }
327 
328     #[cfg(not(CONFIG_CPUMASK_OFFSTACK))]
deref_mut(&mut self) -> &mut Cpumask329     fn deref_mut(&mut self) -> &mut Cpumask {
330         &mut self.mask
331     }
332 }
333 
334 impl Drop for CpumaskVar {
drop(&mut self)335     fn drop(&mut self) {
336         #[cfg(CONFIG_CPUMASK_OFFSTACK)]
337         // SAFETY: By the type invariant, `self.as_raw` is a valid argument to `free_cpumask_var`.
338         unsafe {
339             bindings::free_cpumask_var(self.as_raw())
340         };
341     }
342 }
343