vortex_session/
registry.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Many session types use a registry of objects that can be looked up by name to construct
5//! contexts. This module provides a generic registry type for that purpose.
6
7use std::cmp::Ordering;
8use std::fmt;
9use std::fmt::Debug;
10use std::fmt::Display;
11use std::fmt::Formatter;
12use std::hash::Hash;
13use std::ops::Deref;
14use std::sync::Arc;
15use std::sync::LazyLock;
16use std::sync::OnceLock;
17
18use lasso::Spur;
19use lasso::ThreadedRodeo;
20use parking_lot::RwLock;
21use vortex_error::VortexExpect;
22use vortex_utils::aliases::DefaultHashBuilder;
23use vortex_utils::aliases::dash_map::DashMap;
24use vortex_utils::aliases::hash_set::HashSet;
25
26/// Global string interner for [`Id`] values.
27static INTERNER: LazyLock<ThreadedRodeo<Spur, DefaultHashBuilder>> =
28    LazyLock::new(|| ThreadedRodeo::with_hasher(DefaultHashBuilder::default()));
29
30/// A lightweight, copyable identifier backed by a global string interner.
31///
32/// Used for array encoding IDs, scalar function IDs, layout IDs, and similar
33/// globally-unique string identifiers throughout Vortex. Equality and hashing
34/// are O(1) symbol comparisons.
35#[derive(Clone, Copy, PartialEq, Eq, Hash)]
36pub struct Id(Spur);
37
38impl Id {
39    /// Intern a string and return its `Id`.
40    pub fn new(s: &str) -> Self {
41        Self(INTERNER.get_or_intern(s))
42    }
43
44    /// Intern a string and return its `Id`.
45    pub fn new_static(s: &'static str) -> Self {
46        Self(INTERNER.get_or_intern_static(s))
47    }
48
49    /// Returns the interned string.
50    pub fn as_str(&self) -> &str {
51        let s = INTERNER.resolve(&self.0);
52        // SAFETY: INTERNER is 'static and its arena is append-only, so resolved string
53        // pointers are stable for the lifetime of the program.
54        unsafe { &*(s as *const str) }
55    }
56}
57
58impl From<&str> for Id {
59    #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")]
60    fn from(s: &str) -> Self {
61        Self::new(s)
62    }
63}
64
65impl Display for Id {
66    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
67        f.write_str(self.as_str())
68    }
69}
70
71impl Debug for Id {
72    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
73        write!(f, "Id(\"{}\")", self.as_str())
74    }
75}
76
77impl PartialOrd for Id {
78    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
79        Some(self.cmp(other))
80    }
81}
82
83impl Ord for Id {
84    fn cmp(&self, other: &Self) -> Ordering {
85        self.as_str().cmp(other.as_str())
86    }
87}
88
89impl AsRef<str> for Id {
90    fn as_ref(&self) -> &str {
91        self.as_str()
92    }
93}
94
95impl PartialEq<&Id> for Id {
96    fn eq(&self, other: &&Id) -> bool {
97        self == *other
98    }
99}
100
101impl PartialEq<Id> for &Id {
102    fn eq(&self, other: &Id) -> bool {
103        *self == other
104    }
105}
106
107/// A lazily-initialized, cached [`Id`] for use as a `static`.
108///
109/// Avoids repeated interner write-lock acquisition by storing the interned [`Id`]
110/// on first access and returning the cached copy on all subsequent calls.
111///
112/// # Example
113///
114/// ```
115/// use vortex_session::registry::{CachedId, Id};
116///
117/// static MY_ID: CachedId = CachedId::new("my.encoding");
118///
119/// fn get_id() -> Id {
120///     *MY_ID
121/// }
122/// ```
123pub struct CachedId {
124    s: &'static str,
125    cached: OnceLock<Id>,
126}
127
128impl CachedId {
129    /// Create a new `CachedId` that will intern `s` on first access.
130    pub const fn new(s: &'static str) -> Self {
131        Self {
132            s,
133            cached: OnceLock::new(),
134        }
135    }
136}
137
138impl Deref for CachedId {
139    type Target = Id;
140
141    #[expect(
142        clippy::disallowed_methods,
143        reason = "CachedId interns its static id once here"
144    )]
145    fn deref(&self) -> &Id {
146        self.cached.get_or_init(|| Id::new_static(self.s))
147    }
148}
149
150/// A registry of items that are keyed by a string identifier.
151#[derive(Clone, Debug)]
152pub struct Registry<T>(Arc<DashMap<Id, T>>);
153
154impl<T> Default for Registry<T> {
155    fn default() -> Self {
156        Self(Default::default())
157    }
158}
159
160impl<T: Clone> Registry<T> {
161    pub fn empty() -> Self {
162        Self(Default::default())
163    }
164
165    /// List the IDs in the registry.
166    pub fn ids(&self) -> impl Iterator<Item = Id> + '_ {
167        self.0.iter().map(|i| *i.key())
168    }
169
170    /// List the items in the registry.
171    pub fn items(&self) -> impl Iterator<Item = T> + '_ {
172        self.0.iter().map(|i| i.value().clone())
173    }
174
175    /// Return the items with the given IDs.
176    pub fn find_many<'a>(
177        &self,
178        ids: impl IntoIterator<Item = &'a Id>,
179    ) -> impl Iterator<Item = Option<impl Deref<Target = T>>> {
180        ids.into_iter().map(|id| self.0.get(id))
181    }
182
183    /// Find the item with the given ID.
184    pub fn find(&self, id: &Id) -> Option<T> {
185        self.0.get(id).as_deref().cloned()
186    }
187
188    /// Register a new item, replacing any existing item with the same ID.
189    pub fn register(&self, id: impl Into<Id>, item: impl Into<T>) {
190        self.0.insert(id.into(), item.into());
191    }
192
193    /// Register a new item, replacing any existing item with the same ID, and return self for
194    pub fn with(self, id: impl Into<Id>, item: impl Into<T>) -> Self {
195        self.register(id, item.into());
196        self
197    }
198}
199
200/// A [`ReadContext`] holds a set of interned IDs for use during deserialization, mapping
201/// u16 indices to IDs.
202#[derive(Clone, Debug)]
203pub struct ReadContext {
204    ids: Arc<[Id]>,
205}
206
207impl ReadContext {
208    /// Create a context with the given initial IDs.
209    pub fn new(ids: impl Into<Arc<[Id]>>) -> Self {
210        Self { ids: ids.into() }
211    }
212
213    /// Resolve an interned ID by its index.
214    pub fn resolve(&self, idx: u16) -> Option<Id> {
215        self.ids.get(idx as usize).cloned()
216    }
217
218    pub fn ids(&self) -> &[Id] {
219        &self.ids
220    }
221}
222
223/// A [`Context`] holds a set of interned IDs for use during serialization/deserialization, mapping
224/// IDs to u16 indices.
225///
226/// ## Upcoming Changes
227///
228/// 1. This object holds an Arc of RwLock internally because we need concurrent access from the
229///    layout writer code path. We should update SegmentSink to take an Array rather than
230///    ByteBuffer such that serializing arrays is done sequentially.
231/// 2. The name is terrible. `Interner` is better, but I want to minimize breakage for now.
232#[derive(Clone, Debug)]
233pub struct Context {
234    // TODO(ngates): it's a long story, but if we make SegmentSink and SegmentSource take an
235    //  enum of Segment { Array, DType, Buffer } then we don't actually need a mutable context
236    //  in the LayoutWriter, therefore we don't need a RwLock here and everyone is happier.
237    ids: Arc<RwLock<Vec<Id>>>,
238    // Optional set used to filter the permissible interned IDs.
239    valid_ids: Option<Arc<HashSet<Id>>>,
240}
241
242impl Default for Context {
243    fn default() -> Self {
244        Self {
245            ids: Arc::new(RwLock::new(Vec::new())),
246            valid_ids: None,
247        }
248    }
249}
250
251impl Context {
252    /// Create a context with the given initial IDs.
253    pub fn new(ids: Vec<Id>) -> Self {
254        Self {
255            ids: Arc::new(RwLock::new(ids)),
256            valid_ids: None,
257        }
258    }
259
260    /// Create an empty context.
261    pub fn empty() -> Self {
262        Self::default()
263    }
264
265    /// Restrict the permissible set of interned IDs.
266    pub fn with_valid_ids(mut self, ids: impl IntoIterator<Item = Id>) -> Self {
267        self.valid_ids = Some(Arc::new(ids.into_iter().collect()));
268        self
269    }
270
271    /// Intern an ID, returning its index.
272    pub fn intern(&self, id: &Id) -> Option<u16> {
273        if let Some(valid_ids) = &self.valid_ids
274            && !valid_ids.contains(id)
275        {
276            // ID is not valid, cannot intern.
277            return None;
278        }
279
280        let mut ids = self.ids.write();
281        if let Some(idx) = ids.iter().position(|e| e == id) {
282            return Some(u16::try_from(idx).vortex_expect("Cannot have more than u16::MAX items"));
283        }
284
285        let idx = ids.len();
286        assert!(
287            idx < u16::MAX as usize,
288            "Cannot have more than u16::MAX items"
289        );
290        ids.push(*id);
291        Some(u16::try_from(idx).vortex_expect("checked already"))
292    }
293
294    /// Get the list of interned IDs.
295    pub fn to_ids(&self) -> Vec<Id> {
296        self.ids.read().clone()
297    }
298}
299
300#[cfg(test)]
301mod tests {
302    use super::CachedId;
303    use super::Context;
304
305    static VALID: CachedId = CachedId::new("vortex.test.valid");
306    static INVALID: CachedId = CachedId::new("vortex.test.invalid");
307
308    #[test]
309    fn context_filters_interned_ids() {
310        let valid = *VALID;
311        let invalid = *INVALID;
312        let context = Context::empty().with_valid_ids([valid]);
313
314        assert_eq!(context.intern(&valid), Some(0));
315        assert_eq!(context.intern(&valid), Some(0));
316        assert_eq!(context.intern(&invalid), None);
317        assert_eq!(context.to_ids(), [valid]);
318    }
319}
vortex_session/registry.rs

vortex_session/
registry.rs
