Re: [PATCH v2 1/2] rust: sync: atomic: Add AtomicFlag::get_mut

[Boqun Feng <boqun@kernel.org>]

[For some unknown reasons, I cannot send my reply via gmail hence reply
 via kernel.org account, I might switch from gmail later on]

On Tue, Jan 27, 2026 at 09:52:59PM +0900, FUJITA Tomonori wrote:
> From: FUJITA Tomonori <fujita.tomonori@gmail.com>
> 
> AtomicFlag exposes a bool API, but it lacks a get_mut() equivalent to
> Atomic<T>::get_mut().
> 
> Also add kunit tests for AtomicFlag.
> 
> Signed-off-by: FUJITA Tomonori <fujita.tomonori@gmail.com>
> ---
>  rust/kernel/sync/atomic.rs           | 20 ++++++++++++++++++++
>  rust/kernel/sync/atomic/predefine.rs | 17 +++++++++++++++++
>  2 files changed, 37 insertions(+)
> 
> diff --git a/rust/kernel/sync/atomic.rs b/rust/kernel/sync/atomic.rs
> index 6c46335bdb8c..b6c01d9f3a46 100644
> --- a/rust/kernel/sync/atomic.rs
> +++ b/rust/kernel/sync/atomic.rs
> @@ -591,6 +591,26 @@ pub fn store<Ordering: ordering::ReleaseOrRelaxed>(&self, b: bool, o: Ordering)
>          self.0.store(b.into(), o)
>      }
>  
> +    /// Returns a mutable reference to the underlying flag as a `bool`.
> +    ///
> +    /// This is safe because the mutable reference of the atomic flag guarantees exclusive access.
> +    pub fn get_mut(&mut self) -> &mut bool {
> +        let byte_ptr = {
> +            let ptr = self.0.as_ptr().cast::<u8>();
> +            let offset = if cfg!(target_endian = "big") {
> +                core::mem::size_of::<Flag>() - 1
> +            } else {
> +                0
> +            };

The idea is solid, but I want to avoid endian handling in the function,
I would prefer a "struct declaration" solution like:

    #[cfg(target_endian = "big")]
    #[repr(align(4))]
    pub(super) struct FlagInner {
        _pad: [i8; 3],
	bool_field: bool,
    }

    #[cfg(target_endian = "little")]
    #[repr(align(4))]
    struct FlagInner {
	bool_field: bool,
        _pad: [i8; 3],
    }

redefine `Flag` as `BoolFlag`

    #[repr(i32)]
    pub enum BoolFlag {
        Clear = 0,
	Set = 1,
    }

and `Flag` becomes a union of `BoolFlag` and `FlagInner`:

    /// # Invariants
    /// `Flag` is either 0 or 1 in a i32 representation which implies
    /// that `inner` is always valid as long as `_pad` stays 0.
    pub union Flag {
        pub(super) inner: FlagInner,
        pub flag: BoolFlag,
    }

    // can static_assert that `Flag` and `BoolFlag` has the same
    // alignement and size.

then

    impl AtomicFlag {
        pub fn get_mut(&mut self) -> &mut bool {
	    let flag = self.0.get_mut(); // <- &mut Flag

            // INVARIANTS: flag.inner._pad cannot be modified via the
	    // returned reference.
	    // SAFETY: Per type invariants, `flag.inner.bool_field` is
	    // always a valid bool.
	    unsafe { &mut flag.inner.bool_field }
	}
    }

Thoughts?

Regards,
Boqun

> +
> +            // SAFETY: `ptr` is valid for `size_of::<Flag>()` bytes; `offset` selects the LSB.
> +            unsafe { ptr.add(offset) }
> +        };
> +
> +        // SAFETY: The LSB holds `0`/`1` for `Flag::Clear/Set`, and `bool` is `i8`-sized/aligned.
> +        unsafe { &mut *byte_ptr.cast::<bool>() }
> +    }
> +
>      /// Stores a value to the atomic flag and returns the previous value.
>      #[inline(always)]
>      pub fn xchg<Ordering: ordering::Ordering>(&self, b: bool, o: Ordering) -> bool {
> diff --git a/rust/kernel/sync/atomic/predefine.rs b/rust/kernel/sync/atomic/predefine.rs
> index 11bc67ab70a3..e413b9e9fe21 100644
> --- a/rust/kernel/sync/atomic/predefine.rs
> +++ b/rust/kernel/sync/atomic/predefine.rs
> @@ -351,6 +351,23 @@ fn atomic_bool_tests() {
>          assert_eq!(Ok(false), x.cmpxchg(false, true, Full));
>      }
>  
> +    #[test]
> +    fn atomic_flag_tests() {
> +        let mut flag = AtomicFlag::new(false);
> +
> +        assert_eq!(false, flag.load(Relaxed));
> +
> +        *flag.get_mut() = true;
> +        assert_eq!(true, flag.load(Relaxed));
> +
> +        assert_eq!(true, flag.xchg(false, Relaxed));
> +        assert_eq!(false, flag.load(Relaxed));
> +
> +        *flag.get_mut() = true;
> +        assert_eq!(Ok(true), flag.cmpxchg(true, false, Full));
> +        assert_eq!(false, flag.load(Relaxed));
> +    }
> +
>      #[test]
>      fn atomic_ptr_tests() {
>          let mut v = 42;
> -- 
> 2.43.0
>