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#[derive(Debug)]
59pub struct IdGenerator<T> {
60    counter: AtomicU32,
61    _marker: std::marker::PhantomData<T>,
62}
63
64impl<T> IdGenerator<T> {
65    pub fn new() -> Self {
66        Self {
67            counter: AtomicU32::new(1),
68            _marker: std::marker::PhantomData,
69        }
70    }
71
72    pub fn with_start(start: u32) -> Self {
73        Self {
74            counter: AtomicU32::new(start.max(1)),
75            _marker: std::marker::PhantomData,
76        }
77    }
78
79    pub fn next(&self) -> T
80    where
81        T: From<NonZeroU32>,
82    {
83        let val = self.counter.fetch_add(1, Ordering::SeqCst);
84        let nz = NonZeroU32::new(val)
85            .expect("IdGenerator exhausted u32 space (allocated > u32::MAX ids)");
86        T::from(nz)
87    }
88}
89
90impl<T> Default for IdGenerator<T> {
91    fn default() -> Self {
92        Self::new()
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    use std::collections::HashSet;
100    use std::sync::Arc;
101
102    #[test]
103    fn file_id_sequential() {
104        let idgen = IdGenerator::<FileId>::new();
105        assert_eq!(idgen.next(), FileId::new(1).unwrap());
106        assert_eq!(idgen.next(), FileId::new(2).unwrap());
107        assert_eq!(idgen.next(), FileId::new(3).unwrap());
108    }
109
110    #[test]
111    fn symbol_id_sequential() {
112        let idgen = IdGenerator::<SymbolId>::new();
113        assert_eq!(idgen.next(), SymbolId::new(1).unwrap());
114        assert_eq!(idgen.next(), SymbolId::new(2).unwrap());
115        assert_eq!(idgen.next(), SymbolId::new(3).unwrap());
116    }
117
118    #[test]
119    fn snapshot_id_sequential() {
120        let idgen = IdGenerator::<SnapshotId>::new();
121        assert_eq!(idgen.next(), SnapshotId::new(1).unwrap());
122        assert_eq!(idgen.next(), SnapshotId::new(2).unwrap());
123        assert_eq!(idgen.next(), SnapshotId::new(3).unwrap());
124    }
125
126    #[test]
127    fn id_generator_thread_safe() {
128        let idgen = Arc::new(IdGenerator::<SymbolId>::new());
129        let mut handles = Vec::new();
130
131        for _ in 0..4 {
132            let idgen = Arc::clone(&idgen);
133            handles.push(std::thread::spawn(move || {
134                let mut ids = Vec::with_capacity(250);
135                for _ in 0..250 {
136                    ids.push(idgen.next());
137                }
138                ids
139            }));
140        }
141
142        let all_ids: Vec<SymbolId> = handles
143            .into_iter()
144            .flat_map(|h| h.join().unwrap_or_default())
145            .collect();
146
147        let unique: HashSet<SymbolId> = all_ids.iter().copied().collect();
148        assert_eq!(unique.len(), 1000, "expected 1000 unique IDs");
149        assert!(
150            unique
151                .iter()
152                .all(|id| id.to_raw() >= 1 && id.to_raw() <= 1000),
153            "all IDs must be in range 1..=1000"
154        );
155    }
156
157    #[test]
158    fn file_id_serde_roundtrip() {
159        let original = FileId::new(42).unwrap();
160        let json = serde_json::to_string(&original).unwrap();
161        let roundtrip: FileId = serde_json::from_str(&json).unwrap();
162        assert_eq!(original, roundtrip);
163    }
164
165    #[test]
166    fn symbol_id_serde_roundtrip() {
167        let original = SymbolId::new(99).unwrap();
168        let json = serde_json::to_string(&original).unwrap();
169        let roundtrip: SymbolId = serde_json::from_str(&json).unwrap();
170        assert_eq!(original, roundtrip);
171    }
172
173    #[test]
174    fn file_id_to_raw() {
175        assert_eq!(FileId::new(7).unwrap().to_raw(), 7);
176    }
177
178    #[test]
179    fn id_generator_default() {
180        let idgen = IdGenerator::<FileId>::default();
181        assert_eq!(idgen.next(), FileId::new(1).unwrap());
182    }
183
184    #[test]
185    fn id_generator_with_start_zero_safeguard() {
186        let idgen = IdGenerator::<SymbolId>::with_start(0);
187        let first = idgen.next();
188        assert_eq!(
189            first.to_raw(),
190            1,
191            "with_start(0) must sanitize to start ID 1"
192        );
193    }
194
195    #[test]
196    fn zero_is_invalid() {
197        assert!(FileId::new(0).is_none());
198        assert!(SymbolId::new(0).is_none());
199        assert!(SnapshotId::new(0).is_none());
200        assert!(DataNodeId::new(0).is_none());
201    }
202
203    #[test]
204    fn zero_rejected_on_deserialize() {
205        let res: Result<FileId, _> = serde_json::from_str("0");
206        assert!(res.is_err(), "deserializing 0 must fail");
207    }
208
209    #[test]
210    fn option_id_is_niche_optimized() {
211        // NonZeroU32 niche: Option<Id> collapses to the same size as Id.
212        assert_eq!(
213            std::mem::size_of::<Option<FileId>>(),
214            std::mem::size_of::<FileId>()
215        );
216        assert_eq!(std::mem::size_of::<Option<FileId>>(), 4);
217        assert_eq!(std::mem::size_of::<Option<SymbolId>>(), 4);
218        assert_eq!(std::mem::size_of::<Option<SnapshotId>>(), 4);
219        assert_eq!(std::mem::size_of::<Option<DataNodeId>>(), 4);
220    }
221}