[PATCH 05/10] rust: alloc: introduce allocation flags

[Wedson Almeida Filho <wedsonaf@gmail.com>]

From: Wedson Almeida Filho <walmeida@microsoft.com>

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.

Signed-off-by: Wedson Almeida Filho <walmeida@microsoft.com>
---
 rust/bindings/bindings_helper.h |  3 ++
 rust/kernel/alloc.rs            | 56 +++++++++++++++++++++++++++++++++
 rust/kernel/prelude.rs          |  2 +-
 3 files changed, 60 insertions(+), 1 deletion(-)

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 8ad57a2e693e..ad48ac8dc13d 100644
--- a/rust/kernel/alloc.rs
+++ b/rust/kernel/alloc.rs
@@ -6,3 +6,59 @@
 #[cfg(not(testlib))]
 mod allocator;
 pub mod vecext;
+
+/// 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;
+
+    /// 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/prelude.rs b/rust/kernel/prelude.rs
index a0177f195dec..a7b203f87aa1 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::vecext::VecExt;
+pub use crate::alloc::{flags::*, vecext::VecExt};
 
 #[doc(no_inline)]
 pub use alloc::{boxed::Box, vec::Vec};
-- 
2.34.1