utils_atomics/cell.rs
1#[cfg(feature = "alloc_api")]
2use alloc::alloc::*;
3#[cfg(feature = "alloc_api")]
4use core::mem::ManuallyDrop;
5
6use alloc::boxed::Box;
7use core::sync::atomic::{AtomicPtr, Ordering};
8use docfg::docfg;
9
10/// An atomic cell that can be safely shared between threads and can contain an optional value.
11///
12/// `AtomicCell` provides methods to store, replace, and take values atomically, ensuring safe access
13/// and modification across multiple threads.
14///
15/// # Example
16///
17/// ```rust
18/// use utils_atomics::AtomicCell;
19///
20/// let mut atomic_cell = AtomicCell::new(Some(42));
21///
22/// std::thread::scope(|s| {
23/// // Spawn a thread that replaces the value inside the AtomicCell
24/// s.spawn(|| {
25/// let prev_value = atomic_cell.replace(Some(24));
26/// assert_eq!(prev_value, Some(42));
27/// });
28/// });
29///
30/// // Check that the value was replaced
31/// assert_eq!(atomic_cell.get_mut().copied(), Some(24));
32/// ```
33#[derive(Debug)]
34pub struct AtomicCell<T, #[cfg(feature = "alloc_api")] A: Allocator = Global> {
35 inner: AtomicPtr<T>,
36 #[cfg(feature = "alloc_api")]
37 alloc: ManuallyDrop<A>,
38}
39
40#[docfg(feature = "alloc_api")]
41impl<T, A: Allocator> AtomicCell<T, A> {
42 /// Constructs a new `AtomicCell` containing an optional value t and an allocator alloc.
43 ///
44 /// If the value is `Some(x)`, it is boxed using the allocator.
45 /// If the value is `None`, an empty `AtomicCell` is created with the allocator.
46 ///
47 /// # Example
48 /// ```rust
49 /// #![feature(allocator_api)]
50 ///
51 /// use utils_atomics::AtomicCell;
52 /// use std::alloc::System;
53 ///
54 /// let atomic_cell = AtomicCell::<i32, _>::new_in(Some(42), System);
55 /// ```
56 #[inline]
57 pub fn new_in(t: impl Into<Option<T>>, alloc: A) -> Self {
58 Self::new_boxed_in(match t.into() {
59 Some(x) => Ok(Box::new_in(x, alloc)),
60 None => Err(alloc),
61 })
62 }
63
64 /// Constructs a new `AtomicCell` from a boxed value or an allocator.
65 ///
66 /// If the input is `Ok(t)`, the `AtomicCell` contains the boxed value, and the allocator is extracted from the box.
67 /// If the input is `Err(alloc)`, an empty `AtomicCell` is created with the allocator.
68 ///
69 /// # Example
70 /// ```rust
71 /// #![feature(allocator_api)]
72 /// extern crate alloc;
73 ///
74 /// use utils_atomics::AtomicCell;
75 /// use std::alloc::System;
76 /// use alloc::boxed::Box;
77 ///
78 /// let atomic_cell = AtomicCell::new_boxed_in(Ok(Box::new_in(42, System)));
79 /// ```
80 #[inline]
81 pub fn new_boxed_in(t: Result<Box<T, A>, A>) -> Self {
82 match t {
83 Ok(t) => {
84 let (ptr, alloc) = Box::into_raw_with_allocator(t);
85 Self {
86 inner: AtomicPtr::new(ptr),
87 alloc: ManuallyDrop::new(alloc),
88 }
89 }
90 Err(alloc) => Self {
91 inner: AtomicPtr::new(core::ptr::null_mut()),
92 alloc: ManuallyDrop::new(alloc),
93 },
94 }
95 }
96
97 /// Returns a reference to the allocator associated with the `AtomicCell`.
98 ///
99 /// # Example
100 /// ```rust
101 /// #![feature(allocator_api)]
102 ///
103 /// use utils_atomics::AtomicCell;
104 /// use std::alloc::System;
105 ///
106 /// let atomic_cell = AtomicCell::<i32, System>::new_in(Some(42), System);
107 /// let allocator = atomic_cell.allocator();
108 /// ```
109 #[inline]
110 pub fn allocator(&self) -> &A {
111 core::ops::Deref::deref(&self.alloc)
112 }
113
114 /// Takes the value out of the `AtomicCell`, leaving it empty.
115 ///
116 /// Returns an optional boxed value with a reference to the allocator.
117 /// If the `AtomicCell` is empty, returns `None`.
118 ///
119 /// # Example
120 /// ```rust
121 /// #![feature(allocator_api)]
122 ///
123 /// use utils_atomics::AtomicCell;
124 /// use std::alloc::System;
125 ///
126 /// let atomic_cell = AtomicCell::new_in(Some(42), System);
127 /// let taken_value = atomic_cell.take_in();
128 /// assert_eq!(taken_value, Some(Box::new_in(42, &System)))
129 /// ```
130 #[inline]
131 pub fn take_in(&self) -> Option<Box<T, &A>> {
132 self.replace_in(None)
133 }
134
135 /// Replaces the value inside the `AtomicCell` with a new optional value.
136 ///
137 /// If the new value is `Some(new)`, it is boxed using the allocator.
138 /// If the new value is `None`, the `AtomicCell` is emptied.
139 ///
140 /// Returns the old value as an optional boxed value with a reference to the allocator.
141 /// If the `AtomicCell` was empty, returns None.
142 ///
143 /// # Example
144 /// ```rust
145 /// #![feature(allocator_api)]
146
147 ///
148 /// use utils_atomics::AtomicCell;
149 /// use std::alloc::System;
150 ///
151 /// let atomic_cell = AtomicCell::new_in(Some(42), System);
152 /// let old_value = atomic_cell.replace_in(Some(24));
153 /// assert_eq!(old_value, Some(Box::new_in(42, atomic_cell.allocator())));
154 /// assert_eq!(atomic_cell.take(), Some(24));
155 /// ```
156 #[inline]
157 pub fn replace_in(&self, new: impl Into<Option<T>>) -> Option<Box<T, &A>> {
158 let new = match new.into() {
159 Some(new) => Box::into_raw(Box::new_in(new, core::ops::Deref::deref(&self.alloc))),
160 None => core::ptr::null_mut(),
161 };
162
163 let prev = self.inner.swap(new, Ordering::AcqRel);
164 if prev.is_null() {
165 return None;
166 }
167
168 return unsafe { Some(Box::from_raw_in(prev, core::ops::Deref::deref(&self.alloc))) };
169 }
170}
171
172impl<T> AtomicCell<T> {
173 /// Constructs a new `AtomicCell` containing an optional value `t`.
174 ///
175 /// # Example
176 ///
177 /// ```rust
178 /// use utils_atomics::AtomicCell;
179 ///
180 /// let atomic_cell = AtomicCell::<i32>::new(Some(42));
181 /// ```
182 #[inline]
183 pub fn new(t: impl Into<Option<T>>) -> Self {
184 Self::new_boxed(t.into().map(Box::new))
185 }
186
187 /// Constructs a new `AtomicCell` from an optional boxed value `t`.
188 ///
189 /// # Example
190 ///
191 /// ```rust
192 /// extern crate alloc;
193 ///
194 /// use utils_atomics::AtomicCell;
195 /// use alloc::boxed::Box;
196 ///
197 /// let atomic_cell = AtomicCell::new_boxed(Some(Box::new(42)));
198 /// ```
199 #[inline]
200 pub fn new_boxed(t: impl Into<Option<Box<T>>>) -> Self {
201 match t.into() {
202 Some(t) => Self {
203 inner: AtomicPtr::new(Box::into_raw(t)),
204 #[cfg(feature = "alloc_api")]
205 alloc: ManuallyDrop::new(Global),
206 },
207 None => Self {
208 inner: AtomicPtr::new(core::ptr::null_mut()),
209 #[cfg(feature = "alloc_api")]
210 alloc: ManuallyDrop::new(Global),
211 },
212 }
213 }
214
215 /// Replaces the value inside the `AtomicCell` with a new optional value `new`.
216 /// Returns the old value as an optional value. If the `AtomicCell` was empty, returns `None`.
217 ///
218 /// # Example
219 ///
220 /// ```rust
221 /// use utils_atomics::AtomicCell;
222 ///
223 /// let atomic_cell = AtomicCell::<i32>::new(Some(42));
224 /// let old_value = atomic_cell.replace(Some(24));
225 /// ```
226 #[inline]
227 pub fn replace(&self, new: impl Into<Option<T>>) -> Option<T> {
228 self.replace_boxed(new.into().map(Box::new)).map(|x| *x)
229 }
230
231 /// Replaces the value inside the `AtomicCell` with a new optional boxed value `new`.
232 /// Returns the old value as an optional boxed value. If the `AtomicCell` was empty, returns `None`.
233 ///
234 /// # Example
235 ///
236 /// ```rust
237 /// extern crate alloc;
238 ///
239 /// use utils_atomics::AtomicCell;
240 /// use alloc::boxed::Box;
241 ///
242 /// let atomic_cell = AtomicCell::new_boxed(Some(Box::new(42)));
243 /// let old_value = atomic_cell.replace_boxed(Some(Box::new(24)));
244 /// ```
245 #[inline]
246 pub fn replace_boxed(&self, new: impl Into<Option<Box<T>>>) -> Option<Box<T>> {
247 let new = match new.into() {
248 Some(new) => Box::into_raw(new),
249 None => core::ptr::null_mut(),
250 };
251
252 let prev = self.inner.swap(new, Ordering::AcqRel);
253 if prev.is_null() {
254 return None;
255 }
256 return unsafe { Some(Box::from_raw(prev)) };
257 }
258
259 /// Takes the value out of the `AtomicCell`, leaving it empty.
260 /// Returns an optional boxed value. If the `AtomicCell` is empty, returns `None`.
261 ///
262 /// # Example
263 ///
264 /// ```rust
265 /// use utils_atomics::AtomicCell;
266 ///
267 /// let atomic_cell = AtomicCell::new_boxed(Some(Box::new(42)));
268 /// let taken_value = atomic_cell.take_boxed();
269 /// ```
270 #[inline]
271 pub fn take_boxed(&self) -> Option<Box<T>> {
272 self.replace_boxed(None)
273 }
274}
275
276cfg_if::cfg_if! {
277 if #[cfg(feature = "alloc_api")] {
278 impl<T, A: Allocator> AtomicCell<T, A> {
279 /// Takes the value out of the `AtomicCell`, leaving it empty.
280 /// Returns an optional value. If the `AtomicCell` is empty, returns `None`.
281 ///
282 /// # Examples
283 ///
284 /// ```
285 /// use utils_atomics::AtomicCell;
286 ///
287 /// let atomic_cell = AtomicCell::new(Some(42));
288 /// assert_eq!(atomic_cell.take(), Some(42));
289 /// assert_eq!(atomic_cell.take(), None);
290 /// ```
291 #[inline]
292 pub fn take(&self) -> Option<T> {
293 self.take_in().map(|x| *x)
294 }
295
296 /// Returns a mutable reference to the value inside the `AtomicCell`, if any.
297 /// If the `AtomicCell` is empty, returns `None`.
298 ///
299 /// # Examples
300 ///
301 /// ```
302 /// use utils_atomics::AtomicCell;
303 ///
304 /// let mut atomic_cell = AtomicCell::new(Some(42));
305 /// let value_ref = atomic_cell.get_mut().unwrap();
306 /// *value_ref = 24;
307 /// assert_eq!(*value_ref, 24);
308 /// ```
309 #[inline]
310 pub fn get_mut (&mut self) -> Option<&mut T> {
311 let ptr = *self.inner.get_mut();
312 if ptr.is_null() { return None }
313 return unsafe { Some(&mut *ptr) }
314 }
315
316 /// Returns `true` if the `AtomicCell` contains a value.
317 ///
318 /// # Examples
319 ///
320 /// ```
321 /// use utils_atomics::AtomicCell;
322 ///
323 /// let atomic_cell = AtomicCell::<i32>::new(Some(42));
324 /// assert!(atomic_cell.is_some());
325 /// ```
326 #[inline]
327 pub fn is_some (&self) -> bool {
328 return !self.is_none()
329 }
330
331 /// Returns `true` if the `AtomicCell` is empty.
332 ///
333 /// # Examples
334 ///
335 /// ```
336 /// use utils_atomics::AtomicCell;
337 ///
338 /// let atomic_cell = AtomicCell::<i32>::new(None);
339 /// assert!(atomic_cell.is_none());
340 /// ```
341 #[inline]
342 pub fn is_none (&self) -> bool {
343 return self.inner.load(Ordering::Relaxed).is_null()
344 }
345 }
346
347 impl<T, A: Allocator> Drop for AtomicCell<T, A> {
348 fn drop(&mut self) {
349 unsafe {
350 let ptr = *self.inner.get_mut();
351 if ptr.is_null() {
352 ManuallyDrop::drop(&mut self.alloc);
353 } else {
354 let _ = Box::from_raw_in(ptr, ManuallyDrop::take(&mut self.alloc));
355 }
356 }
357 }
358 }
359
360 unsafe impl<T: Send, A: Allocator + Send> Send for AtomicCell<T, A> {}
361 unsafe impl<T: Sync, A: Allocator + Sync> Sync for AtomicCell<T, A> {}
362 } else {
363 impl<T> AtomicCell<T> {
364 /// Takes the value out of the `AtomicCell`, leaving it empty.
365 /// Returns an optional value. If the `AtomicCell` is empty, returns `None`.
366 ///
367 /// # Examples
368 ///
369 /// ```
370 /// use utils_atomics::AtomicCell;
371 ///
372 /// let atomic_cell = AtomicCell::new(Some(42));
373 /// assert_eq!(atomic_cell.take(), Some(42));
374 /// assert_eq!(atomic_cell.take(), None);
375 /// ```
376 #[inline]
377 pub fn take(&self) -> Option<T> {
378 self.take_boxed().map(|x| *x)
379 }
380
381 /// Returns a mutable reference to the value inside the `AtomicCell`, if any.
382 /// If the `AtomicCell` is empty, returns `None`.
383 ///
384 /// # Examples
385 ///
386 /// ```
387 /// use utils_atomics::AtomicCell;
388 ///
389 /// let mut atomic_cell = AtomicCell::new(Some(42));
390 /// let value_ref = atomic_cell.get_mut().unwrap();
391 /// *value_ref = 24;
392 /// assert_eq!(*value_ref, 24);
393 /// ```
394 #[inline]
395 pub fn get_mut (&mut self) -> Option<&mut T> {
396 let ptr = *self.inner.get_mut();
397 if ptr.is_null() { return None }
398 return unsafe { Some(&mut *ptr) }
399 }
400
401 /// Returns `true` if the `AtomicCell` contains a value.
402 ///
403 /// # Examples
404 ///
405 /// ```
406 /// use utils_atomics::AtomicCell;
407 ///
408 /// let atomic_cell = AtomicCell::<i32>::new(Some(42));
409 /// assert!(atomic_cell.is_some());
410 /// ```
411 #[inline]
412 pub fn is_some (&self) -> bool {
413 return !self.is_none()
414 }
415
416 /// Returns `true` if the `AtomicCell` is empty.
417 ///
418 /// # Examples
419 ///
420 /// ```
421 /// use utils_atomics::AtomicCell;
422 ///
423 /// let atomic_cell = AtomicCell::<i32>::new(None);
424 /// assert!(atomic_cell.is_none());
425 /// ```
426 #[inline]
427 pub fn is_none (&self) -> bool {
428 return self.inner.load(Ordering::Relaxed).is_null()
429 }
430 }
431
432 impl<T> Drop for AtomicCell<T> {
433 fn drop(&mut self) {
434 unsafe {
435 let ptr = *self.inner.get_mut();
436 if !ptr.is_null() {
437 let _: Box<T> = Box::from_raw(ptr);
438 }
439 }
440 }
441 }
442
443 unsafe impl<T: Send> Send for AtomicCell<T> {}
444 unsafe impl<T: Sync> Sync for AtomicCell<T> {}
445 }
446}
447
448// Thanks ChatGPT!
449#[cfg(test)]
450mod tests {
451 use super::AtomicCell;
452
453 #[test]
454 fn create_and_take() {
455 let cell = AtomicCell::<i32>::new(Some(42));
456 assert_eq!(cell.take(), Some(42));
457 assert!(cell.is_none());
458 }
459
460 #[test]
461 fn create_empty_and_take() {
462 let cell = AtomicCell::<i32>::new(None);
463 assert!(cell.is_none());
464 assert_eq!(cell.take(), None);
465 }
466
467 #[test]
468 fn replace() {
469 let cell = AtomicCell::<i32>::new(Some(42));
470 let old_value = cell.replace(Some(13));
471 assert_eq!(old_value, Some(42));
472 assert_eq!(cell.take(), Some(13));
473 }
474
475 #[test]
476 fn replace_with_none() {
477 let cell = AtomicCell::<i32>::new(Some(42));
478 let old_value = cell.replace(None);
479 assert_eq!(old_value, Some(42));
480 assert!(cell.is_none());
481 }
482
483 #[test]
484 fn is_some_and_is_none() {
485 let cell = AtomicCell::<i32>::new(Some(42));
486 assert!(cell.is_some());
487 assert!(!cell.is_none());
488 cell.take();
489 assert!(!cell.is_some());
490 assert!(cell.is_none());
491 }
492
493 // Tests for custom allocator functionality
494 #[cfg(feature = "alloc_api")]
495 mod custom_allocator {
496 use super::*;
497 use alloc::alloc::Global;
498 use alloc::alloc::{Allocator, Layout};
499 use core::{alloc::AllocError, ptr::NonNull};
500
501 #[derive(Debug, Clone, Copy)]
502 pub struct DummyAllocator;
503
504 unsafe impl Allocator for DummyAllocator {
505 fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
506 Global.allocate(layout)
507 }
508
509 unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
510 Global.deallocate(ptr, layout)
511 }
512 }
513
514 #[test]
515 fn create_and_take_with_allocator() {
516 let cell = AtomicCell::<i32, DummyAllocator>::new_in(Some(42), DummyAllocator);
517 assert_eq!(cell.take_in().map(|x| *x), Some(42));
518 assert!(cell.is_none());
519 }
520
521 #[test]
522 fn create_empty_and_take_with_allocator() {
523 let cell = AtomicCell::<i32, DummyAllocator>::new_in(None, DummyAllocator);
524 assert!(cell.is_none());
525 assert_eq!(cell.take_in(), None);
526 }
527
528 #[test]
529 fn replace_with_allocator() {
530 let cell = AtomicCell::<i32, DummyAllocator>::new_in(Some(42), DummyAllocator);
531 let old_value = cell.replace_in(Some(13));
532 assert_eq!(old_value.map(|x| *x), Some(42));
533 assert_eq!(cell.take_in().map(|x| *x), Some(13));
534 }
535
536 #[test]
537 fn replace_with_none_with_allocator() {
538 let cell = AtomicCell::<i32, DummyAllocator>::new_in(Some(42), DummyAllocator);
539 let old_value = cell.replace_in(None);
540 assert_eq!(old_value, Some(Box::new_in(42, cell.allocator())));
541 assert!(cell.is_none());
542 }
543 }
544
545 #[cfg(all(feature = "std", miri))]
546 mod miri {
547 // Add other imports from previous tests
548 use crate::cell::AtomicCell;
549 use std::sync::Arc;
550 use std::thread;
551
552 const NUM_THREADS: usize = 10;
553 const NUM_ITERATIONS: usize = 1000;
554
555 fn stress_test_body(cell: &AtomicCell<Option<i32>>) {
556 for _ in 0..NUM_ITERATIONS {
557 cell.replace(Some(42));
558 cell.take();
559 }
560 }
561
562 #[test]
563 fn miri_stress_test() {
564 let cell = Arc::new(AtomicCell::new(Some(0)));
565 let mut handles = Vec::with_capacity(NUM_THREADS);
566
567 for _ in 0..NUM_THREADS {
568 let cloned_cell = Arc::clone(&cell);
569 let handle = thread::spawn(move || {
570 stress_test_body(&cloned_cell);
571 });
572 handles.push(handle);
573 }
574
575 for handle in handles {
576 handle.join().unwrap();
577 }
578
579 assert!(cell.is_none());
580 }
581 }
582}