Skip to main content

kernel/
bitmap.rs

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
9use crate::alloc::{AllocError, Flags};
10use crate::bindings;
11#[cfg(not(CONFIG_RUST_BITMAP_HARDENED))]
12use crate::pr_err;
13use crate::ptr::Alignment;
14use 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)))]
29pub struct Bitmap {
30    data: [()],
31}
32
33macro_rules! bitmap_assert {
34    ($cond:expr, $($arg:tt)+) => {
35        #[cfg(CONFIG_RUST_BITMAP_HARDENED)]
36        assert!($cond, $($arg)*);
37    }
38}
39
40macro_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
53impl 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)]
145union 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.
185pub struct BitmapVec {
186    /// Representation of bitmap.
187    repr: BitmapRepr,
188    /// Length of this bitmap. Must be `<= MAX_LEN`.
189    nbits: usize,
190}
191
192impl 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
213impl 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.
235unsafe 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.
241unsafe impl Sync for BitmapVec {}
242
243impl 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
257impl 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
323impl 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)]
691mod 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}