Skip to main content

meta_ast/model/
ids.rs

1//! Newtyped ID types with thread-safe generation.
2//!
3//! `FileId`, `SymbolId`, `SnapshotId`, and `DataNodeId` are newtyped
4//! [`NonZeroU32`] values generated by [`IdGenerator<T>`], which wraps an
5//! `AtomicU32` for lock-free allocation across rayon worker threads.
6//!
7//! The generator starts at 1: 0 is the permanently invalid niche value.
8//! Because the inner type is `NonZeroU32`, the compiler niche-optimizes
9//! `Option<$name>` down to 4 bytes (the size of `$name` itself) instead of
10//! the 8 bytes an `Option<u32>` would cost. This benefits structures that
11//! store an optional id, e.g. `DataNode.symbol_id: Option<SymbolId>`.
12// life still has alot to teach you.
13use std::num::NonZeroU32;
14use std::sync::atomic::{AtomicU32, Ordering};
15
16use serde::{Deserialize, Serialize};
17
18macro_rules! define_id_type {
19    ($name:ident) => {
20        #[derive(
21            Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize,
22        )]
23        #[serde(transparent)]
24        pub struct $name(NonZeroU32);
25
26        impl $name {
27            /// Construct from a raw `u32`. Returns `None` for 0, which is the
28            /// permanently invalid niche value reserved for `Option<$name>`.
29            pub fn new(val: u32) -> Option<Self> {
30                NonZeroU32::new(val).map(Self)
31            }
32
33            /// Underlying raw `u32` value (always >= 1).
34            pub fn to_raw(self) -> u32 {
35                self.0.get()
36            }
37        }
38
39        impl From<NonZeroU32> for $name {
40            fn from(val: NonZeroU32) -> Self {
41                Self(val)
42            }
43        }
44    };
45}
46
47define_id_type!(FileId);
48define_id_type!(SymbolId);
49define_id_type!(SnapshotId);
50define_id_type!(DataNodeId);
51
52/// Thread-safe ID allocator.
53///
54/// The counter starts at 1 so every allocated ID is a valid `NonZeroU32`.
55/// Zero is never handed out: it is the niche value that makes `Option<Id>`
56/// free (4 bytes, not 8). Allocation panics only after exhausting the full
57/// `u32` space (>4 billion ids), which is treated as a programmer limit.
58///
59/// The counter uses `Relaxed` ordering: uniqueness needs the atomicity of the
60/// read-modify-write, not an ordering edge to any other memory location.
61#[derive(Debug)]
62pub struct IdGenerator<T> {
63    counter: AtomicU32,
64    _marker: std::marker::PhantomData<T>,
65}
66
67impl<T> IdGenerator<T> {
68    pub fn new() -> Self {
69        Self {
70            counter: AtomicU32::new(1),
71            _marker: std::marker::PhantomData,
72        }
73    }
74
75    pub fn with_start(start: u32) -> Self {
76        Self {
77            counter: AtomicU32::new(start.max(1)),
78            _marker: std::marker::PhantomData,
79        }
80    }
81
82    pub fn next(&self) -> T
83    where
84        T: From<NonZeroU32>,
85    {
86        let val = self.counter.fetch_add(1, Ordering::Relaxed);
87        let nz = NonZeroU32::new(val)
88            .expect("IdGenerator exhausted u32 space (allocated > u32::MAX ids)");
89        T::from(nz)
90    }
91
92    /// Reserve `count` contiguous identifiers and return the first slot.
93    ///
94    /// A caller that numbers a whole unit of work reserves one block, so the
95    /// block order follows the caller's order and never the thread timing.
96    /// An empty reservation returns the next slot without consuming it.
97    pub fn reserve(&self, count: u32) -> u32 {
98        if count == 0 {
99            return self.counter.load(Ordering::Relaxed);
100        }
101        let start = self.counter.fetch_add(count, Ordering::Relaxed);
102        assert!(
103            start != 0 && start.checked_add(count - 1).is_some(),
104            "IdGenerator exhausted u32 space (reserved > u32::MAX ids)"
105        );
106        start
107    }
108}
109
110impl<T> Default for IdGenerator<T> {
111    fn default() -> Self {
112        Self::new()
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119    use std::collections::HashSet;
120    use std::sync::Arc;
121
122    #[test]
123    fn file_id_sequential() {
124        let idgen = IdGenerator::<FileId>::new();
125        assert_eq!(idgen.next(), FileId::new(1).unwrap());
126        assert_eq!(idgen.next(), FileId::new(2).unwrap());
127        assert_eq!(idgen.next(), FileId::new(3).unwrap());
128    }
129
130    #[test]
131    fn symbol_id_sequential() {
132        let idgen = IdGenerator::<SymbolId>::new();
133        assert_eq!(idgen.next(), SymbolId::new(1).unwrap());
134        assert_eq!(idgen.next(), SymbolId::new(2).unwrap());
135        assert_eq!(idgen.next(), SymbolId::new(3).unwrap());
136    }
137
138    #[test]
139    fn snapshot_id_sequential() {
140        let idgen = IdGenerator::<SnapshotId>::new();
141        assert_eq!(idgen.next(), SnapshotId::new(1).unwrap());
142        assert_eq!(idgen.next(), SnapshotId::new(2).unwrap());
143        assert_eq!(idgen.next(), SnapshotId::new(3).unwrap());
144    }
145
146    #[test]
147    fn id_generator_thread_safe() {
148        let idgen = Arc::new(IdGenerator::<SymbolId>::new());
149        let mut handles = Vec::new();
150
151        for _ in 0..4 {
152            let idgen = Arc::clone(&idgen);
153            handles.push(std::thread::spawn(move || {
154                let mut ids = Vec::with_capacity(250);
155                for _ in 0..250 {
156                    ids.push(idgen.next());
157                }
158                ids
159            }));
160        }
161
162        let all_ids: Vec<SymbolId> = handles
163            .into_iter()
164            .flat_map(|h| h.join().unwrap_or_default())
165            .collect();
166
167        let unique: HashSet<SymbolId> = all_ids.iter().copied().collect();
168        assert_eq!(unique.len(), 1000, "expected 1000 unique IDs");
169        assert!(
170            unique
171                .iter()
172                .all(|id| id.to_raw() >= 1 && id.to_raw() <= 1000),
173            "all IDs must be in range 1..=1000"
174        );
175    }
176
177    #[test]
178    fn file_id_serde_roundtrip() {
179        let original = FileId::new(42).unwrap();
180        let json = serde_json::to_string(&original).unwrap();
181        let roundtrip: FileId = serde_json::from_str(&json).unwrap();
182        assert_eq!(original, roundtrip);
183    }
184
185    #[test]
186    fn symbol_id_serde_roundtrip() {
187        let original = SymbolId::new(99).unwrap();
188        let json = serde_json::to_string(&original).unwrap();
189        let roundtrip: SymbolId = serde_json::from_str(&json).unwrap();
190        assert_eq!(original, roundtrip);
191    }
192
193    #[test]
194    fn file_id_to_raw() {
195        assert_eq!(FileId::new(7).unwrap().to_raw(), 7);
196    }
197
198    #[test]
199    fn id_generator_default() {
200        let idgen = IdGenerator::<FileId>::default();
201        assert_eq!(idgen.next(), FileId::new(1).unwrap());
202    }
203
204    #[test]
205    fn id_generator_with_start_zero_safeguard() {
206        let idgen = IdGenerator::<SymbolId>::with_start(0);
207        let first = idgen.next();
208        assert_eq!(
209            first.to_raw(),
210            1,
211            "with_start(0) must sanitize to start ID 1"
212        );
213    }
214
215    #[test]
216    fn zero_is_invalid() {
217        assert!(FileId::new(0).is_none());
218        assert!(SymbolId::new(0).is_none());
219        assert!(SnapshotId::new(0).is_none());
220        assert!(DataNodeId::new(0).is_none());
221    }
222
223    #[test]
224    fn zero_rejected_on_deserialize() {
225        let res: Result<FileId, _> = serde_json::from_str("0");
226        assert!(res.is_err(), "deserializing 0 must fail");
227    }
228
229    #[test]
230    fn option_id_is_niche_optimized() {
231        // NonZeroU32 niche: Option<Id> collapses to the same size as Id.
232        assert_eq!(
233            std::mem::size_of::<Option<FileId>>(),
234            std::mem::size_of::<FileId>()
235        );
236        assert_eq!(std::mem::size_of::<Option<FileId>>(), 4);
237        assert_eq!(std::mem::size_of::<Option<SymbolId>>(), 4);
238        assert_eq!(std::mem::size_of::<Option<SnapshotId>>(), 4);
239        assert_eq!(std::mem::size_of::<Option<DataNodeId>>(), 4);
240    }
241    #[test]
242    fn counters_hand_out_unique_ids_under_contention() {
243        let idgen = Arc::new(IdGenerator::<SymbolId>::new());
244        let handles: Vec<_> = (0..8)
245            .map(|_| {
246                let idgen = Arc::clone(&idgen);
247                std::thread::spawn(move || {
248                    let mut ids = Vec::with_capacity(10_000);
249                    for _ in 0..10_000 {
250                        ids.push(idgen.next());
251                    }
252                    ids
253                })
254            })
255            .collect();
256
257        let all: Vec<SymbolId> = handles
258            .into_iter()
259            .flat_map(|handle| handle.join().unwrap())
260            .collect();
261        let unique: HashSet<SymbolId> = all.iter().copied().collect();
262        assert_eq!(unique.len(), 80_000, "every identifier is handed out once");
263        let raw: Vec<u32> = unique.iter().map(|id| id.to_raw()).collect();
264        assert_eq!(raw.iter().min(), Some(&1));
265        assert_eq!(raw.iter().max(), Some(&80_000));
266    }
267}