cpumask.rs

  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
  7use crate::{
  8    alloc::{AllocError, Flags},
  9    cpu::CpuId,
 10    prelude::*,
 11    types::Opaque,
 12};
 13
 14#[cfg(CONFIG_CPUMASK_OFFSTACK)]
 15use core::ptr::{self, NonNull};
 16
 17use 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)]
 49pub struct Cpumask(Opaque<bindings::cpumask>);
 50
 51impl 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.
 58    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.
 72    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.
 81    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]
 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]
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]
113    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]
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]
131    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]
140    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]
149    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]
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#[repr(transparent)]
216pub struct CpumaskVar {
217    #[cfg(CONFIG_CPUMASK_OFFSTACK)]
218    ptr: NonNull<Cpumask>,
219    #[cfg(not(CONFIG_CPUMASK_OFFSTACK))]
220    mask: Cpumask,
221}
222
223impl CpumaskVar {
224    /// Creates a zero-initialized instance of the [`CpumaskVar`].
225    pub fn new_zero(_flags: Flags) -> Result<Self, AllocError> {
226        Ok(Self {
227            #[cfg(CONFIG_CPUMASK_OFFSTACK)]
228            ptr: {
229                let mut ptr: *mut bindings::cpumask = ptr::null_mut();
230
231                // SAFETY: It is safe to call this method as the reference to `ptr` is valid.
232                //
233                // INVARIANT: The associated memory is freed when the `CpumaskVar` goes out of
234                // scope.
235                unsafe { bindings::zalloc_cpumask_var(&mut ptr, _flags.as_raw()) };
236                NonNull::new(ptr.cast()).ok_or(AllocError)?
237            },
238
239            #[cfg(not(CONFIG_CPUMASK_OFFSTACK))]
240            mask: Cpumask(Opaque::zeroed()),
241        })
242    }
243
244    /// Creates an instance of the [`CpumaskVar`].
245    ///
246    /// # Safety
247    ///
248    /// The caller must ensure that the returned [`CpumaskVar`] is properly initialized before
249    /// getting used.
250    pub unsafe fn new(_flags: Flags) -> Result<Self, AllocError> {
251        Ok(Self {
252            #[cfg(CONFIG_CPUMASK_OFFSTACK)]
253            ptr: {
254                let mut ptr: *mut bindings::cpumask = ptr::null_mut();
255
256                // SAFETY: It is safe to call this method as the reference to `ptr` is valid.
257                //
258                // INVARIANT: The associated memory is freed when the `CpumaskVar` goes out of
259                // scope.
260                unsafe { bindings::alloc_cpumask_var(&mut ptr, _flags.as_raw()) };
261                NonNull::new(ptr.cast()).ok_or(AllocError)?
262            },
263            #[cfg(not(CONFIG_CPUMASK_OFFSTACK))]
264            mask: Cpumask(Opaque::uninit()),
265        })
266    }
267
268    /// Creates a mutable reference to an existing `struct cpumask_var_t` pointer.
269    ///
270    /// # Safety
271    ///
272    /// The caller must ensure that `ptr` is valid for writing and remains valid for the lifetime
273    /// of the returned reference.
274    pub unsafe fn from_raw_mut<'a>(ptr: *mut bindings::cpumask_var_t) -> &'a mut Self {
275        // SAFETY: Guaranteed by the safety requirements of the function.
276        //
277        // INVARIANT: The caller ensures that `ptr` is valid for writing and remains valid for the
278        // lifetime of the returned reference.
279        unsafe { &mut *ptr.cast() }
280    }
281
282    /// Creates a reference to an existing `struct cpumask_var_t` pointer.
283    ///
284    /// # Safety
285    ///
286    /// The caller must ensure that `ptr` is valid for reading and remains valid for the lifetime
287    /// of the returned reference.
288    pub unsafe fn from_raw<'a>(ptr: *const bindings::cpumask_var_t) -> &'a Self {
289        // SAFETY: Guaranteed by the safety requirements of the function.
290        //
291        // INVARIANT: The caller ensures that `ptr` is valid for reading and remains valid for the
292        // lifetime of the returned reference.
293        unsafe { &*ptr.cast() }
294    }
295
296    /// Clones cpumask.
297    pub fn try_clone(cpumask: &Cpumask) -> Result<Self> {
298        // SAFETY: The returned cpumask_var is initialized right after this call.
299        let mut cpumask_var = unsafe { Self::new(GFP_KERNEL) }?;
300
301        cpumask.copy(&mut cpumask_var);
302        Ok(cpumask_var)
303    }
304}
305
306// Make [`CpumaskVar`] behave like a pointer to [`Cpumask`].
307impl Deref for CpumaskVar {
308    type Target = Cpumask;
309
310    #[cfg(CONFIG_CPUMASK_OFFSTACK)]
311    fn deref(&self) -> &Self::Target {
312        // SAFETY: The caller owns CpumaskVar, so it is safe to deref the cpumask.
313        unsafe { &*self.ptr.as_ptr() }
314    }
315
316    #[cfg(not(CONFIG_CPUMASK_OFFSTACK))]
317    fn deref(&self) -> &Self::Target {
318        &self.mask
319    }
320}
321
322impl DerefMut for CpumaskVar {
323    #[cfg(CONFIG_CPUMASK_OFFSTACK)]
324    fn deref_mut(&mut self) -> &mut Cpumask {
325        // SAFETY: The caller owns CpumaskVar, so it is safe to deref the cpumask.
326        unsafe { self.ptr.as_mut() }
327    }
328
329    #[cfg(not(CONFIG_CPUMASK_OFFSTACK))]
330    fn deref_mut(&mut self) -> &mut Cpumask {
331        &mut self.mask
332    }
333}
334
335impl Drop for CpumaskVar {
336    fn drop(&mut self) {
337        #[cfg(CONFIG_CPUMASK_OFFSTACK)]
338        // SAFETY: By the type invariant, `self.as_raw` is a valid argument to `free_cpumask_var`.
339        unsafe {
340            bindings::free_cpumask_var(self.as_raw())
341        };
342    }
343}