rust: alloc: introduce allocation flags

We'll use them when allocating `Box`, `Arc`, and `UniqueArc` instances,
as well as when allocating memory for `Vec` elements. These changes will
come in subsequent patches.

Reviewed-by: Benno Lossin <benno.lossin@proton.me>
Signed-off-by: Wedson Almeida Filho <walmeida@microsoft.com>
Link: https://lore.kernel.org/r/20240328013603.206764-6-wedsonaf@gmail.com
Signed-off-by: Miguel Ojeda <ojeda@kernel.org>

4 files changed, 71 insertions, 12 deletions

diff --git a/rust/bindings/bindings_helper.h b/rust/bindings/bindings_helper.h
index 65b98831b975..ddb5644d4fd9 100644
--- a/rust/bindings/bindings_helper.h
+++ b/rust/bindings/bindings_helper.h
@@ -20,5 +20,8 @@
 
 /* `bindgen` gets confused at certain things. */
 const size_t RUST_CONST_HELPER_ARCH_SLAB_MINALIGN = ARCH_SLAB_MINALIGN;
+const gfp_t RUST_CONST_HELPER_GFP_ATOMIC = GFP_ATOMIC;
 const gfp_t RUST_CONST_HELPER_GFP_KERNEL = GFP_KERNEL;
+const gfp_t RUST_CONST_HELPER_GFP_KERNEL_ACCOUNT = GFP_KERNEL_ACCOUNT;
+const gfp_t RUST_CONST_HELPER_GFP_NOWAIT = GFP_NOWAIT;
 const gfp_t RUST_CONST_HELPER___GFP_ZERO = __GFP_ZERO;
diff --git a/rust/kernel/alloc.rs b/rust/kernel/alloc.rs
index 2c99635f9fd3..fd5541991a22 100644
--- a/rust/kernel/alloc.rs
+++ b/rust/kernel/alloc.rs
@@ -6,3 +6,64 @@
 #[cfg(not(testlib))]
 mod allocator;
 pub mod vec_ext;
+
+/// Flags to be used when allocating memory.
+///
+/// They can be combined with the operators `|`, `&`, and `!`.
+///
+/// Values can be used from the [`flags`] module.
+#[derive(Clone, Copy)]
+pub struct Flags(u32);
+
+impl core::ops::BitOr for Flags {
+    type Output = Self;
+    fn bitor(self, rhs: Self) -> Self::Output {
+        Self(self.0 | rhs.0)
+    }
+}
+
+impl core::ops::BitAnd for Flags {
+    type Output = Self;
+    fn bitand(self, rhs: Self) -> Self::Output {
+        Self(self.0 & rhs.0)
+    }
+}
+
+impl core::ops::Not for Flags {
+    type Output = Self;
+    fn not(self) -> Self::Output {
+        Self(!self.0)
+    }
+}
+
+/// Allocation flags.
+///
+/// These are meant to be used in functions that can allocate memory.
+pub mod flags {
+    use super::Flags;
+    use crate::bindings;
+
+    /// Zeroes out the allocated memory.
+    ///
+    /// This is normally or'd with other flags.
+    pub const __GFP_ZERO: Flags = Flags(bindings::__GFP_ZERO);
+
+    /// Users can not sleep and need the allocation to succeed.
+    ///
+    /// A lower watermark is applied to allow access to "atomic reserves". The current
+    /// implementation doesn't support NMI and few other strict non-preemptive contexts (e.g.
+    /// raw_spin_lock). The same applies to [`GFP_NOWAIT`].
+    pub const GFP_ATOMIC: Flags = Flags(bindings::GFP_ATOMIC);
+
+    /// Typical for kernel-internal allocations. The caller requires ZONE_NORMAL or a lower zone
+    /// for direct access but can direct reclaim.
+    pub const GFP_KERNEL: Flags = Flags(bindings::GFP_KERNEL);
+
+    /// The same as [`GFP_KERNEL`], except the allocation is accounted to kmemcg.
+    pub const GFP_KERNEL_ACCOUNT: Flags = Flags(bindings::GFP_KERNEL_ACCOUNT);
+
+    /// Ror kernel allocations that should not stall for direct reclaim, start physical IO or
+    /// use any filesystem callback.  It is very likely to fail to allocate memory, even for very
+    /// small allocations.
+    pub const GFP_NOWAIT: Flags = Flags(bindings::GFP_NOWAIT);
+}
diff --git a/rust/kernel/alloc/allocator.rs b/rust/kernel/alloc/allocator.rs
index 01ad139e19bc..8cc7099d6ae1 100644
--- a/rust/kernel/alloc/allocator.rs
+++ b/rust/kernel/alloc/allocator.rs
@@ -2,6 +2,7 @@
 
 //! Allocator support.
 
+use super::{flags::*, Flags};
 use core::alloc::{GlobalAlloc, Layout};
 use core::ptr;
 
@@ -15,7 +16,7 @@ struct KernelAllocator;
 ///
 /// - `ptr` can be either null or a pointer which has been allocated by this allocator.
 /// - `new_layout` must have a non-zero size.
-unsafe fn krealloc_aligned(ptr: *mut u8, new_layout: Layout, flags: bindings::gfp_t) -> *mut u8 {
+unsafe fn krealloc_aligned(ptr: *mut u8, new_layout: Layout, flags: Flags) -> *mut u8 {
     // Customized layouts from `Layout::from_size_align()` can have size < align, so pad first.
     let layout = new_layout.pad_to_align();
 
@@ -36,14 +37,14 @@ unsafe fn krealloc_aligned(ptr: *mut u8, new_layout: Layout, flags: bindings::gf
     //   function safety requirement.
     // - `size` is greater than 0 since it's either a `layout.size()` (which cannot be zero
     //   according to the function safety requirement) or a result from `next_power_of_two()`.
-    unsafe { bindings::krealloc(ptr as *const core::ffi::c_void, size, flags) as *mut u8 }
+    unsafe { bindings::krealloc(ptr as *const core::ffi::c_void, size, flags.0) as *mut u8 }
 }
 
 unsafe impl GlobalAlloc for KernelAllocator {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
         // SAFETY: `ptr::null_mut()` is null and `layout` has a non-zero size by the function safety
         // requirement.
-        unsafe { krealloc_aligned(ptr::null_mut(), layout, bindings::GFP_KERNEL) }
+        unsafe { krealloc_aligned(ptr::null_mut(), layout, GFP_KERNEL) }
     }
 
     unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
@@ -64,19 +65,13 @@ unsafe impl GlobalAlloc for KernelAllocator {
         //   requirement.
         // - the size of `layout` is not zero because `new_size` is not zero by the function safety
         //   requirement.
-        unsafe { krealloc_aligned(ptr, layout, bindings::GFP_KERNEL) }
+        unsafe { krealloc_aligned(ptr, layout, GFP_KERNEL) }
     }
 
     unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
         // SAFETY: `ptr::null_mut()` is null and `layout` has a non-zero size by the function safety
         // requirement.
-        unsafe {
-            krealloc_aligned(
-                ptr::null_mut(),
-                layout,
-                bindings::GFP_KERNEL | bindings::__GFP_ZERO,
-            )
-        }
+        unsafe { krealloc_aligned(ptr::null_mut(), layout, GFP_KERNEL | __GFP_ZERO) }
     }
 }
 
diff --git a/rust/kernel/prelude.rs b/rust/kernel/prelude.rs
index c85b5972c0d3..827e4dfd77df 100644
--- a/rust/kernel/prelude.rs
+++ b/rust/kernel/prelude.rs
@@ -14,7 +14,7 @@
 #[doc(no_inline)]
 pub use core::pin::Pin;
 
-pub use crate::alloc::vec_ext::VecExt;
+pub use crate::alloc::{flags::*, vec_ext::VecExt};
 
 #[doc(no_inline)]
 pub use alloc::{boxed::Box, vec::Vec};