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(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
21 #[serde(transparent)]
22 pub struct $name(NonZeroU32);
23
24 impl $name {
25 pub fn new(val: u32) -> Option<Self> {
28 NonZeroU32::new(val).map(Self)
29 }
30
31 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#[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 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}