xref: /linux/rust/kernel/bitmap.rs (revision 12d5fc96ebbb3b948861805b32a4c3a96105ccd9)
1 // SPDX-License-Identifier: GPL-2.0
2 
3 // Copyright (C) 2025 Google LLC.
4 
5 //! Rust API for bitmap.
6 //!
7 //! C headers: [`include/linux/bitmap.h`](srctree/include/linux/bitmap.h).
8 
9 use crate::alloc::{AllocError, Flags};
10 use crate::bindings;
11 #[cfg(not(CONFIG_RUST_BITMAP_HARDENED))]
12 use crate::pr_err;
13 use crate::ptr::Alignment;
14 use core::{
15     num::NonZero,
16     ptr::NonNull, //
17 };
18 
19 /// Represents a C bitmap. Wraps underlying C bitmap API.
20 ///
21 /// # Invariants
22 ///
23 /// Must reference a `[c_ulong]` long enough to fit `data.len()` bits.
24 /// Must not be longer than [`i32::MAX`] bits, so offsets and lengths used with
25 /// [`Bitmap`] functions fit in the `int` and `unsigned int` arguments of the C bitmap API.
26 /// This also matches [`BitmapVec::MAX_LEN`].
27 #[cfg_attr(CONFIG_64BIT, repr(align(8)))]
28 #[cfg_attr(not(CONFIG_64BIT), repr(align(4)))]
29 pub struct Bitmap {
30     data: [()],
31 }
32 
33 macro_rules! bitmap_assert {
34     ($cond:expr, $($arg:tt)+) => {
35         #[cfg(CONFIG_RUST_BITMAP_HARDENED)]
36         assert!($cond, $($arg)*);
37     }
38 }
39 
40 macro_rules! bitmap_assert_return {
41     ($cond:expr, $($arg:tt)+) => {
42         #[cfg(CONFIG_RUST_BITMAP_HARDENED)]
43         assert!($cond, $($arg)*);
44 
45         #[cfg(not(CONFIG_RUST_BITMAP_HARDENED))]
46         if !($cond) {
47             pr_err!($($arg)*);
48             return
49         }
50     }
51 }
52 
53 impl Bitmap {
54     /// Borrows a C bitmap.
55     ///
56     /// # Panics
57     ///
58     /// Panics if CONFIG_RUST_BITMAP_HARDENED is enabled and `nbits` exceeds [`i32::MAX`].
59     ///
60     /// # Safety
61     ///
62     /// * `ptr` holds a non-null address of an initialized array of `unsigned long`
63     ///   that is large enough to hold `nbits` bits.
64     /// * `nbits` must not exceed [`i32::MAX`].
65     /// * the array must not be freed for the lifetime of this [`Bitmap`]
66     /// * concurrent access only happens through atomic operations
67     pub unsafe fn from_raw<'a>(ptr: *const usize, nbits: usize) -> &'a Bitmap {
68         bitmap_assert!(
69             nbits <= i32::MAX as usize,
70             "`nbits` must be <= {}, was {}",
71             i32::MAX,
72             nbits
73         );
74         let data: *const [()] = core::ptr::slice_from_raw_parts(ptr.cast(), nbits);
75         // INVARIANT: `data` references an initialized array that can hold `nbits` bits.
76         // INVARIANT: The caller guarantees that `nbits` does not exceed `i32::MAX`.
77         // SAFETY:
78         // The caller guarantees that `data` (derived from `ptr` and `nbits`)
79         // points to a valid, initialized, and appropriately sized memory region
80         // that will not be freed for the lifetime 'a.
81         // We are casting `*const [()]` to `*const Bitmap`. The `Bitmap`
82         // struct is a ZST with a `data: [()]` field. This means its layout
83         // is compatible with a slice of `()`, and effectively it's a "thin pointer"
84         // (its size is 0 and alignment is 1). The `slice_from_raw_parts`
85         // function correctly encodes the length (number of bits, not elements)
86         // into the metadata of the fat pointer. Therefore, dereferencing this
87         // pointer as `&Bitmap` is safe given the caller's guarantees.
88         unsafe { &*(data as *const Bitmap) }
89     }
90 
91     /// Borrows a C bitmap exclusively.
92     ///
93     /// # Panics
94     ///
95     /// Panics if CONFIG_RUST_BITMAP_HARDENED is enabled and `nbits` exceeds [`i32::MAX`].
96     ///
97     /// # Safety
98     ///
99     /// * `ptr` holds a non-null address of an initialized array of `unsigned long`
100     ///   that is large enough to hold `nbits` bits.
101     /// * `nbits` must not exceed [`i32::MAX`].
102     /// * the array must not be freed for the lifetime of this [`Bitmap`]
103     /// * no concurrent access may happen.
104     pub unsafe fn from_raw_mut<'a>(ptr: *mut usize, nbits: usize) -> &'a mut Bitmap {
105         bitmap_assert!(
106             nbits <= i32::MAX as usize,
107             "`nbits` must be <= {}, was {}",
108             i32::MAX,
109             nbits
110         );
111         let data: *mut [()] = core::ptr::slice_from_raw_parts_mut(ptr.cast(), nbits);
112         // INVARIANT: `data` references an initialized array that can hold `nbits` bits.
113         // INVARIANT: The caller guarantees that `nbits` does not exceed `i32::MAX`.
114         // SAFETY:
115         // The caller guarantees that `data` (derived from `ptr` and `nbits`)
116         // points to a valid, initialized, and appropriately sized memory region
117         // that will not be freed for the lifetime 'a.
118         // Furthermore, the caller guarantees no concurrent access will happen,
119         // which upholds the exclusivity requirement for a mutable reference.
120         // Similar to `from_raw`, casting `*mut [()]` to `*mut Bitmap` is
121         // safe because `Bitmap` is a ZST with a `data: [()]` field,
122         // making its layout compatible with a slice of `()`.
123         unsafe { &mut *(data as *mut Bitmap) }
124     }
125 
126     /// Returns a raw pointer to the backing [`Bitmap`].
127     pub fn as_ptr(&self) -> *const usize {
128         core::ptr::from_ref::<Bitmap>(self).cast::<usize>()
129     }
130 
131     /// Returns a mutable raw pointer to the backing [`Bitmap`].
132     pub fn as_mut_ptr(&mut self) -> *mut usize {
133         core::ptr::from_mut::<Bitmap>(self).cast::<usize>()
134     }
135 
136     /// Returns length of this [`Bitmap`].
137     #[expect(clippy::len_without_is_empty)]
138     pub fn len(&self) -> usize {
139         self.data.len()
140     }
141 }
142 
143 /// Holds either a pointer to array of `unsigned long` or a small bitmap.
144 #[repr(C)]
145 union BitmapRepr {
146     bitmap: usize,
147     ptr: NonNull<usize>,
148 }
149 
150 /// Represents an owned bitmap.
151 ///
152 /// Wraps underlying C bitmap API. See [`Bitmap`] for available
153 /// methods.
154 ///
155 /// # Examples
156 ///
157 /// Basic usage
158 ///
159 /// ```
160 /// use kernel::alloc::flags::GFP_KERNEL;
161 /// use kernel::bitmap::BitmapVec;
162 ///
163 /// let mut b = BitmapVec::new(16, GFP_KERNEL)?;
164 ///
165 /// assert_eq!(16, b.len());
166 /// for i in 0..16 {
167 ///     if i % 4 == 0 {
168 ///       b.set_bit(i);
169 ///     }
170 /// }
171 /// assert_eq!(Some(0), b.next_bit(0));
172 /// assert_eq!(Some(1), b.next_zero_bit(0));
173 /// assert_eq!(Some(4), b.next_bit(1));
174 /// assert_eq!(Some(5), b.next_zero_bit(4));
175 /// assert_eq!(Some(12), b.last_bit());
176 /// # Ok::<(), Error>(())
177 /// ```
178 ///
179 /// # Invariants
180 ///
181 /// * `nbits` is `<= MAX_LEN`.
182 /// * if `nbits <= MAX_INLINE_LEN`, then `repr` is a `usize`.
183 /// * otherwise, `repr` holds a non-null pointer to an initialized
184 ///   array of `unsigned long` that is large enough to hold `nbits` bits.
185 pub struct BitmapVec {
186     /// Representation of bitmap.
187     repr: BitmapRepr,
188     /// Length of this bitmap. Must be `<= MAX_LEN`.
189     nbits: usize,
190 }
191 
192 impl core::ops::Deref for BitmapVec {
193     type Target = Bitmap;
194 
195     fn deref(&self) -> &Bitmap {
196         let ptr = if self.nbits <= BitmapVec::MAX_INLINE_LEN {
197             // SAFETY: Bitmap is represented inline.
198             #[allow(unused_unsafe, reason = "Safe since Rust 1.92.0")]
199             unsafe {
200                 core::ptr::addr_of!(self.repr.bitmap)
201             }
202         } else {
203             // SAFETY: Bitmap is represented as array of `unsigned long`.
204             unsafe { self.repr.ptr.as_ptr() }
205         };
206 
207         // SAFETY: We got the right pointer and invariants of [`Bitmap`] hold.
208         // An inline bitmap is treated like an array with single element.
209         unsafe { Bitmap::from_raw(ptr, self.nbits) }
210     }
211 }
212 
213 impl core::ops::DerefMut for BitmapVec {
214     fn deref_mut(&mut self) -> &mut Bitmap {
215         let ptr = if self.nbits <= BitmapVec::MAX_INLINE_LEN {
216             // SAFETY: Bitmap is represented inline.
217             #[allow(unused_unsafe, reason = "Safe since Rust 1.92.0")]
218             unsafe {
219                 core::ptr::addr_of_mut!(self.repr.bitmap)
220             }
221         } else {
222             // SAFETY: Bitmap is represented as array of `unsigned long`.
223             unsafe { self.repr.ptr.as_ptr() }
224         };
225 
226         // SAFETY: We got the right pointer and invariants of [`BitmapVec`] hold.
227         // An inline bitmap is treated like an array with single element.
228         unsafe { Bitmap::from_raw_mut(ptr, self.nbits) }
229     }
230 }
231 
232 /// Enable ownership transfer to other threads.
233 ///
234 /// SAFETY: We own the underlying bitmap representation.
235 unsafe impl Send for BitmapVec {}
236 
237 /// Enable unsynchronized concurrent access to [`BitmapVec`] through shared references.
238 ///
239 /// SAFETY: `deref()` will return a reference to a [`Bitmap`]. Its methods
240 /// take immutable references are either atomic or read-only.
241 unsafe impl Sync for BitmapVec {}
242 
243 impl Drop for BitmapVec {
244     fn drop(&mut self) {
245         if self.nbits <= BitmapVec::MAX_INLINE_LEN {
246             return;
247         }
248         // SAFETY: `self.ptr` was returned by the C `bitmap_zalloc`.
249         //
250         // INVARIANT: there is no other use of the `self.ptr` after this
251         // call and the value is being dropped so the broken invariant is
252         // not observable on function exit.
253         unsafe { bindings::bitmap_free(self.repr.ptr.as_ptr()) };
254     }
255 }
256 
257 impl BitmapVec {
258     /// The maximum possible length of a `BitmapVec`.
259     pub const MAX_LEN: usize = i32::MAX as usize;
260 
261     /// The maximum length that uses the inline representation.
262     pub const MAX_INLINE_LEN: usize = usize::BITS as usize;
263 
264     /// Construct a longest possible inline [`BitmapVec`].
265     #[inline]
266     pub fn new_inline() -> Self {
267         // INVARIANT: `nbits <= MAX_INLINE_LEN`, so an inline bitmap is the right repr.
268         BitmapVec {
269             repr: BitmapRepr { bitmap: 0 },
270             nbits: BitmapVec::MAX_INLINE_LEN,
271         }
272     }
273 
274     /// Constructs a new [`BitmapVec`].
275     ///
276     /// Fails with [`AllocError`] when the [`BitmapVec`] could not be allocated. This
277     /// includes the case when `nbits` is greater than `MAX_LEN`.
278     #[inline]
279     pub fn new(nbits: usize, flags: Flags) -> Result<Self, AllocError> {
280         if nbits <= BitmapVec::MAX_INLINE_LEN {
281             return Ok(BitmapVec {
282                 repr: BitmapRepr { bitmap: 0 },
283                 nbits,
284             });
285         }
286         if nbits > Self::MAX_LEN {
287             return Err(AllocError);
288         }
289         let nbits_u32 = u32::try_from(nbits).unwrap();
290         // SAFETY: `MAX_INLINE_LEN < nbits` and `nbits <= MAX_LEN`.
291         let ptr = unsafe { bindings::bitmap_zalloc(nbits_u32, flags.as_raw()) };
292         let ptr = NonNull::new(ptr).ok_or(AllocError)?;
293         // INVARIANT: `ptr` returned by C `bitmap_zalloc` and `nbits` checked.
294         Ok(BitmapVec {
295             repr: BitmapRepr { ptr },
296             nbits,
297         })
298     }
299 
300     /// Returns length of this [`Bitmap`].
301     #[allow(clippy::len_without_is_empty)]
302     #[inline]
303     pub fn len(&self) -> usize {
304         self.nbits
305     }
306 
307     /// Fills this `Bitmap` with random bits.
308     #[cfg(CONFIG_FIND_BIT_BENCHMARK_RUST)]
309     pub fn fill_random(&mut self) {
310         // SAFETY: `self.as_mut_ptr` points to either an array of the
311         // appropriate length or one usize.
312         unsafe {
313             bindings::get_random_bytes(
314                 self.as_mut_ptr().cast::<ffi::c_void>(),
315                 usize::div_ceil(self.nbits, bindings::BITS_PER_LONG as usize)
316                     * bindings::BITS_PER_LONG as usize
317                     / 8,
318             );
319         }
320     }
321 }
322 
323 impl Bitmap {
324     /// Set bit with index `index`.
325     ///
326     /// ATTENTION: `set_bit` is non-atomic, which differs from the naming
327     /// convention in C code. The corresponding C function is `__set_bit`.
328     ///
329     /// If CONFIG_RUST_BITMAP_HARDENED is not enabled and `index` is greater than
330     /// or equal to `self.nbits`, does nothing.
331     ///
332     /// # Panics
333     ///
334     /// Panics if CONFIG_RUST_BITMAP_HARDENED is enabled and `index` is greater than
335     /// or equal to `self.nbits`.
336     #[inline]
337     pub fn set_bit(&mut self, index: usize) {
338         bitmap_assert_return!(
339             index < self.len(),
340             "Bit `index` must be < {}, was {}",
341             self.len(),
342             index
343         );
344         // SAFETY: Bit `index` is within bounds.
345         unsafe { bindings::__set_bit(index, self.as_mut_ptr()) };
346     }
347 
348     /// Set bit with index `index`, atomically.
349     ///
350     /// This is a relaxed atomic operation (no implied memory barriers).
351     ///
352     /// ATTENTION: The naming convention differs from C, where the corresponding
353     /// function is called `set_bit`.
354     ///
355     /// If CONFIG_RUST_BITMAP_HARDENED is not enabled and `index` is greater than
356     /// or equal to `self.len()`, does nothing.
357     ///
358     /// # Panics
359     ///
360     /// Panics if CONFIG_RUST_BITMAP_HARDENED is enabled and `index` is greater than
361     /// or equal to `self.len()`.
362     #[inline]
363     pub fn set_bit_atomic(&self, index: usize) {
364         bitmap_assert_return!(
365             index < self.len(),
366             "Bit `index` must be < {}, was {}",
367             self.len(),
368             index
369         );
370         // SAFETY: `index` is within bounds and the caller has ensured that
371         // there is no mix of non-atomic and atomic operations.
372         unsafe { bindings::set_bit(index, self.as_ptr().cast_mut()) };
373     }
374 
375     /// Clear `index` bit.
376     ///
377     /// ATTENTION: `clear_bit` is non-atomic, which differs from the naming
378     /// convention in C code. The corresponding C function is `__clear_bit`.
379     ///
380     /// If CONFIG_RUST_BITMAP_HARDENED is not enabled and `index` is greater than
381     /// or equal to `self.len()`, does nothing.
382     ///
383     /// # Panics
384     ///
385     /// Panics if CONFIG_RUST_BITMAP_HARDENED is enabled and `index` is greater than
386     /// or equal to `self.len()`.
387     #[inline]
388     pub fn clear_bit(&mut self, index: usize) {
389         bitmap_assert_return!(
390             index < self.len(),
391             "Bit `index` must be < {}, was {}",
392             self.len(),
393             index
394         );
395         // SAFETY: `index` is within bounds.
396         unsafe { bindings::__clear_bit(index, self.as_mut_ptr()) };
397     }
398 
399     /// Clear `index` bit, atomically.
400     ///
401     /// This is a relaxed atomic operation (no implied memory barriers).
402     ///
403     /// ATTENTION: The naming convention differs from C, where the corresponding
404     /// function is called `clear_bit`.
405     ///
406     /// If CONFIG_RUST_BITMAP_HARDENED is not enabled and `index` is greater than
407     /// or equal to `self.len()`, does nothing.
408     ///
409     /// # Panics
410     ///
411     /// Panics if CONFIG_RUST_BITMAP_HARDENED is enabled and `index` is greater than
412     /// or equal to `self.len()`.
413     #[inline]
414     pub fn clear_bit_atomic(&self, index: usize) {
415         bitmap_assert_return!(
416             index < self.len(),
417             "Bit `index` must be < {}, was {}",
418             self.len(),
419             index
420         );
421         // SAFETY: `index` is within bounds and the caller has ensured that
422         // there is no mix of non-atomic and atomic operations.
423         unsafe { bindings::clear_bit(index, self.as_ptr().cast_mut()) };
424     }
425 
426     /// Copy `src` into this [`Bitmap`] and set any remaining bits to zero.
427     ///
428     /// # Examples
429     ///
430     /// ```
431     /// use kernel::alloc::{AllocError, flags::GFP_KERNEL};
432     /// use kernel::bitmap::BitmapVec;
433     ///
434     /// let mut long_bitmap = BitmapVec::new(256, GFP_KERNEL)?;
435     ///
436     /// assert_eq!(None, long_bitmap.last_bit());
437     ///
438     /// let mut short_bitmap = BitmapVec::new(16, GFP_KERNEL)?;
439     ///
440     /// short_bitmap.set_bit(7);
441     /// long_bitmap.copy_and_extend(&short_bitmap);
442     /// assert_eq!(Some(7), long_bitmap.last_bit());
443     ///
444     /// # Ok::<(), AllocError>(())
445     /// ```
446     #[inline]
447     pub fn copy_and_extend(&mut self, src: &Bitmap) {
448         let len = core::cmp::min(src.len(), self.len());
449         // SAFETY: Access to `self` and `src` is within bounds. Both lengths fit in `u32`
450         // because a `Bitmap` is at most `i32::MAX` bits, so the casts are lossless.
451         unsafe {
452             bindings::bitmap_copy_and_extend(
453                 self.as_mut_ptr(),
454                 src.as_ptr(),
455                 len as u32,
456                 self.len() as u32,
457             )
458         };
459     }
460 
461     /// Finds last set bit.
462     ///
463     /// # Examples
464     ///
465     /// ```
466     /// use kernel::alloc::{AllocError, flags::GFP_KERNEL};
467     /// use kernel::bitmap::BitmapVec;
468     ///
469     /// let bitmap = BitmapVec::new(64, GFP_KERNEL)?;
470     ///
471     /// match bitmap.last_bit() {
472     ///     Some(idx) => {
473     ///         pr_info!("The last bit has index {idx}.\n");
474     ///     }
475     ///     None => {
476     ///         pr_info!("All bits in this bitmap are 0.\n");
477     ///     }
478     /// }
479     /// # Ok::<(), AllocError>(())
480     /// ```
481     #[inline]
482     pub fn last_bit(&self) -> Option<usize> {
483         // SAFETY: `_find_next_bit` access is within bounds due to invariant.
484         let index = unsafe { bindings::_find_last_bit(self.as_ptr(), self.len()) };
485         if index >= self.len() {
486             None
487         } else {
488             Some(index)
489         }
490     }
491 
492     /// Finds next set bit, starting from `start`.
493     ///
494     /// Returns `None` if no bits are set on or after the given index. The
495     /// index `start` must be in bounds.
496     ///
497     /// # Panics
498     ///
499     /// Panics if `CONFIG_RUST_BITMAP_HARDENED` is enabled and `start` is
500     /// out of bounds.
501     #[inline]
502     pub fn next_bit(&self, start: usize) -> Option<usize> {
503         bitmap_assert!(
504             start < self.len(),
505             "`start` must be < {} was {}",
506             self.len(),
507             start
508         );
509         // SAFETY: `_find_next_bit` tolerates out-of-bounds arguments and returns a
510         // value larger than or equal to `self.len()` in that case.
511         let index = unsafe { bindings::_find_next_bit(self.as_ptr(), self.len(), start) };
512         if index >= self.len() {
513             None
514         } else {
515             Some(index)
516         }
517     }
518 
519     /// Finds next zero bit, starting from `start`.
520     ///
521     /// Returns `None` if no bits are zero on or after the given index. The
522     /// index `start` must be in bounds.
523     ///
524     /// # Panics
525     ///
526     /// Panics if `CONFIG_RUST_BITMAP_HARDENED` is enabled and `start` is
527     /// out of bounds.
528     #[inline]
529     pub fn next_zero_bit(&self, start: usize) -> Option<usize> {
530         bitmap_assert!(
531             start < self.len(),
532             "`start` must be < {} was {}",
533             self.len(),
534             start
535         );
536         // SAFETY: `_find_next_zero_bit` tolerates out-of-bounds arguments and returns a
537         // value larger than or equal to `self.len()` in that case.
538         let index = unsafe { bindings::_find_next_zero_bit(self.as_ptr(), self.len(), start) };
539         if index >= self.len() {
540             None
541         } else {
542             Some(index)
543         }
544     }
545 
546     /// Finds a contiguous area of `nbits` zero bits at or after `start`, where the area plus
547     /// `align_offset` is aligned to `align`.
548     ///
549     /// Returns the bit index of the start of the area, or [`None`] if no such area fitting in
550     /// the bitmap exists.
551     ///
552     /// The returned index plus `align_offset` is a multiple of `align`.
553     ///
554     /// # Panics
555     ///
556     /// Panics if CONFIG_RUST_BITMAP_HARDENED is enabled and `start` is out of bounds.
557     #[inline]
558     pub fn next_zero_area_off(
559         &self,
560         start: usize,
561         nbits: NonZero<usize>,
562         align: Alignment,
563         align_offset: usize,
564     ) -> Option<usize> {
565         bitmap_assert!(
566             start < self.len(),
567             "`start` must be < {}, was {}",
568             self.len(),
569             start
570         );
571 
572         let nr = u32::try_from(nbits.get()).ok()?;
573         let align_mask = align.as_usize() - 1;
574 
575         // The C alignment and end arithmetic must not overflow, or it can read out of bounds.
576         // Overflow is only possible on 32-bit.
577         #[cfg(not(CONFIG_64BIT))]
578         align_mask
579             .checked_add(self.len())?
580             .checked_add(nbits.get())?;
581 
582         // SAFETY: `bitmap_find_next_zero_area_off` is safe to use with an out of bounds `start`
583         // value and, given the overflow check above, never reads beyond `self.len()` bits.
584         let index = unsafe {
585             bindings::bitmap_find_next_zero_area_off(
586                 self.as_ptr().cast_mut(),
587                 self.len(),
588                 start,
589                 nr,
590                 align_mask,
591                 align_offset,
592             )
593         };
594 
595         (index < self.len()).then_some(index)
596     }
597 
598     /// Finds a contiguous area of `nbits` zero bits at or after `start`, aligned to `align`.
599     ///
600     /// Returns the bit index of the start of the area, or [`None`] if no such area fitting in
601     /// the bitmap exists.
602     ///
603     /// The returned index is a multiple of `align`.
604     ///
605     /// # Panics
606     ///
607     /// Panics if CONFIG_RUST_BITMAP_HARDENED is enabled and `start` is out of bounds.
608     ///
609     /// # Examples
610     ///
611     /// ```
612     /// use kernel::{
613     ///     alloc::{AllocError, flags::GFP_KERNEL},
614     ///     bitmap::BitmapVec,
615     ///     ptr::Alignment,
616     ///     sizes::SizeConstants, //
617     /// };
618     ///
619     /// let mut b = BitmapVec::new(64, GFP_KERNEL)?;
620     ///
621     /// assert_eq!(Some(0), b.next_zero_area(0, cv!(8), Alignment::SZ_1));
622     /// b.set(0, cv!(5));
623     /// assert_eq!(Some(5), b.next_zero_area(0, cv!(8), Alignment::SZ_1));
624     /// assert_eq!(Some(8), b.next_zero_area(0, cv!(8), Alignment::SZ_8));
625     /// assert_eq!(None, b.next_zero_area(0, cv!(65), Alignment::SZ_1));
626     /// # Ok::<(), AllocError>(())
627     /// ```
628     #[inline]
629     pub fn next_zero_area(
630         &self,
631         start: usize,
632         nbits: NonZero<usize>,
633         align: Alignment,
634     ) -> Option<usize> {
635         self.next_zero_area_off(start, nbits, align, 0)
636     }
637 
638     /// Sets a contiguous area of `nbits` bits starting at `start`.
639     ///
640     /// If CONFIG_RUST_BITMAP_HARDENED is not enabled and the area `start..start + nbits` is out of
641     /// bounds, does nothing.
642     ///
643     /// # Panics
644     ///
645     /// Panics if CONFIG_RUST_BITMAP_HARDENED is enabled and the area `start..start + nbits` is out
646     /// of bounds.
647     #[inline]
648     pub fn set(&mut self, start: usize, nbits: NonZero<usize>) {
649         bitmap_assert_return!(
650             start
651                 .checked_add(nbits.get())
652                 .is_some_and(|end| end <= self.len()),
653             "Area `start..start + nbits` ({}..{}) must be within bounds {}",
654             start,
655             start.saturating_add(nbits.get()),
656             self.len()
657         );
658         // SAFETY: The area `start..start + nbits` is within bounds and a `Bitmap` is at most
659         // `i32::MAX` bits, so the casts are lossless.
660         unsafe { bindings::__bitmap_set(self.as_mut_ptr(), start as u32, nbits.get() as i32) };
661     }
662 
663     /// Clears a contiguous area of `nbits` bits starting at `start`.
664     ///
665     /// If CONFIG_RUST_BITMAP_HARDENED is not enabled and the area `start..start + nbits` is out of
666     /// bounds, does nothing.
667     ///
668     /// # Panics
669     ///
670     /// Panics if CONFIG_RUST_BITMAP_HARDENED is enabled and the area `start..start + nbits` is out
671     /// of bounds.
672     #[inline]
673     pub fn clear(&mut self, start: usize, nbits: NonZero<usize>) {
674         bitmap_assert_return!(
675             start
676                 .checked_add(nbits.get())
677                 .is_some_and(|end| end <= self.len()),
678             "Area `start..start + nbits` ({}..{}) must be within bounds {}",
679             start,
680             start.saturating_add(nbits.get()),
681             self.len()
682         );
683         // SAFETY: The area `start..start + nbits` is within bounds and a `Bitmap` is at most
684         // `i32::MAX` bits, so the casts are lossless.
685         unsafe { bindings::__bitmap_clear(self.as_mut_ptr(), start as u32, nbits.get() as i32) };
686     }
687 }
688 
689 #[cfg(CONFIG_RUST_BITMAP_KUNIT_TEST)]
690 #[macros::kunit_tests(rust_kernel_bitmap)]
691 mod tests {
692     use super::*;
693     use kernel::{
694         alloc::flags::GFP_KERNEL,
695         num::cv,
696         sizes::SizeConstants, //
697     };
698 
699     #[test]
700     fn bitmap_borrow() {
701         let fake_bitmap: [usize; 2] = [0, 0];
702         let fake_bitmap_len = 2 * usize::BITS as usize;
703         // SAFETY: `fake_c_bitmap` is an array of expected length.
704         let b = unsafe { Bitmap::from_raw(fake_bitmap.as_ptr(), fake_bitmap_len) };
705         assert_eq!(fake_bitmap_len, b.len());
706         assert_eq!(None, b.next_bit(0));
707     }
708 
709     #[test]
710     fn bitmap_copy() {
711         let fake_bitmap: usize = 0xFF;
712         // SAFETY: `fake_c_bitmap` can be used as one-element array of expected length.
713         let b = unsafe { Bitmap::from_raw(core::ptr::addr_of!(fake_bitmap), 8) };
714         assert_eq!(8, b.len());
715         assert_eq!(None, b.next_zero_bit(0));
716     }
717 
718     #[test]
719     fn bitmap_vec_new() -> Result<(), AllocError> {
720         let b = BitmapVec::new(0, GFP_KERNEL)?;
721         assert_eq!(0, b.len());
722 
723         let b = BitmapVec::new(3, GFP_KERNEL)?;
724         assert_eq!(3, b.len());
725 
726         let b = BitmapVec::new(1024, GFP_KERNEL)?;
727         assert_eq!(1024, b.len());
728 
729         // Requesting too large values results in [`AllocError`].
730         let res = BitmapVec::new(1 << 31, GFP_KERNEL);
731         assert!(res.is_err());
732         Ok(())
733     }
734 
735     #[test]
736     fn bitmap_set_clear_find() -> Result<(), AllocError> {
737         let mut b = BitmapVec::new(128, GFP_KERNEL)?;
738 
739         // Zero-initialized
740         assert_eq!(None, b.next_bit(0));
741         assert_eq!(Some(0), b.next_zero_bit(0));
742         assert_eq!(None, b.last_bit());
743 
744         b.set_bit(17);
745 
746         assert_eq!(Some(17), b.next_bit(0));
747         assert_eq!(Some(17), b.next_bit(17));
748         assert_eq!(None, b.next_bit(18));
749         assert_eq!(Some(17), b.last_bit());
750 
751         b.set_bit(107);
752 
753         assert_eq!(Some(17), b.next_bit(0));
754         assert_eq!(Some(17), b.next_bit(17));
755         assert_eq!(Some(107), b.next_bit(18));
756         assert_eq!(Some(107), b.last_bit());
757 
758         b.clear_bit(17);
759 
760         assert_eq!(Some(107), b.next_bit(0));
761         assert_eq!(Some(107), b.last_bit());
762         Ok(())
763     }
764 
765     #[test]
766     #[cfg(not(CONFIG_RUST_BITMAP_HARDENED))]
767     fn owned_bitmap_out_of_bounds() -> Result<(), AllocError> {
768         let mut b = BitmapVec::new(128, GFP_KERNEL)?;
769 
770         b.set_bit(2048);
771         b.set_bit_atomic(2048);
772         b.clear_bit(2048);
773         b.clear_bit_atomic(2048);
774         assert_eq!(None, b.next_bit(2048));
775         assert_eq!(None, b.next_zero_bit(2048));
776         assert_eq!(None, b.last_bit());
777         Ok(())
778     }
779 
780     // TODO: uncomment once kunit supports `#[should_panic]`.
781     // #[cfg(CONFIG_RUST_BITMAP_HARDENED)]
782     // #[test]
783     // #[should_panic]
784     // fn owned_bitmap_out_of_bounds() -> Result<(), AllocError> {
785     //     let mut b = BitmapVec::new(128, GFP_KERNEL)?;
786     //
787     //     b.set_bit(2048);
788     // }
789 
790     #[test]
791     fn bitmap_copy_and_extend() -> Result<(), AllocError> {
792         let mut long_bitmap = BitmapVec::new(256, GFP_KERNEL)?;
793 
794         long_bitmap.set_bit(3);
795         long_bitmap.set_bit(200);
796 
797         let mut short_bitmap = BitmapVec::new(32, GFP_KERNEL)?;
798 
799         short_bitmap.set_bit(17);
800 
801         long_bitmap.copy_and_extend(&short_bitmap);
802 
803         // Previous bits have been cleared.
804         assert_eq!(Some(17), long_bitmap.next_bit(0));
805         assert_eq!(Some(17), long_bitmap.last_bit());
806         Ok(())
807     }
808 
809     #[test]
810     fn bitmap_area_set_clear_find() -> Result<(), AllocError> {
811         let mut b = BitmapVec::new(128, GFP_KERNEL)?;
812 
813         assert_eq!(Some(0), b.next_zero_area(0, cv!(5), Alignment::SZ_1));
814         b.set(0, cv!(5)); // Now contains {[0, 5)}.
815 
816         assert_eq!(Some(0), b.next_bit(0));
817         assert_eq!(Some(4), b.next_bit(4));
818         assert_eq!(Some(5), b.next_zero_bit(0));
819         assert_eq!(Some(5), b.next_zero_area(0, cv!(5), Alignment::SZ_1));
820         assert_eq!(Some(8), b.next_zero_area(0, cv!(5), Alignment::SZ_8));
821 
822         b.set(8, cv!(8)); // Now contains {[0, 5), [8, 16)}.
823         assert_eq!(Some(16), b.next_zero_area(0, cv!(4), Alignment::SZ_16));
824         assert_eq!(Some(16), b.next_zero_area(0, cv!(4), Alignment::SZ_1));
825 
826         b.clear(0, cv!(5)); // Now contains {[8, 16)}.
827         assert_eq!(Some(0), b.next_zero_area(0, cv!(5), Alignment::SZ_1));
828         assert_eq!(Some(8), b.next_bit(0));
829         assert_eq!(Some(15), b.last_bit());
830 
831         b.set(60, cv!(10)); // Now contains {[8, 16), [60, 70)}.
832         assert_eq!(Some(60), b.next_bit(16));
833         assert_eq!(Some(69), b.last_bit());
834         assert_eq!(Some(16), b.next_zero_area(9, cv!(40), Alignment::SZ_1));
835         assert_eq!(Some(70), b.next_zero_area(0, cv!(45), Alignment::SZ_1));
836 
837         b.clear(62, cv!(6)); // Now contains {[8, 16), [60, 62), [68, 70)}.
838         assert_eq!(Some(62), b.next_zero_area(60, cv!(6), Alignment::SZ_1));
839         assert_eq!(Some(61), b.next_bit(61));
840         assert_eq!(Some(69), b.last_bit());
841         Ok(())
842     }
843 
844     #[test]
845     fn bitmap_area_exhaustion() -> Result<(), AllocError> {
846         let mut b = BitmapVec::new(64, GFP_KERNEL)?;
847 
848         assert_eq!(None, b.next_zero_area(0, cv!(65), Alignment::SZ_1));
849         assert_eq!(None, b.next_zero_area(0, cv!(usize::MAX), Alignment::SZ_1));
850         assert_eq!(None, b.next_zero_area(1, cv!(usize::MAX), Alignment::SZ_1));
851 
852         b.set_bit(0); // Now contains {[0, 1)}.
853         assert_eq!(None, b.next_zero_area(0, cv!(usize::MAX), Alignment::SZ_1));
854 
855         b.set(0, cv!(61)); // Now contains {[0, 61)}.
856         assert_eq!(None, b.next_zero_area(0, cv!(4), Alignment::SZ_1));
857         assert_eq!(Some(61), b.next_zero_area(0, cv!(3), Alignment::SZ_1));
858         assert_eq!(None, b.next_zero_area(0, cv!(1), Alignment::SZ_64));
859         Ok(())
860     }
861 
862     #[test]
863     fn bitmap_area_off() -> Result<(), AllocError> {
864         let mut b = BitmapVec::new(64, GFP_KERNEL)?;
865 
866         b.set(0, cv!(5)); // Now contains {[0, 5)}.
867 
868         // The area plus align_offset starts at a multiple of the alignment.
869         assert_eq!(Some(7), b.next_zero_area_off(0, cv!(8), Alignment::SZ_8, 1));
870         assert_eq!(Some(5), b.next_zero_area_off(0, cv!(8), Alignment::SZ_8, 3));
871 
872         // A zero offset behaves like next_zero_area().
873         assert_eq!(
874             b.next_zero_area(0, cv!(8), Alignment::SZ_8),
875             b.next_zero_area_off(0, cv!(8), Alignment::SZ_8, 0)
876         );
877         Ok(())
878     }
879 
880     #[test]
881     #[cfg(not(CONFIG_RUST_BITMAP_HARDENED))]
882     fn owned_bitmap_area_out_of_bounds() -> Result<(), AllocError> {
883         let mut b = BitmapVec::new(64, GFP_KERNEL)?;
884 
885         // Should be ignored since out of bounds.
886         b.set(64, cv!(4));
887         b.set(62, cv!(8));
888         b.set(usize::MAX, cv!(1));
889         b.clear(usize::MAX, cv!(1));
890         b.clear(2048, cv!(8));
891         assert_eq!(None, b.next_bit(0));
892         assert_eq!(None, b.next_zero_area(64, cv!(1), Alignment::SZ_1));
893         Ok(())
894     }
895 }
896