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