triomphe/
arc_union.rs

1use core::fmt;
2use core::marker::PhantomData;
3use core::panic::{RefUnwindSafe, UnwindSafe};
4use core::ptr;
5
6use super::{Arc, ArcBorrow};
7
8/// A tagged union that can represent `Arc<A>` or `Arc<B>` while only consuming a
9/// single word. The type is also `NonNull`, and thus can be stored in an Option
10/// without increasing size.
11///
12/// This is functionally equivalent to
13/// `enum ArcUnion<A, B> { First(Arc<A>), Second(Arc<B>)` but only takes up
14/// up a single word of stack space.
15///
16/// This could probably be extended to support four types if necessary.
17pub struct ArcUnion<A, B> {
18    p: ptr::NonNull<()>,
19    phantom_a: PhantomData<A>,
20    phantom_b: PhantomData<B>,
21}
22
23unsafe impl<A: Sync + Send, B: Send + Sync> Send for ArcUnion<A, B> {}
24unsafe impl<A: Sync + Send, B: Send + Sync> Sync for ArcUnion<A, B> {}
25
26impl<A: RefUnwindSafe, B: RefUnwindSafe> UnwindSafe for ArcUnion<A, B> {}
27
28impl<A: PartialEq, B: PartialEq> PartialEq for ArcUnion<A, B> {
29    fn eq(&self, other: &Self) -> bool {
30        use crate::ArcUnionBorrow::*;
31        match (self.borrow(), other.borrow()) {
32            (First(x), First(y)) => x == y,
33            (Second(x), Second(y)) => x == y,
34            (_, _) => false,
35        }
36    }
37}
38
39/// This represents a borrow of an `ArcUnion`.
40#[derive(Debug)]
41pub enum ArcUnionBorrow<'a, A: 'a, B: 'a> {
42    First(ArcBorrow<'a, A>),
43    Second(ArcBorrow<'a, B>),
44}
45
46impl<A, B> ArcUnion<A, B> {
47    unsafe fn new(ptr: *mut ()) -> Self {
48        ArcUnion {
49            p: ptr::NonNull::new_unchecked(ptr),
50            phantom_a: PhantomData,
51            phantom_b: PhantomData,
52        }
53    }
54
55    /// Returns true if the two values are pointer-equal.
56    #[inline]
57    pub fn ptr_eq(this: &Self, other: &Self) -> bool {
58        this.p == other.p
59    }
60
61    /// Reference count.
62    #[inline]
63    pub fn strong_count(this: &Self) -> usize {
64        ArcUnionBorrow::strong_count(&this.borrow())
65    }
66
67    /// Returns an enum representing a borrow of either A or B.
68    pub fn borrow(&self) -> ArcUnionBorrow<'_, A, B> {
69        if self.is_first() {
70            let ptr = self.p.as_ptr() as *const A;
71            let borrow = unsafe { ArcBorrow::from_ptr(ptr) };
72            ArcUnionBorrow::First(borrow)
73        } else {
74            let ptr = ((self.p.as_ptr() as usize) & !0x1) as *const B;
75            let borrow = unsafe { ArcBorrow::from_ptr(ptr) };
76            ArcUnionBorrow::Second(borrow)
77        }
78    }
79
80    /// Creates an `ArcUnion` from an instance of the first type.
81    #[inline]
82    pub fn from_first(other: Arc<A>) -> Self {
83        unsafe { Self::new(Arc::into_raw(other) as *mut _) }
84    }
85
86    /// Creates an `ArcUnion` from an instance of the second type.
87    #[inline]
88    pub fn from_second(other: Arc<B>) -> Self {
89        unsafe { Self::new(((Arc::into_raw(other) as usize) | 0x1) as *mut _) }
90    }
91
92    /// Returns true if this `ArcUnion` contains the first type.
93    #[inline]
94    pub fn is_first(&self) -> bool {
95        self.p.as_ptr() as usize & 0x1 == 0
96    }
97
98    /// Returns true if this `ArcUnion` contains the second type.
99    #[inline]
100    pub fn is_second(&self) -> bool {
101        !self.is_first()
102    }
103
104    /// Returns a borrow of the first type if applicable, otherwise `None`.
105    pub fn as_first(&self) -> Option<ArcBorrow<'_, A>> {
106        match self.borrow() {
107            ArcUnionBorrow::First(x) => Some(x),
108            ArcUnionBorrow::Second(_) => None,
109        }
110    }
111
112    /// Returns a borrow of the second type if applicable, otherwise None.
113    pub fn as_second(&self) -> Option<ArcBorrow<'_, B>> {
114        match self.borrow() {
115            ArcUnionBorrow::First(_) => None,
116            ArcUnionBorrow::Second(x) => Some(x),
117        }
118    }
119}
120
121impl<A, B> Clone for ArcUnion<A, B> {
122    fn clone(&self) -> Self {
123        match self.borrow() {
124            ArcUnionBorrow::First(x) => ArcUnion::from_first(x.clone_arc()),
125            ArcUnionBorrow::Second(x) => ArcUnion::from_second(x.clone_arc()),
126        }
127    }
128}
129
130impl<A, B> Drop for ArcUnion<A, B> {
131    fn drop(&mut self) {
132        match self.borrow() {
133            ArcUnionBorrow::First(x) => unsafe {
134                let _ = Arc::from_raw(&*x);
135            },
136            ArcUnionBorrow::Second(x) => unsafe {
137                let _ = Arc::from_raw(&*x);
138            },
139        }
140    }
141}
142
143impl<A: fmt::Debug, B: fmt::Debug> fmt::Debug for ArcUnion<A, B> {
144    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
145        fmt::Debug::fmt(&self.borrow(), f)
146    }
147}
148
149impl<'a, A, B> ArcUnionBorrow<'a, A, B> {
150    /// The reference count of this `Arc`.
151    ///
152    /// The number does not include borrowed pointers,
153    /// or temporary `Arc` pointers created with functions like
154    /// [`ArcBorrow::with_arc`].
155    ///
156    /// The function is called `strong_count` to mirror `std::sync::Arc::strong_count`,
157    /// however `triomphe::Arc` does not support weak references.
158    #[inline]
159    pub fn strong_count(this: &Self) -> usize {
160        match this {
161            ArcUnionBorrow::First(arc) => ArcBorrow::strong_count(arc),
162            ArcUnionBorrow::Second(arc) => ArcBorrow::strong_count(arc),
163        }
164    }
165}
triomphe/arc_union.rs

triomphe/
arc_union.rs
