[RFC PATCH 2/2] rust: add maple tree abstractions

[Andrew Ballance <andrewjballance@gmail.com>]

maple trees are sparse array like data structure that maps
non-overlapping ranges to pointers.

implements basic abstractions for maple trees in rust.

Link: https://docs.kernel.org/core-api/maple_tree.html
Signed-off-by: Andrew Ballance <andrewjballance@gmail.com>
---
 rust/helpers/helpers.c    |   1 +
 rust/helpers/maple_tree.c |  25 +++
 rust/kernel/lib.rs        |   1 +
 rust/kernel/maple_tree.rs | 340 ++++++++++++++++++++++++++++++++++++++
 4 files changed, 367 insertions(+)
 create mode 100644 rust/helpers/maple_tree.c
 create mode 100644 rust/kernel/maple_tree.rs

diff --git a/rust/helpers/helpers.c b/rust/helpers/helpers.c
index 0640b7e115be..6bac5ffe1684 100644
--- a/rust/helpers/helpers.c
+++ b/rust/helpers/helpers.c
@@ -18,6 +18,7 @@
 #include "io.c"
 #include "jump_label.c"
 #include "kunit.c"
+#include "maple_tree.c"
 #include "mutex.c"
 #include "page.c"
 #include "platform.c"
diff --git a/rust/helpers/maple_tree.c b/rust/helpers/maple_tree.c
new file mode 100644
index 000000000000..3e70db431616
--- /dev/null
+++ b/rust/helpers/maple_tree.c
@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include <linux/maple_tree.h>
+
+
+void rust_helper_mt_init_flags(struct maple_tree *mt, unsigned int flags)
+{
+	mt_init_flags(mt, flags);
+}
+
+void rust_helper_mtree_lock(struct maple_tree *mt)
+{
+	mtree_lock(mt);
+}
+
+void rust_helper_mtree_unlock(struct maple_tree *mt)
+{
+	mtree_unlock(mt);
+}
+
+struct ma_state rust_helper_MA_STATE(struct maple_tree *mt, unsigned long start, unsigned long end)
+{
+	MA_STATE(mas, mt, start, end);
+	return mas;
+}
diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index 7697c60b2d1a..e84538d7e643 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -57,6 +57,7 @@
 #[cfg(CONFIG_KUNIT)]
 pub mod kunit;
 pub mod list;
+pub mod maple_tree;
 pub mod miscdevice;
 #[cfg(CONFIG_NET)]
 pub mod net;
diff --git a/rust/kernel/maple_tree.rs b/rust/kernel/maple_tree.rs
new file mode 100644
index 000000000000..be107e11a3dc
--- /dev/null
+++ b/rust/kernel/maple_tree.rs
@@ -0,0 +1,340 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Maple trees.
+//!
+//! C header: [`include/linux/maple_tree.h`](srctree/include/linux/maple_tree.h)
+//!
+//! Reference: <https://docs.kernel.org/core-api/maple_tree.html>
+//!
+//! # Examples
+//! ```
+//! # use kernel::maple_tree::*;
+//! # use kernel::alloc::{KBox, flags::GFP_KERNEL};
+//! let mtree = KBox::pin_init(MapleTree::new(flags::DEFAULT_TREE), GFP_KERNEL)?;
+//! let mut guard = mtree.lock();
+//! let entry = KBox::new(5, GFP_KERNEL)?;
+//! guard.insert_range(0, 10, entry, GFP_KERNEL);
+//!
+//! for i in 0..=10 {
+//!    assert_eq!(guard.get(i), Some(&5));
+//! }
+//!
+//! guard.remove(2);
+//!
+//! for i in 0..=10 {
+//!    assert_eq!(guard.get(i), None);
+//! }
+//!
+//! # Ok::<(), Error>(())
+//! ```
+
+// TODO and missing features
+// - the C version suports external locking this only supports the internal spinlock
+// - this version can only have one reader because the safety requirements
+//   means that we need to hold the lock
+// - there is currently no api for the functionality enabled by alloc_trees
+// - future work may use rcu to enable lockless operation.
+// - add a way to iter through ranges
+//
+// SAFETY:
+// we cannot derefrence any pointers without holding the spinlock because
+// all returned pointers are guaranteed to have been inserted by the user
+// but the pointers are not guaranteed to be still be valid
+// another thread may have already removed and dropped the pointers
+// so to safely deref the returned pointers the user must
+// have exclusive write access to the `MapleTree`
+// or rcu_read_lock() is held and updater side uses a synchronize_rcu()
+
+use core::{ffi::c_void, marker::PhantomData, pin::Pin, ptr::NonNull};
+
+use macros::pin_data;
+use macros::pinned_drop;
+
+use crate::error::code;
+use crate::prelude::PinInit;
+
+use crate::{
+    alloc, bindings,
+    error::{self, Error},
+    pin_init,
+    types::{ForeignOwnable, NotThreadSafe, Opaque},
+};
+
+/// A `MapleTree` is a tree like data structure that is optimized for storing
+/// non-overlaping ranges and mapping them to pointers.
+/// They use rcu locks and an internal spinlock to syncronise access.
+///
+/// Note that maple tree ranges are range inclusive.
+// # Invariants
+// self.root is always a valid and initialized `bindings::maple_tree`
+//
+// all values inserted into the tree come from `T::into_foreign`
+#[pin_data(PinnedDrop)]
+pub struct MapleTree<T: ForeignOwnable> {
+    #[pin]
+    root: Opaque<bindings::maple_tree>,
+    _p: PhantomData<T>,
+}
+
+impl<T: ForeignOwnable> MapleTree<T> {
+    /// creates a new `MapleTree` with the specified `flags`
+    ///
+    /// see [`flags`] for the list of flags and their usage
+    pub fn new(flags: Flags) -> impl PinInit<Self> {
+        pin_init!(
+            Self{
+                // SAFETY:
+                // - mt is valid because of ffi_init
+                // - maple_tree contains a lock which should be pinned
+                root <- Opaque::ffi_init(|mt| unsafe {
+                    bindings::mt_init_flags(mt, flags.as_raw())
+                }),
+                _p: PhantomData
+            }
+
+        )
+    }
+
+    /// helper for internal use.
+    /// returns an iterator through the maple tree
+    /// # Safety
+    /// - the user must ensure that it has exclusive access to the tree if no locks are held
+    /// - if the internal lock is held it is safe to deref the returned pointers
+    /// - if the rcu lock is held the pointers will be a value that was inserted
+    ///   by the user but possibly may be invalid
+    unsafe fn iter_no_lock(&self) -> IterNoLock<'_, T> {
+        // SAFETY:
+        // self.root.get() will allways point to a valid maple_tree
+        // by the invariants of MapleTree
+        let ma_state = unsafe { Opaque::new(bindings::MA_STATE(self.root.get(), 0, usize::MAX)) };
+        IterNoLock {
+            ma_state,
+            _p: PhantomData,
+        }
+    }
+
+    /// locks the `Mapletree`'s internal spinlock and returns a [`Guard`].
+    /// When the `Guard` is dropped, the internal spinlock is unlocked
+    /// if the lock is already held by the current thread this will deadlock
+    pub fn lock(&self) -> Guard<'_, T> {
+        // SAFETY:
+        // self.root.get() will allways point to a valid maple_tree
+        // by the invariants of MapleTree
+        unsafe { bindings::mtree_lock(self.root.get()) };
+        Guard {
+            tree: self,
+            _not_send: NotThreadSafe,
+        }
+    }
+}
+
+#[pinned_drop]
+impl<T: ForeignOwnable> PinnedDrop for MapleTree<T> {
+    fn drop(self: Pin<&mut Self>) {
+        // SAFETY:
+        // - we have exclusive access to self because we should have
+        //   exclussive access whenever drop is called
+        //   if drop is called multiple times an invariant is already violated
+        for i in unsafe { self.iter_no_lock() } {
+            //SAFETY:
+            // - we can call from_foreign because all values inserted into a MapleTree
+            //   come from T::into_foreign
+            // - i.as_ptr is guaranteed to not be null because of the invariant of NonNull
+            // - we have exclusive access to self because we should have
+            //   exclussive access whenever drop is called
+            let original = unsafe { T::from_foreign(i.as_ptr()) };
+            drop(original);
+        }
+        // SAFETY:
+        // - self.root.get() will allways point to a valid maple_tree
+        //   by the invariants of MapleTree
+        // - we can call this without taking the spinlock because we should have
+        //   exclusive access whenever drop is called
+        unsafe {
+            bindings::__mt_destroy(self.root.get());
+        }
+    }
+}
+
+// SAFETY: `MapleTree<T>` has no shared mutable state so it is `Send` iff `T` is `Send`.
+// MapleTree is still pin_init so it still cannot be moved but this means that types like
+// Box<MapleTree<T>> can be sent between threads
+unsafe impl<T: ForeignOwnable + Send> Send for MapleTree<T> {}
+
+// SAFETY: `MapleTree<T>` has interior mutability so it is `Sync` iff `T` is `Send`.
+unsafe impl<T: ForeignOwnable + Send> Sync for MapleTree<T> {}
+
+/// an iterator over all of the values in a [`MapleTree`].
+/// this intentionally does not hold the rcu lock or the internal lock
+/// the user must ensure that it exclusive access to the tree if no locks are held
+/// if the internal lock is held it is safe to deref the returned pointers
+/// if the rcu lock is held the pointers will be a value that was inserted
+/// by the user but possibly may be invalid
+struct IterNoLock<'a, T: ForeignOwnable> {
+    ma_state: Opaque<bindings::ma_state>,
+    _p: PhantomData<&'a MapleTree<T>>,
+}
+
+impl<'a, T: ForeignOwnable> Iterator for IterNoLock<'a, T> {
+    type Item = NonNull<c_void>;
+    fn next(&mut self) -> Option<Self::Item> {
+        // SAFETY:
+        // self.ma_state.get() will allways point to a valid ma_state by the invariants of Iter
+        let ptr: *mut c_void = unsafe { bindings::mas_find(self.ma_state.get(), usize::MAX) };
+        NonNull::new(ptr)
+    }
+}
+
+/// A lock guard for a [`MapleTree`]'s internal spinlock
+///
+/// The lock is unlocked when the guard goes out of scope
+// # Invariants
+//
+// `tree` is always a valid reference to a locked `MapleTree`
+// `tree` is unlocked when the guard is dropped
+pub struct Guard<'a, T: ForeignOwnable> {
+    tree: &'a MapleTree<T>,
+    _not_send: NotThreadSafe,
+}
+
+impl<'a, T: ForeignOwnable> Guard<'a, T> {
+    /// Removes a value at the specified index.
+    /// if there is no value at the index returns `None`.
+    /// if there is a value at the index returns it and unmaps the entire range
+    pub fn remove(&mut self, index: usize) -> Option<T> {
+        // SAFETY:
+        // - we can safely call MA_STATE because self.tree.root.get() will be valid
+        //   by the invariants of guard
+        // - we can safely call mas_erase because the lock is held and mas is valid
+        // - we can call try_from_foreign because all values inserted into a MapleTree
+        //   come from T::into_foreign
+        unsafe {
+            let mas = Opaque::new(bindings::MA_STATE(self.tree.root.get(), index, index));
+            let removed = bindings::mas_erase(mas.get());
+            T::try_from_foreign(removed)
+        }
+    }
+
+    /// Returns a reference to the `T` at `index` if it exists,
+    /// otherwise returns `None`
+    pub fn get(&self, index: usize) -> Option<T::Borrowed<'_>> {
+        // SAFETY:
+        // self.tree.root.get() will always be valid because of the invariants of MapleTree
+        let ptr = unsafe { bindings::mtree_load(self.tree.root.get(), index) };
+        if ptr.is_null() {
+            return None;
+        }
+        // SAFETY:
+        // - we can safely call borrow because all values inserted into a MapleTree
+        //   come from T::into_foreign
+        // - ptr is not null by the check above
+        Some(unsafe { T::borrow(ptr) })
+    }
+
+    /// Returns a mut reference to the `T` at `index` if it exists,
+    /// otherwise returns `None`
+    pub fn get_mut(&mut self, index: usize) -> Option<T::BorrowedMut<'_>> {
+        // SAFETY:
+        // self.tree.root.get() will always be valid because of the invariants of MapleTree
+        let ptr = unsafe { bindings::mtree_load(self.tree.root.get(), index) };
+        if ptr.is_null() {
+            return None;
+        }
+        // SAFETY:
+        // - we can safely call borrowmut because all values inserted into a MapleTree
+        //   come from T::into_foreign
+        // - ptr is not null by the check above
+        // - we have exclusive ownership becauce this function takes `&mut self`
+        Some(unsafe { T::borrow_mut(ptr) })
+    }
+
+    /// a convenience alias for [`Self::insert_range`] where `start == end`
+    pub fn insert(&mut self, index: usize, value: T, gfp: alloc::Flags) -> Result<(), (T, Error)> {
+        self.insert_range(index, index, value, gfp)
+    }
+
+    /// Maps the range `[start, end]` to `value` in the MapleTree.
+    /// if `[start, end]` overlaps with any range already inserted, then `value` will
+    /// not be inserted.
+    ///
+    /// * Returns `Ok(())` if insertion is successful
+    ///
+    /// * Returns `Err((value, EINVAL))` if `T::into_foreign` is `NULL`
+    ///   or is in [0, 4096] with the bottom two bits set to `10` (ie 2, 6, 10 .. 4094)
+    ///   or `start` > `end`.
+    ///
+    /// * Returns `Err((value, EEXISTS))` if there is any entry already within the range.
+    ///
+    /// * Returns `Err((value, ENOMEM))` if allocation failed.
+    pub fn insert_range(
+        &mut self,
+        start: usize,
+        end: usize,
+        value: T,
+        gfp: alloc::Flags,
+    ) -> Result<(), (T, Error)> {
+        let ptr = value.into_foreign();
+        // a insertion of NULL will succeed even if there are values stored there (i tested this)
+        // this may remove values that we do not want to remove
+        // and any stored T that gets overwritten will not be dropped
+        // which we do not want to happen
+        // so return early if ptr is NULL
+        if ptr.is_null() {
+            // SAFETY:
+            // ptr is from T::into_foreign so we can safely convert it back
+            let original = unsafe { T::from_foreign(ptr) };
+            return Err((original, code::EINVAL));
+        }
+
+        // SAFETY:
+        // - we can call __mtree_insert_range because we hold the lock because of the
+        //   invariants of self
+        // - self.tree.root.get() will always be valid because of the invariants of MapleTree
+        let errno = unsafe {
+            bindings::__mtree_insert_range(self.tree.root.get(), start, end, ptr, gfp.as_raw())
+        };
+
+        let err = error::to_result(errno);
+        // SAFETY:
+        // - we can call from_foreign because all values inserted into a MapleTree
+        //   come from T::into_foreign
+        // - we have exclusive ownership of ptr because if err is an error then, ptr was
+        //   not inserted into the MapleTree
+        //
+        err.map_err(|e| unsafe { (T::from_foreign(ptr), e) })
+    }
+}
+
+impl<T: ForeignOwnable> Drop for Guard<'_, T> {
+    fn drop(&mut self) {
+        // SAFETY:
+        // - unlock the MapleTree because the guard is being dropped
+        // - self.tree.root.get() will always be valid because of the invariants of MapleTree
+        unsafe { bindings::mtree_unlock(self.tree.root.get()) };
+    }
+}
+
+#[derive(Clone, Copy, PartialEq)]
+/// flags to be used with [`MapleTree::new`].
+///
+/// values can used from the [`flags`] module.
+pub struct Flags(u32);
+
+impl Flags {
+    pub(crate) fn as_raw(self) -> u32 {
+        self.0
+    }
+}
+
+/// flags to be used with [`MapleTree::new`]
+pub mod flags {
+    use super::Flags;
+
+    /// Creates a default MapleTree
+    pub const DEFAULT_TREE: Flags = Flags(0);
+
+    /// if a `MapleTree` is created with ALLOC_TREE the `MapleTree` will be a alloc tree.
+    /// alloc trees have a lower branching factor but allows the user to search
+    /// for a gap of a given size.
+    pub const ALLOC_TREE: Flags = Flags(bindings::MT_FLAGS_ALLOC_RANGE);
+}
-- 
2.49.0