magic_orb/
magic_orb.rs

1#[cfg(feature = "not_nightly")]
2use crate::sync_unsafe_cell::SyncUnsafeCell;
3#[cfg(not(feature = "not_nightly"))]
4use std::cell::SyncUnsafeCell;
5
6use std::{
7    cmp::min,
8    fmt::Debug,
9    sync::{
10        Arc,
11        atomic::{AtomicBool, AtomicUsize, Ordering},
12    },
13};
14
15#[derive(Debug, Clone)]
16pub struct MagicOrb<T>
17where
18    T: Clone + Send,
19{
20    buf: Arc<SyncUnsafeCell<Vec<T>>>,
21    write: Arc<AtomicUsize>,
22    lock: Arc<AtomicBool>,
23    len: Arc<AtomicUsize>,
24    capacity: Arc<AtomicUsize>,
25}
26
27impl<T: Default + Clone + Send> MagicOrb<T> {
28    pub fn new_default(size: usize) -> Self {
29        if size == 0 {
30            panic!("Can't crate MagicOrb with buffer size 0");
31        }
32        MagicOrb {
33            buf: Arc::new(SyncUnsafeCell::new(vec![T::default(); size])),
34            write: Arc::new(AtomicUsize::new(0)),
35            lock: Arc::new(AtomicBool::new(false)),
36            len: Arc::new(AtomicUsize::new(size)),
37            capacity: Arc::new(AtomicUsize::new(size)),
38        }
39    }
40}
41
42impl<T: Clone + Send> From<Vec<T>> for MagicOrb<T> {
43    fn from(value: Vec<T>) -> Self {
44        if value.is_empty() {
45            panic!("Can't crate MagicOrb with buffer size 0");
46        }
47        MagicOrb {
48            capacity: Arc::new(AtomicUsize::new(value.len())),
49            len: Arc::new(AtomicUsize::new(value.len())),
50            buf: Arc::new(SyncUnsafeCell::new(value)),
51            write: Arc::new(AtomicUsize::new(0)),
52            lock: Arc::new(AtomicBool::new(false)),
53        }
54    }
55}
56
57impl<T: Clone + Send + Debug> MagicOrb<T> {
58    pub fn new(size: usize, default_val: T) -> Self {
59        if size == 0 {
60            panic!("Can't crate MagicOrb with buffer size 0");
61        }
62        MagicOrb {
63            buf: Arc::new(SyncUnsafeCell::new(vec![default_val; size])),
64            write: Arc::new(AtomicUsize::new(0)),
65            lock: Arc::new(AtomicBool::new(false)),
66            len: Arc::new(AtomicUsize::new(size)),
67            capacity: Arc::new(AtomicUsize::new(size)),
68        }
69    }
70
71    pub fn push_slice_overwrite(&self, data: &[T]) {
72        let occupied = self.capacity.load(Ordering::Acquire);
73        let data = if data.len() > occupied {
74            &data[data.len() - occupied..]
75        } else {
76            data
77        };
78
79        self.take_lock();
80        {
81            // SAFETY: Lock prevents aliasing &mut T
82            // Guarantees required: should be between self.take_lock() and self.return_lock()
83            let buf = unsafe { self.buf.get().as_mut().unwrap() };
84            let write = self.write.load(Ordering::Relaxed);
85
86            if data.len() + write <= occupied {
87                buf[write..write + data.len()].clone_from_slice(data);
88                self.write
89                    .store((write + data.len()) % occupied, Ordering::Relaxed);
90            } else {
91                let first_len = occupied - write;
92                buf[write..].clone_from_slice(&data[..first_len]);
93                buf[..data.len() - first_len].clone_from_slice(&data[first_len..]);
94                self.write.store(data.len() - first_len, Ordering::Relaxed);
95            }
96
97            let capacity = self.capacity();
98
99            _ = self
100                .len
101                .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |cur_val| {
102                    if cur_val < capacity {
103                        Some(min(cur_val + data.len(), capacity))
104                    } else {
105                        None
106                    }
107                });
108        }
109        self.return_lock();
110    }
111
112    pub fn get_contiguous(&self) -> Vec<T> {
113        let capacity = self.capacity();
114        let mut ret = Vec::with_capacity(capacity);
115
116        self.take_lock();
117        if self.is_empty() {
118            self.return_lock();
119            return ret;
120        }
121
122        let vacant_amount = {
123            // SAFETY: Lock prevents aliasing &mut T.
124            // Guarantees required: should be between self.take_lock() and self.return_lock()
125            let buf = unsafe { self.buf.get().as_mut().unwrap() };
126            let write = self.write.load(Ordering::Relaxed);
127            let vacant_amount = capacity - self.len();
128            ret.extend_from_slice(&buf[write..]);
129            ret.extend_from_slice(&buf[..write]);
130            vacant_amount
131        };
132        self.return_lock();
133
134        if vacant_amount > 0 {
135            ret = ret.split_off(vacant_amount);
136        }
137
138        ret
139    }
140
141    pub fn pop_back(&self) {
142        self.take_lock();
143        {
144            if self
145                .len
146                .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |cur_val| {
147                    if cur_val == 0 {
148                        None
149                    } else {
150                        Some(cur_val - 1)
151                    }
152                })
153                .is_ok()
154            {
155                let max_len = self.capacity();
156                _ = self
157                    .write
158                    .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |idx| {
159                        Some((idx + max_len - 1) % max_len)
160                    });
161            }
162        }
163        self.return_lock();
164    }
165
166    pub fn capacity(&self) -> usize {
167        self.capacity.load(Ordering::Relaxed)
168    }
169
170    pub fn len(&self) -> usize {
171        self.len.load(Ordering::Relaxed)
172    }
173
174    pub fn is_empty(&self) -> bool {
175        if self.len() == 0 {
176            return true;
177        }
178        false
179    }
180
181    fn take_lock(&self) {
182        while self
183            .lock
184            .compare_exchange(false, true, Ordering::Release, Ordering::Acquire)
185            .is_err()
186        {
187            std::hint::spin_loop();
188        }
189    }
190
191    fn return_lock(&self) {
192        self.lock.store(false, Ordering::Release);
193    }
194}