1use 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 pub fn new(val: u32) -> Option<Self> {
30 NonZeroU32::new(val).map(Self)
31 }
32
33 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#[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 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}