projects/linux·blob v6.18 · refs

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