Re: [PATCH v2 4/5] rust: maple_tree: add MapleTreeAlloc

[Daniel Almeida <daniel.almeida@collabora.com>]

> On 19 Aug 2025, at 07:34, Alice Ryhl <aliceryhl@google.com> wrote:
> 
> To support allocation trees, we introduce a new type MapleTreeAlloc for
> the case where the tree is created using MT_FLAGS_ALLOC_RANGE. To ensure
> that you can only call mtree_alloc_range on an allocation tree, we
> restrict thta method to the new MapleTreeAlloc type. However, all
> methods on MapleTree remain accessible to MapleTreeAlloc as allocation
> trees can use the other methods without issues.
> 
> Signed-off-by: Alice Ryhl <aliceryhl@google.com>
> ---
> rust/kernel/maple_tree.rs | 158 ++++++++++++++++++++++++++++++++++++++++++++++
> 1 file changed, 158 insertions(+)
> 
> diff --git a/rust/kernel/maple_tree.rs b/rust/kernel/maple_tree.rs
> index 17e4d8586ebad56aee87a97befdfec5741f147de..1a32960e6e721ca32ca45d8bb63fcffedeae3424 100644
> --- a/rust/kernel/maple_tree.rs
> +++ b/rust/kernel/maple_tree.rs
> @@ -33,6 +33,26 @@ pub struct MapleTree<T: ForeignOwnable> {
>     _p: PhantomData<T>,
> }
> 
> +/// A maple tree with `MT_FLAGS_ALLOC_RANGE` set.
> +///
> +/// All methods on [`MapleTree`] are also accessible on this type.
> +#[pin_data]
> +#[repr(transparent)]
> +pub struct MapleTreeAlloc<T: ForeignOwnable> {
> +    #[pin]
> +    tree: MapleTree<T>,
> +}
> +
> +// Make MapleTree methods usable on MapleTreeAlloc.
> +impl<T: ForeignOwnable> core::ops::Deref for MapleTreeAlloc<T> {
> +    type Target = MapleTree<T>;
> +
> +    #[inline]
> +    fn deref(&self) -> &MapleTree<T> {
> +        &self.tree
> +    }
> +}
> +
> /// A helper type used for navigating a [`MapleTree`].
> ///
> /// # Invariants
> @@ -364,6 +384,107 @@ pub fn load(&mut self, index: usize) -> Option<T::BorrowedMut<'_>> {
>     }
> }
> 
> +impl<T: ForeignOwnable> MapleTreeAlloc<T> {
> +    /// Create a new allocation tree.
> +    pub fn new() -> impl PinInit<Self> {
> +        let tree = pin_init!(MapleTree {
> +            // SAFETY: This initializes a maple tree into a pinned slot. The maple tree will be
> +            // destroyed in Drop before the memory location becomes invalid.
> +            tree <- Opaque::ffi_init(|slot| unsafe {
> +                bindings::mt_init_flags(slot, bindings::MT_FLAGS_ALLOC_RANGE)
> +            }),
> +            _p: PhantomData,
> +        });
> +
> +        pin_init!(MapleTreeAlloc { tree <- tree })
> +    }
> +
> +    /// Insert an entry with the given size somewhere in the given range.
> +    ///
> +    /// The maple tree will search for a location in the given range where there is space to insert
> +    /// the new range. If there is not enough available space, then an error will be returned.
> +    ///
> +    /// The index of the new range is returned.
> +    ///
> +    /// # Examples
> +    ///
> +    /// ```
> +    /// use kernel::maple_tree::{MapleTreeAlloc, AllocErrorKind};
> +    ///
> +    /// let tree = KBox::pin_init(MapleTreeAlloc::<KBox<i32>>::new(), GFP_KERNEL)?;
> +    ///
> +    /// let ten = KBox::new(10, GFP_KERNEL)?;
> +    /// let twenty = KBox::new(20, GFP_KERNEL)?;
> +    /// let thirty = KBox::new(30, GFP_KERNEL)?;
> +    /// let hundred = KBox::new(100, GFP_KERNEL)?;
> +    ///
> +    /// // Allocate three ranges.
> +    /// let idx1 = tree.alloc_range(100, ten, ..1000, GFP_KERNEL)?;
> +    /// let idx2 = tree.alloc_range(100, twenty, ..1000, GFP_KERNEL)?;
> +    /// let idx3 = tree.alloc_range(100, thirty, ..1000, GFP_KERNEL)?;
> +    ///
> +    /// assert_eq!(idx1, 0);
> +    /// assert_eq!(idx2, 100);
> +    /// assert_eq!(idx3, 200);
> +    ///
> +    /// // This will fail because the remaining space is too small.
> +    /// assert_eq!(
> +    ///     tree.alloc_range(800, hundred, ..1000, GFP_KERNEL).unwrap_err().cause,
> +    ///     AllocErrorKind::Busy,
> +    /// );
> +    /// # Ok::<_, Error>(())
> +    /// ```
> +    pub fn alloc_range<R>(
> +        &self,
> +        size: usize,
> +        value: T,
> +        range: R,
> +        gfp: Flags,
> +    ) -> Result<usize, AllocError<T>>
> +    where
> +        R: RangeBounds<usize>,
> +    {
> +        let Some((min, max)) = to_maple_range(range) else {

Ok, I see that the returned range can mean multiple things, depending on
context. Fine, you can disregard my previous comment about this function.

> +            return Err(AllocError {
> +                value,
> +                cause: AllocErrorKind::InvalidRequest,
> +            });
> +        };
> +
> +        let ptr = T::into_foreign(value);
> +        let mut index = 0;
> +
> +        // SAFETY: The tree is valid, and we are passing a pointer to an owned instance of `T`.
> +        let res = to_result(unsafe {
> +            bindings::mtree_alloc_range(
> +                self.tree.tree.get(),
> +                &mut index,
> +                ptr,
> +                size,
> +                min,
> +                max,
> +                gfp.as_raw(),
> +            )
> +        });
> +
> +        if let Err(err) = res {
> +            // SAFETY: As `mtree_alloc_range` failed, it is safe to take back ownership.
> +            let value = unsafe { T::from_foreign(ptr) };
> +
> +            let cause = if err == ENOMEM {
> +                AllocErrorKind::AllocError(kernel::alloc::AllocError)
> +            } else if err == EBUSY {
> +                AllocErrorKind::Busy
> +            } else {
> +                AllocErrorKind::InvalidRequest
> +            };
> +            Err(AllocError { value, cause })
> +        } else {
> +            Ok(index)
> +        }
> +    }
> +}
> +
> impl<'tree, T: ForeignOwnable> MaState<'tree, T> {
>     /// Initialize a new `MaState` with the given tree.
>     ///
> @@ -480,3 +601,40 @@ fn from(insert_err: InsertError<T>) -> Error {
>         Error::from(insert_err.cause)
>     }
> }
> +
> +/// Error type for failure to insert a new value.
> +pub struct AllocError<T> {
> +    /// The value that could not be inserted.
> +    pub value: T,
> +    /// The reason for the failure to insert.
> +    pub cause: AllocErrorKind,
> +}
> +
> +/// The reason for the failure to insert.
> +#[derive(PartialEq, Eq, Copy, Clone)]
> +pub enum AllocErrorKind {
> +    /// There is not enough space for the requested allocation.
> +    Busy,
> +    /// Failure to allocate memory.
> +    AllocError(kernel::alloc::AllocError),
> +    /// The insertion request was invalid.
> +    InvalidRequest,
> +}
> +
> +impl From<AllocErrorKind> for Error {
> +    #[inline]
> +    fn from(kind: AllocErrorKind) -> Error {
> +        match kind {
> +            AllocErrorKind::Busy => EBUSY,
> +            AllocErrorKind::AllocError(kernel::alloc::AllocError) => ENOMEM,
> +            AllocErrorKind::InvalidRequest => EINVAL,
> +        }
> +    }
> +}
> +
> +impl<T> From<AllocError<T>> for Error {
> +    #[inline]
> +    fn from(insert_err: AllocError<T>) -> Error {
> +        Error::from(insert_err.cause)
> +    }
> +}
> 
> -- 
> 2.51.0.rc1.167.g924127e9c0-goog
> 
> 

Reviewed-by: Daniel Almeida <daniel.almeida@collabora.com>