1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166

use super::IntoOptionArc;
use std::mem;
use std::ptr::null_mut;
use std::sync::atomic::{AtomicPtr, Ordering};
use std::sync::Arc;

/// An array of references in which elements may be updated and retrieved atomically.
///
/// This is a Rust interpretation of [AtomicReferenceArray](https://docs.oracle.com/javase/7/docs/api/java/util/concurrent/atomic/AtomicReferenceArray.html) from Java.
pub struct AtomicOptionRefArray<T> {
    buf: Box<[AtomicPtr<T>]>,
}

impl<T> AtomicOptionRefArray<T> {
    /// Constructs a new array with the specified length.
    /// All values will be `None`.
    pub fn new(len: usize) -> Self {
        let mut buf = Vec::with_capacity(len);

        for _ in 0..len {
            buf.push(AtomicPtr::new(null_mut()));
        }

        Self {
            buf: buf.into_boxed_slice(),
        }
    }

    /// Constructs a new array with the specified length.
    /// Uses the given function to construct each value.
    pub fn new_with<U: IntoOptionArc<T>>(len: usize, f: impl Fn(usize) -> U) -> Self {
        let mut buf = Vec::with_capacity(len);

        for i in 0..len {
            let value = f(i).into_option_arc()
                .map_or_else(null_mut, |value| Arc::into_raw(value) as *mut _);

            buf.push(AtomicPtr::new(value));
        }

        Self {
            buf: buf.into_boxed_slice(),
        }
    }

    /// Returns the number of elements in the array.
    pub fn len(&self) -> usize {
        self.buf.len()
    }

    /// Returns `true` if the array has a length of 0.
    pub fn is_empty(&self) -> bool {
        self.buf.is_empty()
    }

    /// Loads and returns a reference to an value at the given position or `None`
    /// if the value at the index is not set.
    ///
    /// Panics if `index` is out of bounds.
    pub fn load(&self, index: usize) -> Option<Arc<T>> {
        let ptr = self.buf[index].load(Ordering::SeqCst);

        ptr_to_option_arc(ptr)
    }

    /// Stores the value at the given position.
    ///
    /// Panics if `index` is out of bounds.
    pub fn store(&self, index: usize, value: impl IntoOptionArc<T>) {
        let ptr = option_arc_to_ptr(value);

        self.buf[index].store(ptr, Ordering::SeqCst);
    }

    /// Swaps the value at the given position, returning the previous value.
    ///
    /// Panics if `index` is out of bounds.
    pub fn swap(&self, index: usize, value: impl IntoOptionArc<T>) -> Option<Arc<T>> {
        let ptr = option_arc_to_ptr(value);
        let ptr = self.buf[index].swap(ptr, Ordering::SeqCst);

        ptr_to_option_arc(ptr)
    }
}

impl<T> Drop for AtomicOptionRefArray<T> {
    fn drop(&mut self) {
        for value in self.buf.iter() {
            let value = value.swap(null_mut(), Ordering::SeqCst);
            if !value.is_null() {
                unsafe {
                    // Reconstruct the Arc from the raw ptr which will trigger our destructor
                    // if there is one
                    let _ = Arc::from_raw(value);
                }
            }
        }
    }
}

fn option_arc_to_ptr<T>(value: impl IntoOptionArc<T>) -> *mut T {
    if let Some(value) = value.into_option_arc() {
        Arc::into_raw(value) as *mut _
    } else {
        null_mut()
    }
}

fn ptr_to_option_arc<T>(ptr: *mut T) -> Option<Arc<T>> {
    if ptr.is_null() {
        // Return `None` if null is stored in the AtomicPtr
        None
    } else {
        // Otherwise, reconstruct the stored Arc
        let value = unsafe { Arc::from_raw(ptr) };

        // And create a new reference to return
        let value_ = Arc::clone(&value);

        // Forget the reconstructed Arc (as its still in the array as a raw ptr)
        mem::forget(value);

        // And return our new reference
        Some(value_)
    }
}

#[cfg(test)]
mod tests {
    use super::AtomicOptionRefArray;

    #[test]
    fn test_store_load() {
        let m = AtomicOptionRefArray::<String>::new(5);

        // Store several values
        for i in 0..5 {
            m.store(i, format!("{}", i));
        }

        // Load the values and assert the values
        for i in 0..5 {
            assert_eq!(m.load(i).unwrap().as_ref(), &format!("{}", i));
        }
    }

    #[test]
    fn test_overwrite() {
        let m = AtomicOptionRefArray::<String>::new(5);

        // Store at #0
        m.store(0, String::from("Hello World"));

        // Take a reference to #0
        let m0 = m.load(0);

        // Store at #0 (again)
        m.store(0, String::from("Goodbye World"));

        // Compare value of stored #0
        assert_eq!(m0.unwrap().as_ref(), "Hello World");

        // Compare value of new #0
        assert_eq!(m.load(0).unwrap().as_ref(), "Goodbye World");
    }
}