projects/linux·blob v7.1 · refs

621 lines · 20.2 KiB

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 core::ptr::NonNull;
14
15/// Represents a C bitmap. Wraps underlying C bitmap API.
16///
17/// # Invariants
18///
19/// Must reference a `[c_ulong]` long enough to fit `data.len()` bits.
20#[cfg_attr(CONFIG_64BIT, repr(align(8)))]
21#[cfg_attr(not(CONFIG_64BIT), repr(align(4)))]
22pub struct Bitmap {
23 data: [()],
24}
25
26impl Bitmap {
27 /// Borrows a C bitmap.
28 ///
29 /// # Safety
30 ///
31 /// * `ptr` holds a non-null address of an initialized array of `unsigned long`
32 /// that is large enough to hold `nbits` bits.
33 /// * the array must not be freed for the lifetime of this [`Bitmap`]
34 /// * concurrent access only happens through atomic operations
35 pub unsafe fn from_raw<'a>(ptr: *const usize, nbits: usize) -> &'a Bitmap {
36 let data: *const [()] = core::ptr::slice_from_raw_parts(ptr.cast(), nbits);
37 // INVARIANT: `data` references an initialized array that can hold `nbits` bits.
38 // SAFETY:
39 // The caller guarantees that `data` (derived from `ptr` and `nbits`)
40 // points to a valid, initialized, and appropriately sized memory region
41 // that will not be freed for the lifetime 'a.
42 // We are casting `*const [()]` to `*const Bitmap`. The `Bitmap`
43 // struct is a ZST with a `data: [()]` field. This means its layout
44 // is compatible with a slice of `()`, and effectively it's a "thin pointer"
45 // (its size is 0 and alignment is 1). The `slice_from_raw_parts`
46 // function correctly encodes the length (number of bits, not elements)
47 // into the metadata of the fat pointer. Therefore, dereferencing this
48 // pointer as `&Bitmap` is safe given the caller's guarantees.
49 unsafe { &*(data as *const Bitmap) }
50 }
51
52 /// Borrows a C bitmap exclusively.
53 ///
54 /// # Safety
55 ///
56 /// * `ptr` holds a non-null address of an initialized array of `unsigned long`
57 /// that is large enough to hold `nbits` bits.
58 /// * the array must not be freed for the lifetime of this [`Bitmap`]
59 /// * no concurrent access may happen.
60 pub unsafe fn from_raw_mut<'a>(ptr: *mut usize, nbits: usize) -> &'a mut Bitmap {
61 let data: *mut [()] = core::ptr::slice_from_raw_parts_mut(ptr.cast(), nbits);
62 // INVARIANT: `data` references an initialized array that can hold `nbits` bits.
63 // SAFETY:
64 // The caller guarantees that `data` (derived from `ptr` and `nbits`)
65 // points to a valid, initialized, and appropriately sized memory region
66 // that will not be freed for the lifetime 'a.
67 // Furthermore, the caller guarantees no concurrent access will happen,
68 // which upholds the exclusivity requirement for a mutable reference.
69 // Similar to `from_raw`, casting `*mut [()]` to `*mut Bitmap` is
70 // safe because `Bitmap` is a ZST with a `data: [()]` field,
71 // making its layout compatible with a slice of `()`.
72 unsafe { &mut *(data as *mut Bitmap) }
73 }
74
75 /// Returns a raw pointer to the backing [`Bitmap`].
76 pub fn as_ptr(&self) -> *const usize {
77 core::ptr::from_ref::<Bitmap>(self).cast::<usize>()
78 }
79
80 /// Returns a mutable raw pointer to the backing [`Bitmap`].
81 pub fn as_mut_ptr(&mut self) -> *mut usize {
82 core::ptr::from_mut::<Bitmap>(self).cast::<usize>()
83 }
84
85 /// Returns length of this [`Bitmap`].
86 #[expect(clippy::len_without_is_empty)]
87 pub fn len(&self) -> usize {
88 self.data.len()
89 }
90}
91
92/// Holds either a pointer to array of `unsigned long` or a small bitmap.
93#[repr(C)]
94union BitmapRepr {
95 bitmap: usize,
96 ptr: NonNull<usize>,
97}
98
99macro_rules! bitmap_assert {
100 ($cond:expr, $($arg:tt)+) => {
101 #[cfg(CONFIG_RUST_BITMAP_HARDENED)]
102 assert!($cond, $($arg)*);
103 }
104}
105
106macro_rules! bitmap_assert_return {
107 ($cond:expr, $($arg:tt)+) => {
108 #[cfg(CONFIG_RUST_BITMAP_HARDENED)]
109 assert!($cond, $($arg)*);
110
111 #[cfg(not(CONFIG_RUST_BITMAP_HARDENED))]
112 if !($cond) {
113 pr_err!($($arg)*);
114 return
115 }
116 }
117}
118
119/// Represents an owned bitmap.
120///
121/// Wraps underlying C bitmap API. See [`Bitmap`] for available
122/// methods.
123///
124/// # Examples
125///
126/// Basic usage
127///
128/// ```
129/// use kernel::alloc::flags::GFP_KERNEL;
130/// use kernel::bitmap::BitmapVec;
131///
132/// let mut b = BitmapVec::new(16, GFP_KERNEL)?;
133///
134/// assert_eq!(16, b.len());
135/// for i in 0..16 {
136/// if i % 4 == 0 {
137/// b.set_bit(i);
138/// }
139/// }
140/// assert_eq!(Some(0), b.next_bit(0));
141/// assert_eq!(Some(1), b.next_zero_bit(0));
142/// assert_eq!(Some(4), b.next_bit(1));
143/// assert_eq!(Some(5), b.next_zero_bit(4));
144/// assert_eq!(Some(12), b.last_bit());
145/// # Ok::<(), Error>(())
146/// ```
147///
148/// # Invariants
149///
150/// * `nbits` is `<= MAX_LEN`.
151/// * if `nbits <= MAX_INLINE_LEN`, then `repr` is a `usize`.
152/// * otherwise, `repr` holds a non-null pointer to an initialized
153/// array of `unsigned long` that is large enough to hold `nbits` bits.
154pub struct BitmapVec {
155 /// Representation of bitmap.
156 repr: BitmapRepr,
157 /// Length of this bitmap. Must be `<= MAX_LEN`.
158 nbits: usize,
159}
160
161impl core::ops::Deref for BitmapVec {
162 type Target = Bitmap;
163
164 fn deref(&self) -> &Bitmap {
165 let ptr = if self.nbits <= BitmapVec::MAX_INLINE_LEN {
166 // SAFETY: Bitmap is represented inline.
167 #[allow(unused_unsafe, reason = "Safe since Rust 1.92.0")]
168 unsafe {
169 core::ptr::addr_of!(self.repr.bitmap)
170 }
171 } else {
172 // SAFETY: Bitmap is represented as array of `unsigned long`.
173 unsafe { self.repr.ptr.as_ptr() }
174 };
175
176 // SAFETY: We got the right pointer and invariants of [`Bitmap`] hold.
177 // An inline bitmap is treated like an array with single element.
178 unsafe { Bitmap::from_raw(ptr, self.nbits) }
179 }
180}
181
182impl core::ops::DerefMut for BitmapVec {
183 fn deref_mut(&mut self) -> &mut Bitmap {
184 let ptr = if self.nbits <= BitmapVec::MAX_INLINE_LEN {
185 // SAFETY: Bitmap is represented inline.
186 #[allow(unused_unsafe, reason = "Safe since Rust 1.92.0")]
187 unsafe {
188 core::ptr::addr_of_mut!(self.repr.bitmap)
189 }
190 } else {
191 // SAFETY: Bitmap is represented as array of `unsigned long`.
192 unsafe { self.repr.ptr.as_ptr() }
193 };
194
195 // SAFETY: We got the right pointer and invariants of [`BitmapVec`] hold.
196 // An inline bitmap is treated like an array with single element.
197 unsafe { Bitmap::from_raw_mut(ptr, self.nbits) }
198 }
199}
200
201/// Enable ownership transfer to other threads.
202///
203/// SAFETY: We own the underlying bitmap representation.
204unsafe impl Send for BitmapVec {}
205
206/// Enable unsynchronized concurrent access to [`BitmapVec`] through shared references.
207///
208/// SAFETY: `deref()` will return a reference to a [`Bitmap`]. Its methods
209/// take immutable references are either atomic or read-only.
210unsafe impl Sync for BitmapVec {}
211
212impl Drop for BitmapVec {
213 fn drop(&mut self) {
214 if self.nbits <= BitmapVec::MAX_INLINE_LEN {
215 return;
216 }
217 // SAFETY: `self.ptr` was returned by the C `bitmap_zalloc`.
218 //
219 // INVARIANT: there is no other use of the `self.ptr` after this
220 // call and the value is being dropped so the broken invariant is
221 // not observable on function exit.
222 unsafe { bindings::bitmap_free(self.repr.ptr.as_ptr()) };
223 }
224}
225
226impl BitmapVec {
227 /// The maximum possible length of a `BitmapVec`.
228 pub const MAX_LEN: usize = i32::MAX as usize;
229
230 /// The maximum length that uses the inline representation.
231 pub const MAX_INLINE_LEN: usize = usize::BITS as usize;
232
233 /// Construct a longest possible inline [`BitmapVec`].
234 #[inline]
235 pub fn new_inline() -> Self {
236 // INVARIANT: `nbits <= MAX_INLINE_LEN`, so an inline bitmap is the right repr.
237 BitmapVec {
238 repr: BitmapRepr { bitmap: 0 },
239 nbits: BitmapVec::MAX_INLINE_LEN,
240 }
241 }
242
243 /// Constructs a new [`BitmapVec`].
244 ///
245 /// Fails with [`AllocError`] when the [`BitmapVec`] could not be allocated. This
246 /// includes the case when `nbits` is greater than `MAX_LEN`.
247 #[inline]
248 pub fn new(nbits: usize, flags: Flags) -> Result<Self, AllocError> {
249 if nbits <= BitmapVec::MAX_INLINE_LEN {
250 return Ok(BitmapVec {
251 repr: BitmapRepr { bitmap: 0 },
252 nbits,
253 });
254 }
255 if nbits > Self::MAX_LEN {
256 return Err(AllocError);
257 }
258 let nbits_u32 = u32::try_from(nbits).unwrap();
259 // SAFETY: `MAX_INLINE_LEN < nbits` and `nbits <= MAX_LEN`.
260 let ptr = unsafe { bindings::bitmap_zalloc(nbits_u32, flags.as_raw()) };
261 let ptr = NonNull::new(ptr).ok_or(AllocError)?;
262 // INVARIANT: `ptr` returned by C `bitmap_zalloc` and `nbits` checked.
263 Ok(BitmapVec {
264 repr: BitmapRepr { ptr },
265 nbits,
266 })
267 }
268
269 /// Returns length of this [`Bitmap`].
270 #[allow(clippy::len_without_is_empty)]
271 #[inline]
272 pub fn len(&self) -> usize {
273 self.nbits
274 }
275
276 /// Fills this `Bitmap` with random bits.
277 #[cfg(CONFIG_FIND_BIT_BENCHMARK_RUST)]
278 pub fn fill_random(&mut self) {
279 // SAFETY: `self.as_mut_ptr` points to either an array of the
280 // appropriate length or one usize.
281 unsafe {
282 bindings::get_random_bytes(
283 self.as_mut_ptr().cast::<ffi::c_void>(),
284 usize::div_ceil(self.nbits, bindings::BITS_PER_LONG as usize)
285 * bindings::BITS_PER_LONG as usize
286 / 8,
287 );
288 }
289 }
290}
291
292impl Bitmap {
293 /// Set bit with index `index`.
294 ///
295 /// ATTENTION: `set_bit` is non-atomic, which differs from the naming
296 /// convention in C code. The corresponding C function is `__set_bit`.
297 ///
298 /// If CONFIG_RUST_BITMAP_HARDENED is not enabled and `index` is greater than
299 /// or equal to `self.nbits`, does nothing.
300 ///
301 /// # Panics
302 ///
303 /// Panics if CONFIG_RUST_BITMAP_HARDENED is enabled and `index` is greater than
304 /// or equal to `self.nbits`.
305 #[inline]
306 pub fn set_bit(&mut self, index: usize) {
307 bitmap_assert_return!(
308 index < self.len(),
309 "Bit `index` must be < {}, was {}",
310 self.len(),
311 index
312 );
313 // SAFETY: Bit `index` is within bounds.
314 unsafe { bindings::__set_bit(index, self.as_mut_ptr()) };
315 }
316
317 /// Set bit with index `index`, atomically.
318 ///
319 /// This is a relaxed atomic operation (no implied memory barriers).
320 ///
321 /// ATTENTION: The naming convention differs from C, where the corresponding
322 /// function is called `set_bit`.
323 ///
324 /// If CONFIG_RUST_BITMAP_HARDENED is not enabled and `index` is greater than
325 /// or equal to `self.len()`, does nothing.
326 ///
327 /// # Panics
328 ///
329 /// Panics if CONFIG_RUST_BITMAP_HARDENED is enabled and `index` is greater than
330 /// or equal to `self.len()`.
331 #[inline]
332 pub fn set_bit_atomic(&self, index: usize) {
333 bitmap_assert_return!(
334 index < self.len(),
335 "Bit `index` must be < {}, was {}",
336 self.len(),
337 index
338 );
339 // SAFETY: `index` is within bounds and the caller has ensured that
340 // there is no mix of non-atomic and atomic operations.
341 unsafe { bindings::set_bit(index, self.as_ptr().cast_mut()) };
342 }
343
344 /// Clear `index` bit.
345 ///
346 /// ATTENTION: `clear_bit` is non-atomic, which differs from the naming
347 /// convention in C code. The corresponding C function is `__clear_bit`.
348 ///
349 /// If CONFIG_RUST_BITMAP_HARDENED is not enabled and `index` is greater than
350 /// or equal to `self.len()`, does nothing.
351 ///
352 /// # Panics
353 ///
354 /// Panics if CONFIG_RUST_BITMAP_HARDENED is enabled and `index` is greater than
355 /// or equal to `self.len()`.
356 #[inline]
357 pub fn clear_bit(&mut self, index: usize) {
358 bitmap_assert_return!(
359 index < self.len(),
360 "Bit `index` must be < {}, was {}",
361 self.len(),
362 index
363 );
364 // SAFETY: `index` is within bounds.
365 unsafe { bindings::__clear_bit(index, self.as_mut_ptr()) };
366 }
367
368 /// Clear `index` bit, atomically.
369 ///
370 /// This is a relaxed atomic operation (no implied memory barriers).
371 ///
372 /// ATTENTION: The naming convention differs from C, where the corresponding
373 /// function is called `clear_bit`.
374 ///
375 /// If CONFIG_RUST_BITMAP_HARDENED is not enabled and `index` is greater than
376 /// or equal to `self.len()`, does nothing.
377 ///
378 /// # Panics
379 ///
380 /// Panics if CONFIG_RUST_BITMAP_HARDENED is enabled and `index` is greater than
381 /// or equal to `self.len()`.
382 #[inline]
383 pub fn clear_bit_atomic(&self, index: usize) {
384 bitmap_assert_return!(
385 index < self.len(),
386 "Bit `index` must be < {}, was {}",
387 self.len(),
388 index
389 );
390 // SAFETY: `index` is within bounds and the caller has ensured that
391 // there is no mix of non-atomic and atomic operations.
392 unsafe { bindings::clear_bit(index, self.as_ptr().cast_mut()) };
393 }
394
395 /// Copy `src` into this [`Bitmap`] and set any remaining bits to zero.
396 ///
397 /// # Examples
398 ///
399 /// ```
400 /// use kernel::alloc::{AllocError, flags::GFP_KERNEL};
401 /// use kernel::bitmap::BitmapVec;
402 ///
403 /// let mut long_bitmap = BitmapVec::new(256, GFP_KERNEL)?;
404 ///
405 /// assert_eq!(None, long_bitmap.last_bit());
406 ///
407 /// let mut short_bitmap = BitmapVec::new(16, GFP_KERNEL)?;
408 ///
409 /// short_bitmap.set_bit(7);
410 /// long_bitmap.copy_and_extend(&short_bitmap);
411 /// assert_eq!(Some(7), long_bitmap.last_bit());
412 ///
413 /// # Ok::<(), AllocError>(())
414 /// ```
415 #[inline]
416 pub fn copy_and_extend(&mut self, src: &Bitmap) {
417 let len = core::cmp::min(src.len(), self.len());
418 // SAFETY: access to `self` and `src` is within bounds.
419 unsafe {
420 bindings::bitmap_copy_and_extend(
421 self.as_mut_ptr(),
422 src.as_ptr(),
423 len as u32,
424 self.len() as u32,
425 )
426 };
427 }
428
429 /// Finds last set bit.
430 ///
431 /// # Examples
432 ///
433 /// ```
434 /// use kernel::alloc::{AllocError, flags::GFP_KERNEL};
435 /// use kernel::bitmap::BitmapVec;
436 ///
437 /// let bitmap = BitmapVec::new(64, GFP_KERNEL)?;
438 ///
439 /// match bitmap.last_bit() {
440 /// Some(idx) => {
441 /// pr_info!("The last bit has index {idx}.\n");
442 /// }
443 /// None => {
444 /// pr_info!("All bits in this bitmap are 0.\n");
445 /// }
446 /// }
447 /// # Ok::<(), AllocError>(())
448 /// ```
449 #[inline]
450 pub fn last_bit(&self) -> Option<usize> {
451 // SAFETY: `_find_next_bit` access is within bounds due to invariant.
452 let index = unsafe { bindings::_find_last_bit(self.as_ptr(), self.len()) };
453 if index >= self.len() {
454 None
455 } else {
456 Some(index)
457 }
458 }
459
460 /// Finds next set bit, starting from `start`.
461 ///
462 /// Returns `None` if `start` is greater or equal to `self.nbits`.
463 #[inline]
464 pub fn next_bit(&self, start: usize) -> Option<usize> {
465 bitmap_assert!(
466 start < self.len(),
467 "`start` must be < {} was {}",
468 self.len(),
469 start
470 );
471 // SAFETY: `_find_next_bit` tolerates out-of-bounds arguments and returns a
472 // value larger than or equal to `self.len()` in that case.
473 let index = unsafe { bindings::_find_next_bit(self.as_ptr(), self.len(), start) };
474 if index >= self.len() {
475 None
476 } else {
477 Some(index)
478 }
479 }
480
481 /// Finds next zero bit, starting from `start`.
482 /// Returns `None` if `start` is greater than or equal to `self.len()`.
483 #[inline]
484 pub fn next_zero_bit(&self, start: usize) -> Option<usize> {
485 bitmap_assert!(
486 start < self.len(),
487 "`start` must be < {} was {}",
488 self.len(),
489 start
490 );
491 // SAFETY: `_find_next_zero_bit` tolerates out-of-bounds arguments and returns a
492 // value larger than or equal to `self.len()` in that case.
493 let index = unsafe { bindings::_find_next_zero_bit(self.as_ptr(), self.len(), start) };
494 if index >= self.len() {
495 None
496 } else {
497 Some(index)
498 }
499 }
500}
501
502use macros::kunit_tests;
503
504#[kunit_tests(rust_kernel_bitmap)]
505mod tests {
506 use super::*;
507 use kernel::alloc::flags::GFP_KERNEL;
508
509 #[test]
510 fn bitmap_borrow() {
511 let fake_bitmap: [usize; 2] = [0, 0];
512 let fake_bitmap_len = 2 * usize::BITS as usize;
513 // SAFETY: `fake_c_bitmap` is an array of expected length.
514 let b = unsafe { Bitmap::from_raw(fake_bitmap.as_ptr(), fake_bitmap_len) };
515 assert_eq!(fake_bitmap_len, b.len());
516 assert_eq!(None, b.next_bit(0));
517 }
518
519 #[test]
520 fn bitmap_copy() {
521 let fake_bitmap: usize = 0xFF;
522 // SAFETY: `fake_c_bitmap` can be used as one-element array of expected length.
523 let b = unsafe { Bitmap::from_raw(core::ptr::addr_of!(fake_bitmap), 8) };
524 assert_eq!(8, b.len());
525 assert_eq!(None, b.next_zero_bit(0));
526 }
527
528 #[test]
529 fn bitmap_vec_new() -> Result<(), AllocError> {
530 let b = BitmapVec::new(0, GFP_KERNEL)?;
531 assert_eq!(0, b.len());
532
533 let b = BitmapVec::new(3, GFP_KERNEL)?;
534 assert_eq!(3, b.len());
535
536 let b = BitmapVec::new(1024, GFP_KERNEL)?;
537 assert_eq!(1024, b.len());
538
539 // Requesting too large values results in [`AllocError`].
540 let res = BitmapVec::new(1 << 31, GFP_KERNEL);
541 assert!(res.is_err());
542 Ok(())
543 }
544
545 #[test]
546 fn bitmap_set_clear_find() -> Result<(), AllocError> {
547 let mut b = BitmapVec::new(128, GFP_KERNEL)?;
548
549 // Zero-initialized
550 assert_eq!(None, b.next_bit(0));
551 assert_eq!(Some(0), b.next_zero_bit(0));
552 assert_eq!(None, b.last_bit());
553
554 b.set_bit(17);
555
556 assert_eq!(Some(17), b.next_bit(0));
557 assert_eq!(Some(17), b.next_bit(17));
558 assert_eq!(None, b.next_bit(18));
559 assert_eq!(Some(17), b.last_bit());
560
561 b.set_bit(107);
562
563 assert_eq!(Some(17), b.next_bit(0));
564 assert_eq!(Some(17), b.next_bit(17));
565 assert_eq!(Some(107), b.next_bit(18));
566 assert_eq!(Some(107), b.last_bit());
567
568 b.clear_bit(17);
569
570 assert_eq!(Some(107), b.next_bit(0));
571 assert_eq!(Some(107), b.last_bit());
572 Ok(())
573 }
574
575 #[test]
576 fn owned_bitmap_out_of_bounds() -> Result<(), AllocError> {
577 // TODO: Kunit #[test]s do not support `cfg` yet,
578 // so we add it here in the body.
579 #[cfg(not(CONFIG_RUST_BITMAP_HARDENED))]
580 {
581 let mut b = BitmapVec::new(128, GFP_KERNEL)?;
582 b.set_bit(2048);
583 b.set_bit_atomic(2048);
584 b.clear_bit(2048);
585 b.clear_bit_atomic(2048);
586 assert_eq!(None, b.next_bit(2048));
587 assert_eq!(None, b.next_zero_bit(2048));
588 assert_eq!(None, b.last_bit());
589 }
590 Ok(())
591 }
592
593 // TODO: uncomment once kunit supports [should_panic] and `cfg`.
594 // #[cfg(CONFIG_RUST_BITMAP_HARDENED)]
595 // #[test]
596 // #[should_panic]
597 // fn owned_bitmap_out_of_bounds() -> Result<(), AllocError> {
598 // let mut b = BitmapVec::new(128, GFP_KERNEL)?;
599 //
600 // b.set_bit(2048);
601 // }
602
603 #[test]
604 fn bitmap_copy_and_extend() -> Result<(), AllocError> {
605 let mut long_bitmap = BitmapVec::new(256, GFP_KERNEL)?;
606
607 long_bitmap.set_bit(3);
608 long_bitmap.set_bit(200);
609
610 let mut short_bitmap = BitmapVec::new(32, GFP_KERNEL)?;
611
612 short_bitmap.set_bit(17);
613
614 long_bitmap.copy_and_extend(&short_bitmap);
615
616 // Previous bits have been cleared.
617 assert_eq!(Some(17), long_bitmap.next_bit(0));
618 assert_eq!(Some(17), long_bitmap.last_bit());
619 Ok(())
620 }
621}