atomic_box_or_small_int.rs

use std::boxed::Box;
use std::sync::atomic::AtomicPtr;
use std::sync::atomic::Ordering;
use std::ptr;
use std::mem;


/// Max value for flag stored in `AtomicBoxOrFlag`
/// `3` is safe value, because pointers are at least 4-byte aligned.
/// Although in practice `malloc` never return small addresses,
/// so this value could be something like `0x100`.
pub const MAX_FLAG: usize = 3;

#[derive(Debug)]
pub enum BoxOrSmallInt<T> {
    Box(Box<T>),
    SmallInt(usize),
}

impl<T> BoxOrSmallInt<T> {
    /// Decode raw value
    unsafe fn from_raw(raw: *mut T) -> BoxOrSmallInt<T> {
        let u = raw as usize;
        if u <= MAX_FLAG {
            BoxOrSmallInt::SmallInt(u)
        } else {
            BoxOrSmallInt::Box(Box::from_raw(raw))
        }
    }
}


/// Atomic holder for either pointer of `Box<T>`.
/// It is safe, although it exposes its internal `AtomicPtr<T>` field.
///
/// `AtomicBoxOrFlag` owns a pointer, thus it `drop` it self `drop`.
#[derive(Debug)]
pub struct AtomicBoxOrSmallInt<T>(pub AtomicPtr<T>);

impl<T> AtomicBoxOrSmallInt<T> {
    /// Create an empty object with value `0`.
    pub fn new() -> AtomicBoxOrSmallInt<T> {
        AtomicBoxOrSmallInt(AtomicPtr::new(ptr::null_mut()))
    }

    /// Safe version of `compare_exchange`: it properly manages object ownership.
    pub fn compare_exchange_box(
        &self, compare: *mut T, exchange: Box<T>, success: Ordering, failure: Ordering)
            -> Result<BoxOrSmallInt<T>, (*mut T, Box<T>)>
    {
        let exchange_raw = exchange.as_ref() as *const T as *mut T;
        match self.0.compare_exchange(compare, exchange_raw, success, failure) {
            Ok(_) => {
                mem::forget(exchange);
                Ok(unsafe { BoxOrSmallInt::from_raw(compare) })
            },
            Err(v) => Err((v, exchange)),
        }
    }
}

impl<T> Drop for AtomicBoxOrSmallInt<T> {
    fn drop(&mut self) {
        let ptr = self.0.load(Ordering::SeqCst);
        if ptr as usize > MAX_FLAG {
            unsafe { Box::from_raw(ptr); }
        }
    }
}