How do I store an atomic pointer with sentinel values in Rust?

Question

In some unsafe Rust I need to store an atomic pointer that normally refers to a real object, but also has four sentinel values for signaling things between the communicating threads. What's the idiomatic way to do this?

In C++ I would make sure that the alignment of the pointee is at least four and then use std::atomic<uintptr_t> with the sentinels being the integers zero to three. What are the corresponding relevant APIs for this in Rust, and are there any extra UB issues (pointer provenance?) that I'm likely to run into?

Answer 1

Provenance

As per the guidance provided in the Pointers vs Integers section of the std::ptr module:

Note that a pointer can represent a usize (via without_provenance), so the right type to use in situations where a value is “sometimes a pointer and sometimes a bare usize” is a pointer type.

That is, in Rust, you should therefore use an AtomicPtr, not an AtomicUsize.

Manipulation

Since materializing a pointer from an integer would lose provenance, specific facilities are provided to forge an arbitrary pointer with the provenance of an existing pointer. Specifically, <*[const|mut] T>::with_addr(self, addr: usize) -> *[const|mut] T.

However, there's an even simpler solution for bit-masking map_addr:

let ptr = ptr.map_addr(|addr| addr | 0xf);
let ptr = ptr.map_addr(|addr| addr & !0xf);

Conclusion

Write two little wrapper types:

use core::sync::atomic::{AtomicPtr, Ordering};

/// A pointer with some tag value stored in the low bits.
///
/// - `T`: a type whose alignment must be greater than or equal to `MASK + 1`.
/// - `MASK`: the bits which may be set in the pointer.
///
/// `MASK + 1` SHALL not overflow and SHALL be a power of 2.
pub struct TaggedPtr<T, const MASK: usize>(*mut T);

impl<T, const MASK: usize> TaggedPtr<T, MASK> {
    /// Creates a new tagged pointer from the raw pointer and tag.
    ///
    /// # Panics
    ///
    /// When debug assertions are enabled, if the low bits of `ptr` are not 0
    /// or if the `tag` value overflows `MASK`.
    ///
    /// If debug assertions are disabled, the low bits of `ptr` are masked away
    /// and the `tag` value is truncated.
    pub fn new(ptr: *mut T, tag: usize) -> Self {
        const {
            assert!((MASK + 1).is_power_of_two());
            assert!((MASK + 1) <= core::mem::align_of::<T>());
        };

        debug_assert_eq!(0, (ptr as usize) & MASK);
        debug_assert_eq!(0, tag & !MASK);

        let tagged = ptr.map_addr(|addr| addr | (tag & MASK));

        Self(tagged)
    }

    /// Gets the raw pointer and tag.
    ///
    /// # Safety Guarantee
    ///
    /// The raw pointer and tag are guaranteed to match the latest set value _iff_
    /// the raw pointer low bits were 0 and the tag did not overflow `MASK`.
    #[inline(always)]
    pub fn get(&self) -> (*mut T, usize) {
        let ptr = self.0.map_addr(|addr| addr & !MASK);
        let tag = (self.0 as usize) & MASK;
        
        (ptr, tag)
    }


    /// Set the value to a new raw pointer and tag.
    ///
    /// # Panics
    ///
    /// When debug assertions are enabled, if the low bits of `ptr` are not 0
    /// or if the `tag` value overflows `MASK`.
    ///
    /// If debug assertions are disabled, the low bits of `ptr` are masked away
    /// and the `tag` value is truncated.
    #[inline(always)]
    pub fn set(&mut self, ptr: *mut T, tag: usize) {
        debug_assert_eq!(0, tag & !MASK);

        self.0 = ptr.map_addr(|addr| addr | (tag & MASK));
    }
}

/// An atomic `TaggedPtr<T, MASK>`.
pub struct AtomicTaggedPtr<T, const MASK: usize>(AtomicPtr<T>);

impl<T, const MASK: usize> AtomicTaggedPtr<T, MASK> {
    pub fn new(ptr: TaggedPtr<T, MASK>) -> Self {
        Self(AtomicPtr::new(ptr.0))
    }

    #[inline(always)]
    pub fn load(&self, ordering: Ordering) -> TaggedPtr<T, MASK> {
        TaggedPtr(self.0.load(ordering))
    }
    
    #[inline(always)]
    pub fn store(&self, ptr: TaggedPtr<T, MASK>, ordering: Ordering) {
        self.0.store(ptr.0, ordering);
    }

    //  And more utility functions, as desired. See AtomicPtr API as reference.
}