// SPDX-License-Identifier: GPL-2.0 // Copyright (C) 2025 Google LLC. //! Rust API for an ID pool backed by a [`BitmapVec`]. use crate::alloc::{AllocError, Flags};use crate::bitmap::BitmapVec; /// Represents a dynamic ID pool backed by a [`BitmapVec`].////// Clients acquire and release IDs from unset bits in a bitmap.////// The capacity of the ID pool may be adjusted by users as/// needed. The API supports the scenario where users need precise control/// over the time of allocation of a new backing bitmap, which may require/// release of spinlock./// Due to concurrent updates, all operations are re-verified to determine/// if the grow or shrink is sill valid.////// # Examples////// Basic usage////// ```/// use kernel::alloc::AllocError;/// use kernel::id_pool::{IdPool, UnusedId};////// let mut pool = IdPool::with_capacity(64, GFP_KERNEL)?;/// for i in 0..64 {/// assert_eq!(i, pool.find_unused_id(i).ok_or(ENOSPC)?.acquire());/// }////// pool.release_id(23);/// assert_eq!(23, pool.find_unused_id(0).ok_or(ENOSPC)?.acquire());////// assert!(pool.find_unused_id(0).is_none()); // time to realloc./// let resizer = pool.grow_request().ok_or(ENOSPC)?.realloc(GFP_KERNEL)?;/// pool.grow(resizer);////// assert_eq!(pool.find_unused_id(0).ok_or(ENOSPC)?.acquire(), 64);/// # Ok::<(), Error>(())/// ```////// Releasing spinlock to grow the pool////// ```no_run/// use kernel::alloc::{AllocError, flags::GFP_KERNEL};/// use kernel::sync::{new_spinlock, SpinLock};/// use kernel::id_pool::IdPool;////// fn get_id_maybe_realloc(guarded_pool: &SpinLock<IdPool>) -> Result<usize, AllocError> {/// let mut pool = guarded_pool.lock();/// loop {/// match pool.find_unused_id(0) {/// Some(index) => return Ok(index.acquire()),/// None => {/// let alloc_request = pool.grow_request();/// drop(pool);/// let resizer = alloc_request.ok_or(AllocError)?.realloc(GFP_KERNEL)?;/// pool = guarded_pool.lock();/// pool.grow(resizer)/// }/// }/// }/// }/// ```pub struct IdPool { map: BitmapVec,} /// Indicates that an [`IdPool`] should change to a new target size.pub struct ReallocRequest { num_ids: usize,} /// Contains a [`BitmapVec`] of a size suitable for reallocating [`IdPool`].pub struct PoolResizer { new: BitmapVec,} impl ReallocRequest { /// Allocates a new backing [`BitmapVec`] for [`IdPool`]. /// /// This method only prepares reallocation and does not complete it. /// Reallocation will complete after passing the [`PoolResizer`] to the /// [`IdPool::grow`] or [`IdPool::shrink`] operation, which will check /// that reallocation still makes sense. pub fn realloc(&self, flags: Flags) -> Result<PoolResizer, AllocError> { let new = BitmapVec::new(self.num_ids, flags)?; Ok(PoolResizer { new }) }} impl IdPool { /// Constructs a new [`IdPool`]. /// /// The pool will have a capacity of [`MAX_INLINE_LEN`]. /// /// [`MAX_INLINE_LEN`]: BitmapVec::MAX_INLINE_LEN #[inline] pub fn new() -> Self { Self { map: BitmapVec::new_inline(), } } /// Constructs a new [`IdPool`] with space for a specific number of bits. /// /// A capacity below [`MAX_INLINE_LEN`] is adjusted to [`MAX_INLINE_LEN`]. /// /// [`MAX_INLINE_LEN`]: BitmapVec::MAX_INLINE_LEN #[inline] pub fn with_capacity(num_ids: usize, flags: Flags) -> Result<Self, AllocError> { let num_ids = usize::max(num_ids, BitmapVec::MAX_INLINE_LEN); let map = BitmapVec::new(num_ids, flags)?; Ok(Self { map }) } /// Returns how many IDs this pool can currently have. #[inline] pub fn capacity(&self) -> usize { self.map.len() } /// Returns a [`ReallocRequest`] if the [`IdPool`] can be shrunk, [`None`] otherwise. /// /// The capacity of an [`IdPool`] cannot be shrunk below [`MAX_INLINE_LEN`]. /// /// [`MAX_INLINE_LEN`]: BitmapVec::MAX_INLINE_LEN /// /// # Examples /// /// ``` /// use kernel::{ /// alloc::AllocError, /// bitmap::BitmapVec, /// id_pool::{ /// IdPool, /// ReallocRequest, /// }, /// }; /// /// let mut pool = IdPool::with_capacity(1024, GFP_KERNEL)?; /// let alloc_request = pool.shrink_request().ok_or(AllocError)?; /// let resizer = alloc_request.realloc(GFP_KERNEL)?; /// pool.shrink(resizer); /// assert_eq!(pool.capacity(), BitmapVec::MAX_INLINE_LEN); /// # Ok::<(), AllocError>(()) /// ``` #[inline] pub fn shrink_request(&self) -> Option<ReallocRequest> { let cap = self.capacity(); // Shrinking below `MAX_INLINE_LEN` is never possible. if cap <= BitmapVec::MAX_INLINE_LEN { return None; } // Determine if the bitmap can shrink based on the position of // its last set bit. If the bit is within the first quarter of // the bitmap then shrinking is possible. In this case, the // bitmap should shrink to half its current size. let Some(bit) = self.map.last_bit() else { return Some(ReallocRequest { num_ids: BitmapVec::MAX_INLINE_LEN, }); }; if bit >= (cap / 4) { return None; } let num_ids = usize::max(BitmapVec::MAX_INLINE_LEN, cap / 2); Some(ReallocRequest { num_ids }) } /// Shrinks pool by using a new [`BitmapVec`], if still possible. #[inline] pub fn shrink(&mut self, mut resizer: PoolResizer) { // Between request to shrink that led to allocation of `resizer` and now, // bits may have changed. // Verify that shrinking is still possible. In case shrinking to // the size of `resizer` is no longer possible, do nothing, // drop `resizer` and move on. let Some(updated) = self.shrink_request() else { return; }; if updated.num_ids > resizer.new.len() { return; } resizer.new.copy_and_extend(&self.map); self.map = resizer.new; } /// Returns a [`ReallocRequest`] for growing this [`IdPool`], if possible. /// /// The capacity of an [`IdPool`] cannot be grown above [`MAX_LEN`]. /// /// [`MAX_LEN`]: BitmapVec::MAX_LEN #[inline] pub fn grow_request(&self) -> Option<ReallocRequest> { let num_ids = self.capacity() * 2; if num_ids > BitmapVec::MAX_LEN { return None; } Some(ReallocRequest { num_ids }) } /// Grows pool by using a new [`BitmapVec`], if still necessary. /// /// The `resizer` arguments has to be obtained by calling [`Self::grow_request`] /// on this object and performing a [`ReallocRequest::realloc`]. #[inline] pub fn grow(&mut self, mut resizer: PoolResizer) { // Between request to grow that led to allocation of `resizer` and now, // another thread may have already grown the capacity. // In this case, do nothing, drop `resizer` and move on. if resizer.new.len() <= self.capacity() { return; } resizer.new.copy_and_extend(&self.map); self.map = resizer.new; } /// Finds an unused ID in the bitmap. /// /// Upon success, returns its index. Otherwise, returns [`None`] /// to indicate that a [`Self::grow_request`] is needed. #[inline] #[must_use] pub fn find_unused_id(&mut self, offset: usize) -> Option<UnusedId<'_>> { // INVARIANT: `next_zero_bit()` returns None or an integer less than `map.len()` Some(UnusedId { id: self.map.next_zero_bit(offset)?, pool: self, }) } /// Releases an ID. #[inline] pub fn release_id(&mut self, id: usize) { self.map.clear_bit(id); }} /// Represents an unused id in an [`IdPool`].////// # Invariants////// The value of `id` is less than `pool.map.len()`.pub struct UnusedId<'pool> { id: usize, pool: &'pool mut IdPool,} impl<'pool> UnusedId<'pool> { /// Get the unused id as an usize. /// /// Be aware that the id has not yet been acquired in the pool. The /// [`acquire`] method must be called to prevent others from taking the id. /// /// [`acquire`]: UnusedId::acquire() #[inline] #[must_use] pub fn as_usize(&self) -> usize { self.id } /// Get the unused id as an u32. /// /// Be aware that the id has not yet been acquired in the pool. The /// [`acquire`] method must be called to prevent others from taking the id. /// /// [`acquire`]: UnusedId::acquire() #[inline] #[must_use] pub fn as_u32(&self) -> u32 { // CAST: By the type invariants: // `self.id < pool.map.len() <= BitmapVec::MAX_LEN = i32::MAX`. self.id as u32 } /// Acquire the unused id. #[inline] pub fn acquire(self) -> usize { self.pool.map.set_bit(self.id); self.id }} impl Default for IdPool { #[inline] fn default() -> Self { Self::new() }}
// SPDX-License-Identifier: GPL-2.0
// Copyright (C) 2025 Google LLC.
//! Rust API for an ID pool backed by a [`BitmapVec`].
use crate::alloc::{AllocError, Flags};
use crate::bitmap::BitmapVec;
/// Represents a dynamic ID pool backed by a [`BitmapVec`].
///
/// Clients acquire and release IDs from unset bits in a bitmap.
///
/// The capacity of the ID pool may be adjusted by users as
/// needed. The API supports the scenario where users need precise control
/// over the time of allocation of a new backing bitmap, which may require
/// release of spinlock.
/// Due to concurrent updates, all operations are re-verified to determine
/// if the grow or shrink is sill valid.
///
/// # Examples
///
/// Basic usage
///
/// ```
/// use kernel::alloc::AllocError;
/// use kernel::id_pool::{IdPool, UnusedId};
///
/// let mut pool = IdPool::with_capacity(64, GFP_KERNEL)?;
/// for i in 0..64 {
/// assert_eq!(i, pool.find_unused_id(i).ok_or(ENOSPC)?.acquire());
/// }
///
/// pool.release_id(23);
/// assert_eq!(23, pool.find_unused_id(0).ok_or(ENOSPC)?.acquire());
///
/// assert!(pool.find_unused_id(0).is_none()); // time to realloc.
/// let resizer = pool.grow_request().ok_or(ENOSPC)?.realloc(GFP_KERNEL)?;
/// pool.grow(resizer);
///
/// assert_eq!(pool.find_unused_id(0).ok_or(ENOSPC)?.acquire(), 64);
/// # Ok::<(), Error>(())
/// ```
///
/// Releasing spinlock to grow the pool
///
/// ```no_run
/// use kernel::alloc::{AllocError, flags::GFP_KERNEL};
/// use kernel::sync::{new_spinlock, SpinLock};
/// use kernel::id_pool::IdPool;
///
/// fn get_id_maybe_realloc(guarded_pool: &SpinLock<IdPool>) -> Result<usize, AllocError> {
/// let mut pool = guarded_pool.lock();
/// loop {
/// match pool.find_unused_id(0) {
/// Some(index) => return Ok(index.acquire()),
/// None => {
/// let alloc_request = pool.grow_request();
/// drop(pool);
/// let resizer = alloc_request.ok_or(AllocError)?.realloc(GFP_KERNEL)?;
/// pool = guarded_pool.lock();
/// pool.grow(resizer)
/// }
/// }
/// }
/// }
/// ```
pub struct IdPool {
map: BitmapVec,
}
/// Indicates that an [`IdPool`] should change to a new target size.
pub struct ReallocRequest {
num_ids: usize,
}
/// Contains a [`BitmapVec`] of a size suitable for reallocating [`IdPool`].
pub struct PoolResizer {
new: BitmapVec,
}
impl ReallocRequest {
/// Allocates a new backing [`BitmapVec`] for [`IdPool`].
///
/// This method only prepares reallocation and does not complete it.
/// Reallocation will complete after passing the [`PoolResizer`] to the
/// [`IdPool::grow`] or [`IdPool::shrink`] operation, which will check
/// that reallocation still makes sense.
pub fn realloc(&self, flags: Flags) -> Result<PoolResizer, AllocError> {
let new = BitmapVec::new(self.num_ids, flags)?;
Ok(PoolResizer { new })
}
}
impl IdPool {
/// Constructs a new [`IdPool`].
///
/// The pool will have a capacity of [`MAX_INLINE_LEN`].
///
/// [`MAX_INLINE_LEN`]: BitmapVec::MAX_INLINE_LEN
#[inline]
pub fn new() -> Self {
Self {
map: BitmapVec::new_inline(),
}
}
/// Constructs a new [`IdPool`] with space for a specific number of bits.
///
/// A capacity below [`MAX_INLINE_LEN`] is adjusted to [`MAX_INLINE_LEN`].
///
/// [`MAX_INLINE_LEN`]: BitmapVec::MAX_INLINE_LEN
#[inline]
pub fn with_capacity(num_ids: usize, flags: Flags) -> Result<Self, AllocError> {
let num_ids = usize::max(num_ids, BitmapVec::MAX_INLINE_LEN);
let map = BitmapVec::new(num_ids, flags)?;
Ok(Self { map })
}
/// Returns how many IDs this pool can currently have.
#[inline]
pub fn capacity(&self) -> usize {
self.map.len()
}
/// Returns a [`ReallocRequest`] if the [`IdPool`] can be shrunk, [`None`] otherwise.
///
/// The capacity of an [`IdPool`] cannot be shrunk below [`MAX_INLINE_LEN`].
///
/// [`MAX_INLINE_LEN`]: BitmapVec::MAX_INLINE_LEN
///
/// # Examples
///
/// ```
/// use kernel::{
/// alloc::AllocError,
/// bitmap::BitmapVec,
/// id_pool::{
/// IdPool,
/// ReallocRequest,
/// },
/// };
///
/// let mut pool = IdPool::with_capacity(1024, GFP_KERNEL)?;
/// let alloc_request = pool.shrink_request().ok_or(AllocError)?;
/// let resizer = alloc_request.realloc(GFP_KERNEL)?;
/// pool.shrink(resizer);
/// assert_eq!(pool.capacity(), BitmapVec::MAX_INLINE_LEN);
/// # Ok::<(), AllocError>(())
/// ```
#[inline]
pub fn shrink_request(&self) -> Option<ReallocRequest> {
let cap = self.capacity();
// Shrinking below `MAX_INLINE_LEN` is never possible.
if cap <= BitmapVec::MAX_INLINE_LEN {
return None;
}
// Determine if the bitmap can shrink based on the position of
// its last set bit. If the bit is within the first quarter of
// the bitmap then shrinking is possible. In this case, the
// bitmap should shrink to half its current size.
let Some(bit) = self.map.last_bit() else {
return Some(ReallocRequest {
num_ids: BitmapVec::MAX_INLINE_LEN,
});
};
if bit >= (cap / 4) {
return None;
}
let num_ids = usize::max(BitmapVec::MAX_INLINE_LEN, cap / 2);
Some(ReallocRequest { num_ids })
}
/// Shrinks pool by using a new [`BitmapVec`], if still possible.
#[inline]
pub fn shrink(&mut self, mut resizer: PoolResizer) {
// Between request to shrink that led to allocation of `resizer` and now,
// bits may have changed.
// Verify that shrinking is still possible. In case shrinking to
// the size of `resizer` is no longer possible, do nothing,
// drop `resizer` and move on.
let Some(updated) = self.shrink_request() else {
return;
};
if updated.num_ids > resizer.new.len() {
return;
}
resizer.new.copy_and_extend(&self.map);
self.map = resizer.new;
}
/// Returns a [`ReallocRequest`] for growing this [`IdPool`], if possible.
///
/// The capacity of an [`IdPool`] cannot be grown above [`MAX_LEN`].
///
/// [`MAX_LEN`]: BitmapVec::MAX_LEN
#[inline]
pub fn grow_request(&self) -> Option<ReallocRequest> {
let num_ids = self.capacity() * 2;
if num_ids > BitmapVec::MAX_LEN {
return None;
}
Some(ReallocRequest { num_ids })
}
/// Grows pool by using a new [`BitmapVec`], if still necessary.
///
/// The `resizer` arguments has to be obtained by calling [`Self::grow_request`]
/// on this object and performing a [`ReallocRequest::realloc`].
#[inline]
pub fn grow(&mut self, mut resizer: PoolResizer) {
// Between request to grow that led to allocation of `resizer` and now,
// another thread may have already grown the capacity.
// In this case, do nothing, drop `resizer` and move on.
if resizer.new.len() <= self.capacity() {
return;
}
resizer.new.copy_and_extend(&self.map);
self.map = resizer.new;
}
/// Finds an unused ID in the bitmap.
///
/// Upon success, returns its index. Otherwise, returns [`None`]
/// to indicate that a [`Self::grow_request`] is needed.
#[inline]
#[must_use]
pub fn find_unused_id(&mut self, offset: usize) -> Option<UnusedId<'_>> {
// INVARIANT: `next_zero_bit()` returns None or an integer less than `map.len()`
Some(UnusedId {
id: self.map.next_zero_bit(offset)?,
pool: self,
})
}
/// Releases an ID.
#[inline]
pub fn release_id(&mut self, id: usize) {
self.map.clear_bit(id);
}
}
/// Represents an unused id in an [`IdPool`].
///
/// # Invariants
///
/// The value of `id` is less than `pool.map.len()`.
pub struct UnusedId<'pool> {
id: usize,
pool: &'pool mut IdPool,
}
impl<'pool> UnusedId<'pool> {
/// Get the unused id as an usize.
///
/// Be aware that the id has not yet been acquired in the pool. The
/// [`acquire`] method must be called to prevent others from taking the id.
///
/// [`acquire`]: UnusedId::acquire()
#[inline]
#[must_use]
pub fn as_usize(&self) -> usize {
self.id
}
/// Get the unused id as an u32.
///
/// Be aware that the id has not yet been acquired in the pool. The
/// [`acquire`] method must be called to prevent others from taking the id.
///
/// [`acquire`]: UnusedId::acquire()
#[inline]
#[must_use]
pub fn as_u32(&self) -> u32 {
// CAST: By the type invariants:
// `self.id < pool.map.len() <= BitmapVec::MAX_LEN = i32::MAX`.
self.id as u32
}
/// Acquire the unused id.
#[inline]
pub fn acquire(self) -> usize {
self.pool.map.set_bit(self.id);
self.id
}
}
impl Default for IdPool {
#[inline]
fn default() -> Self {
Self::new()
}
}