[PATCH v4 3/3] Add DynamicIdPool Rust API.

[Burak Emir <bqe@google.com>]

This is a port of the Binder data structure introded in commit
15d9da3f818c ("binder: use bitmap for faster descriptor lookup") to
Rust.

The file drivers/android/dbitmap.h uses bitmaps for allocating and
releasing IDs. We provide an equivalent Rust API that gives clients
fine-grained control over when to allocate a new bitmap.

Suggested-by: Alice Ryhl <aliceryhl@google.com>
Signed-off-by: Burak Emir <bqe@google.com>
---
 MAINTAINERS                    |   3 +-
 rust/kernel/dynamic_id_pool.rs | 191 +++++++++++++++++++++++++++++++++
 rust/kernel/lib.rs             |   1 +
 3 files changed, 194 insertions(+), 1 deletion(-)
 create mode 100644 rust/kernel/dynamic_id_pool.rs

diff --git a/MAINTAINERS b/MAINTAINERS
index b3bbce9274f0..d429ede24d3c 100644
--- a/MAINTAINERS
+++ b/MAINTAINERS
@@ -4036,12 +4036,13 @@ F:	rust/helpers/bitmap.c
 F:	rust/helpers/cpumask.c
 
 BITMAP API [RUST]
-M:	Alice Ryhl <aliceryhl@google.com> (bitmap)
+M:	Alice Ryhl <aliceryhl@google.com> (bitmap,dynamic_id_pool)
 M:	Viresh Kumar <viresh.kumar@linaro.org> (cpumask)
 R:	Yury Norov <yury.norov@gmail.com>
 S:	Maintained
 F:	rust/kernel/bitmap.rs
 F:	rust/kernel/cpumask.rs
+F:	rust/kernel/dynamic_id_pool.rs
 
 BITOPS API
 M:	Yury Norov <yury.norov@gmail.com>
diff --git a/rust/kernel/dynamic_id_pool.rs b/rust/kernel/dynamic_id_pool.rs
new file mode 100644
index 000000000000..3e243358cd6b
--- /dev/null
+++ b/rust/kernel/dynamic_id_pool.rs
@@ -0,0 +1,191 @@
+// SPDX-License-Identifier: GPL-2.0
+
+// Copyright (C) 2025 Google LLC.
+
+//! Rust API for an ID pool backed by a `Bitmap`.
+
+use crate::alloc::{AllocError, Flags};
+use crate::bitmap::Bitmap;
+
+/// Represents a dynamic ID pool backed by a `Bitmap`.
+///
+/// Clients acquire and release IDs from zero bits in a bitmap.
+///
+/// The ID pool can grow or shrink as needed. It has been designed
+/// to support the scenario where users need to control the time
+/// of allocation of a new backing bitmap, which may require release
+/// of locks.
+/// These operations then, are verified to determine if the grow or
+/// shrink is sill valid.
+///
+/// # Examples
+///
+/// Basic usage
+///
+/// ```
+/// use kernel::alloc::{AllocError, flags::GFP_KERNEL};
+/// use kernel::dynamic_id_pool::DynamicIdPool;
+///
+/// let mut pool = DynamicIdPool::new(64, GFP_KERNEL)?;
+/// for i in 0..64 {
+///   pool.acquire_next_id(i).ok_or(ENOSPC)?;
+/// }
+/// assert_eq!(pool.acquire_next_id(63), None);
+/// let resizer = pool.grow_alloc().alloc(GFP_KERNEL)?;
+/// pool.grow(resizer);
+/// assert_eq!(pool.acquire_next_id(63), Some(64));
+/// # Ok::<(), Error>(())
+/// ```
+///
+/// Releasing spinlock to grow the pool
+///
+/// ```
+/// use kernel::alloc::{AllocError, flags::GFP_KERNEL};
+/// use kernel::sync::{new_spinlock, SpinLock};
+/// use kernel::dynamic_id_pool::DynamicIdPool;
+///
+/// struct Example {
+///   pool: DynamicIdPool
+/// }
+///
+/// fn get_id_maybe_alloc(s: &SpinLock<Example>) -> Result<usize, AllocError> {
+///   let mut guard = s.lock();
+///   let mut resizer = None;
+///   loop {
+///     match guard.pool.acquire_next_id(0) {
+///       index @ Some(_) => return Ok(index),
+///       None => {
+///         let alloc_request = guard.pool.grow_alloc();
+///         drop(guard);
+///         let resizer = alloc_request.alloc(GFP_KERNEL)?;
+///         guard = s.lock();
+///         guard.pool.grow(resizer)
+///       }
+///     }
+///   }
+/// }
+/// ```
+pub struct DynamicIdPool {
+    map: Bitmap,
+}
+
+/// Returned when the `DynamicIdPool` should change size.
+pub struct AllocRequest {
+    nbits: usize,
+}
+
+/// Contains an allocated `Bitmap` for resizing `DynamicIdPool`.
+pub struct PoolResizer {
+    new: Bitmap,
+}
+
+impl AllocRequest {
+    /// Allocates a new `Bitmap` for `DynamicIdPool`.
+    pub fn alloc(&self, flags: Flags) -> Result<PoolResizer, AllocError> {
+        let new = Bitmap::new(self.nbits, flags)?;
+        Ok(PoolResizer { new })
+    }
+}
+
+/// Minimum size we want to use for our backing `Bitmap`.
+const NBITS_MIN: usize = bindings::BITS_PER_LONG as usize;
+
+impl DynamicIdPool {
+    ///
+    #[inline]
+    pub fn new(mut nbits: usize, flags: Flags) -> Result<Self, AllocError> {
+        if nbits < NBITS_MIN {
+            nbits = NBITS_MIN
+        }
+        let map = Bitmap::new(nbits, flags)?;
+        Ok(Self { map })
+    }
+
+    /// Returns how many IDs this pool can currently have.
+    #[inline]
+    pub fn len(&self) -> usize {
+        self.map.len()
+    }
+
+    /// Returns an `AllocRequest` if the can be shrunk, or None if not possible.
+    #[inline]
+    pub fn shrink_alloc(&self) -> Option<AllocRequest> {
+        let len = self.map.len();
+        if len == NBITS_MIN {
+            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.
+         */
+        match self.map.find_last_bit() {
+            Some(bit) => {
+                if bit < (len >> 2) {
+                    Some(AllocRequest { nbits: len >> 1 })
+                } else {
+                    None
+                }
+            }
+            None => Some(AllocRequest { nbits: NBITS_MIN }),
+        }
+    }
+
+    /// Shrinks pool by using a new `Bitmap`, if still possible.
+    #[inline]
+    pub fn shrink(&mut self, mut resizer: PoolResizer) {
+        // Verify that shrinking is still possible. The `resizer`
+        // bitmap might have been allocated without locks, so this call
+        // could now be outdated. In this case, drop `resizer` and move on.
+        if let Some(AllocRequest { nbits }) = self.shrink_alloc() {
+            if nbits <= resizer.new.len() {
+                resizer.new.copy_prefix_from_bitmap(&self.map);
+                self.map = resizer.new;
+                return;
+            }
+        }
+    }
+
+    /// Returns an `AllocRequest` for growing this `DynamicIdPool`.
+    #[inline]
+    pub fn grow_alloc(&self) -> AllocRequest {
+        AllocRequest {
+            nbits: self.map.len() << 1,
+        }
+    }
+
+    /// Grows pool by using a new `Bitmap`, if still necessary.
+    #[inline]
+    pub fn grow(&mut self, mut resizer: PoolResizer) {
+        // `resizer` bitmap might have been allocated without locks,
+        // so this call could now be outdated. In this case, drop
+        // `resizer` and move on.
+        if resizer.new.len() <= self.map.len() {
+            return;
+        }
+
+        resizer.new.copy_from_bitmap_and_extend(&self.map);
+        self.map = resizer.new;
+    }
+
+    /// Acquires a new ID by finding and setting the next zero bit in the
+    /// bitmap. Upon success, returns its index. Otherwise, returns `None`
+    /// to indicate that a `grow_alloc` is needed.
+    #[inline]
+    pub fn acquire_next_id(&mut self, offset: usize) -> Option<usize> {
+        match self.map.find_next_zero_bit(offset) {
+            res @ Some(nr) => {
+                self.map.set_bit(nr);
+                res
+            }
+            None => None,
+        }
+    }
+
+    /// Releases an ID.
+    #[inline]
+    pub fn release_id(&mut self, id: usize) {
+        self.map.clear_bit(id);
+    }
+}
diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index be06ffc47473..4789e707dacc 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -47,6 +47,7 @@
 pub mod device_id;
 pub mod devres;
 pub mod driver;
+pub mod dynamic_id_pool;
 pub mod error;
 #[cfg(CONFIG_RUST_FW_LOADER_ABSTRACTIONS)]
 pub mod firmware;
-- 
2.49.0.rc1.451.g8f38331e32-goog