rust/kernel/id_pool.rs - pub/scm/fs/xfs/xfs-linux - Git at Google

 // 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()
	}
	}