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(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
21        #[serde(transparent)]
22        pub struct $name(NonZeroU32);
23
24        impl $name {
25            /// Construct from a raw `u32`. Returns `None` for 0, which is the
26            /// permanently invalid niche value reserved for `Option<$name>`.
27            pub fn new(val: u32) -> Option<Self> {
28                NonZeroU32::new(val).map(Self)
29            }
30
31            /// Underlying raw `u32` value (always >= 1).
32            pub fn to_raw(self) -> u32 {
33                self.0.get()
34            }
35        }
36
37        impl From<NonZeroU32> for $name {
38            fn from(val: NonZeroU32) -> Self {
39                Self(val)
40            }
41        }
42    };
43}
44
45define_id_type!(FileId);
46define_id_type!(SymbolId);
47define_id_type!(SnapshotId);
48define_id_type!(DataNodeId);
49
50/// Thread-safe ID allocator.
51///
52/// The counter starts at 1 so every allocated ID is a valid `NonZeroU32`.
53/// Zero is never handed out: it is the niche value that makes `Option<Id>`
54/// free (4 bytes, not 8). Allocation panics only after exhausting the full
55/// `u32` space (>4 billion ids), which is treated as a programmer limit.
56#[derive(Debug)]
57pub struct IdGenerator<T> {
58    counter: AtomicU32,
59    _marker: std::marker::PhantomData<T>,
60}
61
62impl<T> IdGenerator<T> {
63    pub fn new() -> Self {
64        Self {
65            counter: AtomicU32::new(1),
66            _marker: std::marker::PhantomData,
67        }
68    }
69
70    pub fn with_start(start: u32) -> Self {
71        Self {
72            counter: AtomicU32::new(start.max(1)),
73            _marker: std::marker::PhantomData,
74        }
75    }
76
77    pub fn next(&self) -> T
78    where
79        T: From<NonZeroU32>,
80    {
81        let val = self.counter.fetch_add(1, Ordering::SeqCst);
82        let nz = NonZeroU32::new(val)
83            .expect("IdGenerator exhausted u32 space (allocated > u32::MAX ids)");
84        T::from(nz)
85    }
86}
87
88impl<T> Default for IdGenerator<T> {
89    fn default() -> Self {
90        Self::new()
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97    use std::collections::HashSet;
98    use std::sync::Arc;
99
100    #[test]
101    fn file_id_sequential() {
102        let idgen = IdGenerator::<FileId>::new();
103        assert_eq!(idgen.next(), FileId::new(1).unwrap());
104        assert_eq!(idgen.next(), FileId::new(2).unwrap());
105        assert_eq!(idgen.next(), FileId::new(3).unwrap());
106    }
107
108    #[test]
109    fn symbol_id_sequential() {
110        let idgen = IdGenerator::<SymbolId>::new();
111        assert_eq!(idgen.next(), SymbolId::new(1).unwrap());
112        assert_eq!(idgen.next(), SymbolId::new(2).unwrap());
113        assert_eq!(idgen.next(), SymbolId::new(3).unwrap());
114    }
115
116    #[test]
117    fn snapshot_id_sequential() {
118        let idgen = IdGenerator::<SnapshotId>::new();
119        assert_eq!(idgen.next(), SnapshotId::new(1).unwrap());
120        assert_eq!(idgen.next(), SnapshotId::new(2).unwrap());
121        assert_eq!(idgen.next(), SnapshotId::new(3).unwrap());
122    }
123
124    #[test]
125    fn id_generator_thread_safe() {
126        let idgen = Arc::new(IdGenerator::<SymbolId>::new());
127        let mut handles = Vec::new();
128
129        for _ in 0..4 {
130            let idgen = Arc::clone(&idgen);
131            handles.push(std::thread::spawn(move || {
132                let mut ids = Vec::with_capacity(250);
133                for _ in 0..250 {
134                    ids.push(idgen.next());
135                }
136                ids
137            }));
138        }
139
140        let all_ids: Vec<SymbolId> = handles
141            .into_iter()
142            .flat_map(|h| h.join().unwrap_or_default())
143            .collect();
144
145        let unique: HashSet<SymbolId> = all_ids.iter().copied().collect();
146        assert_eq!(unique.len(), 1000, "expected 1000 unique IDs");
147        assert!(
148            unique
149                .iter()
150                .all(|id| id.to_raw() >= 1 && id.to_raw() <= 1000),
151            "all IDs must be in range 1..=1000"
152        );
153    }
154
155    #[test]
156    fn file_id_serde_roundtrip() {
157        let original = FileId::new(42).unwrap();
158        let json = serde_json::to_string(&original).unwrap();
159        let roundtrip: FileId = serde_json::from_str(&json).unwrap();
160        assert_eq!(original, roundtrip);
161    }
162
163    #[test]
164    fn symbol_id_serde_roundtrip() {
165        let original = SymbolId::new(99).unwrap();
166        let json = serde_json::to_string(&original).unwrap();
167        let roundtrip: SymbolId = serde_json::from_str(&json).unwrap();
168        assert_eq!(original, roundtrip);
169    }
170
171    #[test]
172    fn file_id_to_raw() {
173        assert_eq!(FileId::new(7).unwrap().to_raw(), 7);
174    }
175
176    #[test]
177    fn id_generator_default() {
178        let idgen = IdGenerator::<FileId>::default();
179        assert_eq!(idgen.next(), FileId::new(1).unwrap());
180    }
181
182    #[test]
183    fn id_generator_with_start_zero_safeguard() {
184        let idgen = IdGenerator::<SymbolId>::with_start(0);
185        let first = idgen.next();
186        assert_eq!(
187            first.to_raw(),
188            1,
189            "with_start(0) must sanitize to start ID 1"
190        );
191    }
192
193    #[test]
194    fn zero_is_invalid() {
195        assert!(FileId::new(0).is_none());
196        assert!(SymbolId::new(0).is_none());
197        assert!(SnapshotId::new(0).is_none());
198        assert!(DataNodeId::new(0).is_none());
199    }
200
201    #[test]
202    fn zero_rejected_on_deserialize() {
203        let res: Result<FileId, _> = serde_json::from_str("0");
204        assert!(res.is_err(), "deserializing 0 must fail");
205    }
206
207    #[test]
208    fn option_id_is_niche_optimized() {
209        // NonZeroU32 niche: Option<Id> collapses to the same size as Id.
210        assert_eq!(
211            std::mem::size_of::<Option<FileId>>(),
212            std::mem::size_of::<FileId>()
213        );
214        assert_eq!(std::mem::size_of::<Option<FileId>>(), 4);
215        assert_eq!(std::mem::size_of::<Option<SymbolId>>(), 4);
216        assert_eq!(std::mem::size_of::<Option<SnapshotId>>(), 4);
217        assert_eq!(std::mem::size_of::<Option<DataNodeId>>(), 4);
218    }
219}