1 files changed, 329 insertions, 0 deletions
diff --git a/drivers/android/binder/range_alloc/mod.rs b/drivers/android/binder/range_alloc/mod.rs
new file mode 100644
index 000000000000..2301e2bc1a1f
--- /dev/null
+++ b/drivers/android/binder/range_alloc/mod.rs
@@ -0,0 +1,329 @@
+// SPDX-License-Identifier: GPL-2.0
+
+// Copyright (C) 2025 Google LLC.
+
+use kernel::{page::PAGE_SIZE, prelude::*, seq_file::SeqFile, task::Pid};
+
+mod tree;
+use self::tree::{FromArrayAllocs, ReserveNewTreeAlloc, TreeRangeAllocator};
+
+mod array;
+use self::array::{ArrayRangeAllocator, EmptyArrayAlloc};
+
+enum DescriptorState<T> {
+    Reserved(Reservation),
+    Allocated(Allocation<T>),
+}
+
+impl<T> DescriptorState<T> {
+    fn new(is_oneway: bool, debug_id: usize, pid: Pid) -> Self {
+        DescriptorState::Reserved(Reservation {
+            debug_id,
+            is_oneway,
+            pid,
+        })
+    }
+
+    fn pid(&self) -> Pid {
+        match self {
+            DescriptorState::Reserved(inner) => inner.pid,
+            DescriptorState::Allocated(inner) => inner.reservation.pid,
+        }
+    }
+
+    fn is_oneway(&self) -> bool {
+        match self {
+            DescriptorState::Reserved(inner) => inner.is_oneway,
+            DescriptorState::Allocated(inner) => inner.reservation.is_oneway,
+        }
+    }
+}
+
+#[derive(Clone)]
+struct Reservation {
+    debug_id: usize,
+    is_oneway: bool,
+    pid: Pid,
+}
+
+impl Reservation {
+    fn allocate<T>(self, data: Option<T>) -> Allocation<T> {
+        Allocation {
+            data,
+            reservation: self,
+        }
+    }
+}
+
+struct Allocation<T> {
+    reservation: Reservation,
+    data: Option<T>,
+}
+
+impl<T> Allocation<T> {
+    fn deallocate(self) -> (Reservation, Option<T>) {
+        (self.reservation, self.data)
+    }
+
+    fn debug_id(&self) -> usize {
+        self.reservation.debug_id
+    }
+
+    fn take(&mut self) -> Option<T> {
+        self.data.take()
+    }
+}
+
+/// The array implementation must switch to the tree if it wants to go beyond this number of
+/// ranges.
+const TREE_THRESHOLD: usize = 8;
+
+/// Represents a range of pages that have just become completely free.
+#[derive(Copy, Clone)]
+pub(crate) struct FreedRange {
+    pub(crate) start_page_idx: usize,
+    pub(crate) end_page_idx: usize,
+}
+
+impl FreedRange {
+    fn interior_pages(offset: usize, size: usize) -> FreedRange {
+        FreedRange {
+            // Divide round up
+            start_page_idx: offset.div_ceil(PAGE_SIZE),
+            // Divide round down
+            end_page_idx: (offset + size) / PAGE_SIZE,
+        }
+    }
+}
+
+struct Range<T> {
+    offset: usize,
+    size: usize,
+    state: DescriptorState<T>,
+}
+
+impl<T> Range<T> {
+    fn endpoint(&self) -> usize {
+        self.offset + self.size
+    }
+}
+
+pub(crate) struct RangeAllocator<T> {
+    inner: Impl<T>,
+}
+
+enum Impl<T> {
+    Empty(usize),
+    Array(ArrayRangeAllocator<T>),
+    Tree(TreeRangeAllocator<T>),
+}
+
+impl<T> RangeAllocator<T> {
+    pub(crate) fn new(size: usize) -> Self {
+        Self {
+            inner: Impl::Empty(size),
+        }
+    }
+
+    pub(crate) fn free_oneway_space(&self) -> usize {
+        match &self.inner {
+            Impl::Empty(size) => size / 2,
+            Impl::Array(array) => array.free_oneway_space(),
+            Impl::Tree(tree) => tree.free_oneway_space(),
+        }
+    }
+
+    pub(crate) fn count_buffers(&self) -> usize {
+        match &self.inner {
+            Impl::Empty(_size) => 0,
+            Impl::Array(array) => array.count_buffers(),
+            Impl::Tree(tree) => tree.count_buffers(),
+        }
+    }
+
+    pub(crate) fn debug_print(&self, m: &SeqFile) -> Result<()> {
+        match &self.inner {
+            Impl::Empty(_size) => Ok(()),
+            Impl::Array(array) => array.debug_print(m),
+            Impl::Tree(tree) => tree.debug_print(m),
+        }
+    }
+
+    /// Try to reserve a new buffer, using the provided allocation if necessary.
+    pub(crate) fn reserve_new(&mut self, mut args: ReserveNewArgs<T>) -> Result<ReserveNew<T>> {
+        match &mut self.inner {
+            Impl::Empty(size) => {
+                let empty_array = match args.empty_array_alloc.take() {
+                    Some(empty_array) => ArrayRangeAllocator::new(*size, empty_array),
+                    None => {
+                        return Ok(ReserveNew::NeedAlloc(ReserveNewNeedAlloc {
+                            args,
+                            need_empty_array_alloc: true,
+                            need_new_tree_alloc: false,
+                            need_tree_alloc: false,
+                        }))
+                    }
+                };
+
+                self.inner = Impl::Array(empty_array);
+                self.reserve_new(args)
+            }
+            Impl::Array(array) if array.is_full() => {
+                let allocs = match args.new_tree_alloc {
+                    Some(ref mut allocs) => allocs,
+                    None => {
+                        return Ok(ReserveNew::NeedAlloc(ReserveNewNeedAlloc {
+                            args,
+                            need_empty_array_alloc: false,
+                            need_new_tree_alloc: true,
+                            need_tree_alloc: true,
+                        }))
+                    }
+                };
+
+                let new_tree =
+                    TreeRangeAllocator::from_array(array.total_size(), &mut array.ranges, allocs);
+
+                self.inner = Impl::Tree(new_tree);
+                self.reserve_new(args)
+            }
+            Impl::Array(array) => {
+                let offset =
+                    array.reserve_new(args.debug_id, args.size, args.is_oneway, args.pid)?;
+                Ok(ReserveNew::Success(ReserveNewSuccess {
+                    offset,
+                    oneway_spam_detected: false,
+                    _empty_array_alloc: args.empty_array_alloc,
+                    _new_tree_alloc: args.new_tree_alloc,
+                    _tree_alloc: args.tree_alloc,
+                }))
+            }
+            Impl::Tree(tree) => {
+                let alloc = match args.tree_alloc {
+                    Some(alloc) => alloc,
+                    None => {
+                        return Ok(ReserveNew::NeedAlloc(ReserveNewNeedAlloc {
+                            args,
+                            need_empty_array_alloc: false,
+                            need_new_tree_alloc: false,
+                            need_tree_alloc: true,
+                        }));
+                    }
+                };
+                let (offset, oneway_spam_detected) =
+                    tree.reserve_new(args.debug_id, args.size, args.is_oneway, args.pid, alloc)?;
+                Ok(ReserveNew::Success(ReserveNewSuccess {
+                    offset,
+                    oneway_spam_detected,
+                    _empty_array_alloc: args.empty_array_alloc,
+                    _new_tree_alloc: args.new_tree_alloc,
+                    _tree_alloc: None,
+                }))
+            }
+        }
+    }
+
+    /// Deletes the allocations at `offset`.
+    pub(crate) fn reservation_abort(&mut self, offset: usize) -> Result<FreedRange> {
+        match &mut self.inner {
+            Impl::Empty(_size) => Err(EINVAL),
+            Impl::Array(array) => array.reservation_abort(offset),
+            Impl::Tree(tree) => {
+                let freed_range = tree.reservation_abort(offset)?;
+                if tree.is_empty() {
+                    self.inner = Impl::Empty(tree.total_size());
+                }
+                Ok(freed_range)
+            }
+        }
+    }
+
+    /// Called when an allocation is no longer in use by the kernel.
+    ///
+    /// The value in `data` will be stored, if any. A mutable reference is used to avoid dropping
+    /// the `T` when an error is returned.
+    pub(crate) fn reservation_commit(&mut self, offset: usize, data: &mut Option<T>) -> Result {
+        match &mut self.inner {
+            Impl::Empty(_size) => Err(EINVAL),
+            Impl::Array(array) => array.reservation_commit(offset, data),
+            Impl::Tree(tree) => tree.reservation_commit(offset, data),
+        }
+    }
+
+    /// Called when the kernel starts using an allocation.
+    ///
+    /// Returns the size of the existing entry and the data associated with it.
+    pub(crate) fn reserve_existing(&mut self, offset: usize) -> Result<(usize, usize, Option<T>)> {
+        match &mut self.inner {
+            Impl::Empty(_size) => Err(EINVAL),
+            Impl::Array(array) => array.reserve_existing(offset),
+            Impl::Tree(tree) => tree.reserve_existing(offset),
+        }
+    }
+
+    /// Call the provided callback at every allocated region.
+    ///
+    /// This destroys the range allocator. Used only during shutdown.
+    pub(crate) fn take_for_each<F: Fn(usize, usize, usize, Option<T>)>(&mut self, callback: F) {
+        match &mut self.inner {
+            Impl::Empty(_size) => {}
+            Impl::Array(array) => array.take_for_each(callback),
+            Impl::Tree(tree) => tree.take_for_each(callback),
+        }
+    }
+}
+
+/// The arguments for `reserve_new`.
+#[derive(Default)]
+pub(crate) struct ReserveNewArgs<T> {
+    pub(crate) size: usize,
+    pub(crate) is_oneway: bool,
+    pub(crate) debug_id: usize,
+    pub(crate) pid: Pid,
+    pub(crate) empty_array_alloc: Option<EmptyArrayAlloc<T>>,
+    pub(crate) new_tree_alloc: Option<FromArrayAllocs<T>>,
+    pub(crate) tree_alloc: Option<ReserveNewTreeAlloc<T>>,
+}
+
+/// The return type of `ReserveNew`.
+pub(crate) enum ReserveNew<T> {
+    Success(ReserveNewSuccess<T>),
+    NeedAlloc(ReserveNewNeedAlloc<T>),
+}
+
+/// Returned by `reserve_new` when the reservation was successul.
+pub(crate) struct ReserveNewSuccess<T> {
+    pub(crate) offset: usize,
+    pub(crate) oneway_spam_detected: bool,
+
+    // If the user supplied an allocation that we did not end up using, then we return it here.
+    // The caller will kfree it outside of the lock.
+    _empty_array_alloc: Option<EmptyArrayAlloc<T>>,
+    _new_tree_alloc: Option<FromArrayAllocs<T>>,
+    _tree_alloc: Option<ReserveNewTreeAlloc<T>>,
+}
+
+/// Returned by `reserve_new` to request the caller to make an allocation before calling the method
+/// again.
+pub(crate) struct ReserveNewNeedAlloc<T> {
+    args: ReserveNewArgs<T>,
+    need_empty_array_alloc: bool,
+    need_new_tree_alloc: bool,
+    need_tree_alloc: bool,
+}
+
+impl<T> ReserveNewNeedAlloc<T> {
+    /// Make the necessary allocations for another call to `reserve_new`.
+    pub(crate) fn make_alloc(mut self) -> Result<ReserveNewArgs<T>> {
+        if self.need_empty_array_alloc && self.args.empty_array_alloc.is_none() {
+            self.args.empty_array_alloc = Some(EmptyArrayAlloc::try_new(TREE_THRESHOLD)?);
+        }
+        if self.need_new_tree_alloc && self.args.new_tree_alloc.is_none() {
+            self.args.new_tree_alloc = Some(FromArrayAllocs::try_new(TREE_THRESHOLD)?);
+        }
+        if self.need_tree_alloc && self.args.tree_alloc.is_none() {
+            self.args.tree_alloc = Some(ReserveNewTreeAlloc::try_new()?);
+        }
+        Ok(self.args)
+    }
+}