atomic_bitflags.rs

extern crate atomig; // v0.4.1
extern crate bitflags; // v2.5.0

use core::sync::atomic::Ordering;
use atomig::*;

/// Extension trait to add flags related functions to [`atomig::Atomic`]
pub trait AtomicFlagsExt<Flags>
{
    fn fetch_insert(&self, val: Flags, order: Ordering) -> Flags;
    fn fetch_remove(&self, val: Flags, order: Ordering) -> Flags;
    fn fetch_toggle(&self, val: Flags, order: Ordering) -> Flags;
    fn fetch_set(&self, val: Flags, value: bool, order: Ordering) -> Flags;
}

// Blanket implementation Atomic<F> where F: AtomLogic + bitflags::Flags
impl<T> AtomicFlagsExt<T> for Atomic<T>
where 
    T: AtomLogic + bitflags::Flags,
    <T as Atom>::Repr: atomig::impls::PrimitiveAtomLogic
{
    fn fetch_insert(&self, val: T, order: Ordering) -> T { 
        Atomic::<T>::fetch_or(&self, val, order)
    }
    fn fetch_remove(&self, val: T, order: Ordering) -> T { 
        Atomic::<T>::fetch_and(&self, T::complement(val), order)
    }
    fn fetch_toggle(&self, val: T, order: Ordering) -> T { 
        Atomic::<T>::fetch_xor(&self, val, order)
    }
    fn fetch_set(&self, val: T, set: bool, order: Ordering) -> T { 
        if set {
            self.fetch_insert(val, order)
        } else {
            self.fetch_remove(val, order)
        }
    }
}

bitflags::bitflags! {
    // NOTE: derive(Atom) cannot be used here since the internal bits type will doesnt implement Atom
    #[derive(Debug, PartialEq)]
    struct Example: u8 {
        const ATOMIC = 1;
        const FLAG = 2;
    }
}

/// Instead use this macro to generate the Atom glue
macro_rules! impl_atomic_flags {
    ($flags:ty) => {
        impl atomig::Atom for $flags {
            type Repr = <<$flags as bitflags::Flags>::Bits as atomig::Atom>::Repr;
            fn pack(self) -> Self::Repr {
                <<$flags as bitflags::Flags>::Bits as atomig::Atom>::pack(self.bits())
            }
            fn unpack(src: Self::Repr) -> Self {
                let bits = <<$flags as bitflags::Flags>::Bits as atomig::Atom>::unpack(src);
                Self::from_bits_retain(bits)
            }
        }
        impl atomig::AtomLogic for $flags {}
    };
}
impl_atomic_flags!(Example);

#[test]
fn test() {
    let value = Atomic::new(Example::FLAG);
    
    set_atomic(&value);
    
    let load = value.load(Ordering::Relaxed);
    assert_eq!(load, Example::all());
}

fn set_atomic(atomic: &Atomic<Example>) {
    atomic.fetch_insert(Example::ATOMIC, Ordering::Relaxed);    
}