polars_utils/
unique_id.rs

1use std::any::Any;
2use std::fmt::LowerHex;
3use std::sync::Arc;
4
5/// Unique identifier potentially backed by an `Arc` address to protect against collisions.
6///
7/// Note that a serialization roundtrip will force this to become a [`UniqueId::Plain`] variant.
8#[derive(Clone)]
9pub enum UniqueId {
10    /// ID that derives from the memory address of an `Arc` to protect against collisions.
11    /// Note that it internally stores as a `dyn Any` to allow it to re-use an existing `Arc` holding
12    /// any type.
13    MemoryRef(Arc<dyn Any + Send + Sync>),
14
15    /// Stores a plain `usize`. Unlike the `MemoryRef` variant, there is no internal protection against
16    /// collisions - this must handled separately.
17    ///
18    /// Note: This repr may also be constructed as the result of a serialization round-trip.
19    Plain(usize),
20}
21
22impl UniqueId {
23    #[inline]
24    pub fn to_usize(&self) -> usize {
25        match self {
26            Self::MemoryRef(v) => Arc::as_ptr(v) as *const () as usize,
27            Self::Plain(v) => *v,
28        }
29    }
30
31    /// Use an existing `Arc<T>` as backing for an ID.
32    pub fn from_arc<T: Any + Send + Sync>(arc: Arc<T>) -> Self {
33        Self::MemoryRef(arc.clone())
34    }
35
36    /// Downcasts to a concrete Arc type. Returns None `Self` is `Plain`.
37    ///
38    /// # Panics
39    /// On a debug build, panics if `Self` is an `Arc` but does not contain `T`. On a release build
40    /// this will instead `None`.
41    pub fn downcast_arc<T: Any>(self) -> Option<Arc<T>> {
42        match self {
43            Self::MemoryRef(inner) => {
44                // Note, ref type here must match exactly with T.
45                let v: &dyn Any = inner.as_ref();
46
47                if v.type_id() != std::any::TypeId::of::<T>() {
48                    if cfg!(debug_assertions) {
49                        panic!("invalid downcast of UniqueId")
50                    } else {
51                        // Just return None on release.
52                        return None;
53                    }
54                }
55
56                // Safety: Type IDs checked above.
57                let ptr: *const dyn Any = Arc::into_raw(inner);
58                let ptr: *const T = ptr as _;
59                Some(unsafe { Arc::from_raw(ptr) })
60            },
61
62            Self::Plain(_) => None,
63        }
64    }
65}
66
67impl Default for UniqueId {
68    fn default() -> Self {
69        Self::MemoryRef(Arc::new(()))
70    }
71}
72
73impl PartialEq for UniqueId {
74    fn eq(&self, other: &Self) -> bool {
75        self.to_usize() == other.to_usize()
76    }
77}
78
79impl Eq for UniqueId {}
80
81impl PartialOrd for UniqueId {
82    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
83        Some(Ord::cmp(self, other))
84    }
85}
86
87impl Ord for UniqueId {
88    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
89        Ord::cmp(&self.to_usize(), &other.to_usize())
90    }
91}
92
93impl std::hash::Hash for UniqueId {
94    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
95        self.to_usize().hash(state)
96    }
97}
98
99impl std::fmt::Display for UniqueId {
100    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
101        std::fmt::Display::fmt(&self.to_usize(), f)
102    }
103}
104
105impl std::fmt::Debug for UniqueId {
106    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
107        use UniqueId::*;
108
109        write!(
110            f,
111            "UniqueId::{}({})",
112            match self {
113                MemoryRef(_) => "MemoryRef",
114                Plain(_) => "Plain",
115            },
116            self.to_usize()
117        )
118    }
119}
120
121impl LowerHex for UniqueId {
122    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
123        LowerHex::fmt(&self.to_usize(), f)
124    }
125}
126
127#[cfg(feature = "serde")]
128mod _serde_impl {
129    use super::UniqueId;
130
131    impl serde::ser::Serialize for UniqueId {
132        fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
133        where
134            S: serde::Serializer,
135        {
136            usize::serialize(&self.to_usize(), serializer)
137        }
138    }
139
140    impl<'de> serde::de::Deserialize<'de> for UniqueId {
141        fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
142        where
143            D: serde::Deserializer<'de>,
144        {
145            usize::deserialize(deserializer).map(Self::Plain)
146        }
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use std::any::Any;
153    use std::sync::Arc;
154
155    use super::UniqueId;
156
157    #[test]
158    fn test_unique_id() {
159        let id = UniqueId::default();
160
161        assert!(matches!(id, UniqueId::MemoryRef(_)));
162
163        assert_eq!(id, id);
164        assert_ne!(id, UniqueId::default());
165
166        assert_eq!(id, UniqueId::Plain(id.to_usize()));
167
168        // Following code shows the memory layout
169        let UniqueId::MemoryRef(arc_ref) = &id else {
170            unreachable!()
171        };
172        let inner_ref: &dyn Any = arc_ref.as_ref();
173
174        assert_eq!(std::mem::size_of_val(inner_ref), 0);
175        assert_eq!(std::mem::size_of::<Arc<dyn Any>>(), 16);
176
177        assert_eq!(
178            Arc::as_ptr(arc_ref) as *const () as usize,
179            inner_ref as *const _ as *const () as usize,
180        );
181    }
182
183    #[test]
184    fn test_unique_id_downcast() {
185        let id = UniqueId::default();
186        let _: Arc<()> = id.downcast_arc().unwrap();
187
188        let inner: Arc<usize> = Arc::new(37);
189        let id = UniqueId::from_arc(inner);
190
191        let out: Arc<usize> = id.downcast_arc().unwrap();
192        assert_eq!(*out.as_ref(), 37);
193    }
194}
polars_utils/unique_id.rs

polars_utils/
unique_id.rs
