Atomic ARC Reference

atomic_arc_reference.h

#include <stdatomic.h>
#include <stdint.h>
#include <stdbool.h>

/*!
 An atomic ARC reference. Here's how it works:

 The lowest bit of the pointer (mask 0x1) is used as a tag for the pointer being stored.
 if the tag bit is set, that means that one thread currently owns this pointer, and will release it if the value has
 been concurrently updated.

 This works for the same reason that tagged pointers work in objc, because allocations in OSX/iOS are 16-byte aligned.

 This also means that this type cannot store tagged pointers, which you shouldn't be doing anyway.

 NOTE: We can make all atomic operations here 'relaxed', as we do not require sequential ordering for loads/stores here,
 due to our use of compare_exchange before releasing the value.
 */
typedef struct atomic_arc_reference {
    union {
        _Atomic uintptr_t _value;
        __unsafe_unretained id _asObject;
    };
} atomic_arc_reference;

#define ATOMIC_ARC_MASK 0x1

/*!
 Acquires (atomically) the reference held inside this object

 @param ref The reference to fetch from

 @return The value held by this reference
 */
static inline id atomic_arc_get(struct atomic_arc_reference *ref) NS_RETURNS_RETAINED {
    extern void *objc_retain(void *);
    extern void objc_release(void *);
    extern void *objc_autorelease(void *);

    // While we retain the value, use a temp bit to indicate that another user is holding this value.
    uintptr_t oldValue = atomic_fetch_or_explicit(&ref->_value, ATOMIC_ARC_MASK, memory_order_relaxed);
    bool tempHeld = !(oldValue & ATOMIC_ARC_MASK);

    void *results = objc_retain((void *)(oldValue & (~ATOMIC_ARC_MASK)));

    // Now that we've retained the value, unset the temp bit.
    if (tempHeld) {
        uintptr_t expected = (oldValue | ATOMIC_ARC_MASK);

        // NOTE: If this fails, because that means another thread set a new value in the meantime.
        // In this case, we should release the old value here, as the assignment will NOT release while a load
        // is in progress.
        if (!atomic_compare_exchange_strong_explicit(&ref->_value, &expected, oldValue, memory_order_relaxed,
                                                     memory_order_relaxed)) {
            // NOTE: This release balances the retain in operator =, NOT the retain a few lines above.
            objc_release((void *)oldValue);
        }
    }

    // This balances the retain a few lines above.
    return (__bridge id)results;
}

/*!
 Update the pointer held by this atomic reference.

 @param ref The reference to update
 @param newValue The new value to assign.

 @return This atomic reference.
 */
static inline void atomic_arc_set(struct atomic_arc_reference *ref, id value) {
    extern void *objc_retain(void *);
    extern void objc_release(void *);
    extern void *objc_autorelease(void *);

    // We alway need to retain the value.
    uintptr_t newValue = (uintptr_t)objc_retain((__bridge void *)value);

    // Update the value of the reference, using an atomic exchange operation.
    uintptr_t oldValue = atomic_exchange_explicit(&ref->_value, newValue, memory_order_relaxed);

    // As long as another user doesn't currently own the value, we should release it afterward.
    if (!(oldValue & ATOMIC_ARC_MASK)) {
        objc_release((void *)oldValue);
    }
}

/*!
 Update the pointer held by this atomic reference, if it's equal to an expected value..

 @param ref The reference to update
 @param newValue The new value to assign.

 @return Whether or not the operation succeeded. 
 */
static inline bool atomic_arc_compare_exchange(struct atomic_arc_reference *ref, id expected, id value) {
    extern void *objc_retain(void *);
    extern void objc_release(void *);
    extern void *objc_autorelease(void *);

    // We always need to retain the value.
    uintptr_t expectedValue = (uintptr_t)expected;
    uintptr_t newValue = (uintptr_t)objc_retain((__bridge void *)value);

    while (true) {
        // Load the old value.
        uintptr_t oldValue = atomic_load_explicit(&ref->_value, memory_order_relaxed);

        // If the value has changed, then release our temp value and return.
        if ((oldValue & (~ATOMIC_ARC_MASK)) != expectedValue) {
            objc_release((void *)newValue);
            return false;
        }

        // If this compare and swap fails, it doesn't necessarily mean that the value was changed, it could have been
        // updated by a GET operation, which means we can't simply return here.
        if (atomic_compare_exchange_strong_explicit(&ref->_value, &oldValue, newValue, memory_order_relaxed,
                                                                                       memory_order_relaxed)) {
            if (!(oldValue & ATOMIC_ARC_MASK)) {
                objc_release((void *)oldValue);
            }
            return true;
        }
    }
}