triomphe/
arc_union.rs

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