obj_alloc/
id_map.rs

1
2use core::fmt;
3use serde::{Deserialize, Deserializer, Serialize, Serializer};
4use std::collections::HashMap;
5use std::marker::PhantomData;
6use std::ops::{Index, IndexMut};
7
8/// Id 基础 trait，所有自定义 Id 需实现此 trait
9pub trait Id: Copy + Clone + Eq + PartialEq + fmt::Debug + Into<u64> + From<u64> {
10    const NONE: Self;
11
12    fn none() -> Self {
13        Self::NONE
14    }
15    
16    fn is_none(&self) -> bool {
17        self.eq(&Self::NONE)
18    }
19    
20    /// 快速转换为 u64
21    fn as_u64(&self) -> u64 {
22        (*self).into()
23    }
24    
25    /// 从 u64 构建 Id
26    fn from_u64(val: u64) -> Self {
27        Self::from(val)
28    }
29}
30
31
32/// 生成自定义 Id 类型的极简宏
33#[macro_export]
34macro_rules! new_id_type {
35    () => {};
36
37    (
38        $(#[$meta:meta])*
39        $vis:vis struct $name:ident;
40        $($rest:tt)*
41    ) => {
42        $(#[$meta])*
43        #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
44        $vis struct $name(pub u64);
45
46        impl From<u64> for $name {
47            #[inline]
48            fn from(val: u64) -> Self {
49                Self(val)
50            }
51        }
52
53        impl From<$name> for u64 {
54            #[inline]
55            fn from(id: $name) -> Self {
56                id.0
57            }
58        }
59
60        impl $crate::Id for $name {
61            const NONE: Self = Self(0);
62        }
63
64        impl serde::Serialize for $name {
65            #[inline]
66            fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
67            where
68                S: serde::Serializer,
69            {
70                self.0.serialize(serializer)
71            }
72        }
73
74        impl<'de> serde::Deserialize<'de> for $name {
75            #[inline]
76            fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
77            where
78                D: serde::Deserializer<'de>,
79            {
80                let val = u64::deserialize(deserializer)?;
81                Ok(Self(val))
82            }
83        }
84
85        $crate::new_id_type!($($rest)*);
86    };
87}
88
89new_id_type!{
90    pub struct DefaultId;
91}
92
93
94#[derive(Debug, Clone)]
95pub struct IdMap<K: Id, V> {
96    pub(crate) inner: HashMap<u64, V>,
97    max_id: u64,
98    _marker: PhantomData<K>,
99}
100
101impl<V> IdMap<DefaultId, V> {
102    /// 创建空的 IdMap（初始 max_id = 0）
103    pub fn new() -> Self { Self::with_id_capacity(0) }
104    
105    /// 创建指定初始容量的 IdMap
106    pub fn with_capacity(capacity: usize) -> Self { Self::with_id_capacity(capacity) }
107}
108
109impl<K: Id ,V> Default for IdMap<K, V> {
110    fn default() -> Self {
111        Self::with_id()
112    }
113}
114
115impl<K: Id, V> IdMap<K, V> {
116    /// 为自定义 Id 类型创建空 IdMap
117    pub fn with_id() -> Self {
118        Self {
119            inner: HashMap::new(),
120            max_id: 0,
121            _marker: PhantomData,
122        }
123    }
124    
125    /// 自定义 Id 类型创建指定初始容量的 IdMap
126    pub fn with_id_capacity(capacity: usize) -> Self {
127        Self {
128            inner: HashMap::with_capacity(capacity),
129            max_id: 0,
130            _marker: PhantomData,
131        }
132    }
133    
134    /// 插入值，自动生成递增 Id 并返回
135    pub fn insert(&mut self, value: V) -> K {
136        self.max_id += 1; // 递增生成新 Id（从 1 开始，避免 0 作为初始值）
137        let id_u64 = self.max_id;
138        self.inner.insert(id_u64, value); // 存储值
139        K::from_u64(id_u64) // 转换为指定 Id 类型并返回
140    }
141    
142    
143    /// 【手动指定 Id】插入键值对，返回旧值（若存在）
144    ///
145    /// 注意：若手动传入的 Id 大于当前 max_id，会更新 max_id 以保证自动生成的 Id 不重复
146    pub fn insert_with_id(&mut self, id: K, value: V) -> Option<V> {
147        let id_u64 = id.as_u64();
148        // 若手动传入的 Id 更大，更新 max_id，避免自动生成 Id 重复
149        if id_u64 > self.max_id {
150            self.max_id = id_u64;
151        }
152        self.inner.insert(id_u64, value)
153    }
154    
155    /// 从 Vec<V> 批量插入值，自动生成递增 Id，返回对应的 Id 列表
156    /// 生成的 Id 从当前 max_id + 1 开始连续递增
157    pub fn from_vec(values: Vec<V>) -> (Self, Vec<K>) {
158        let mut map = Self {
159            inner: HashMap::with_capacity(values.len()),
160            max_id: 0,
161            _marker: PhantomData,
162        };
163        let ids = values
164            .into_iter()
165            .map(|val| {
166                map.max_id += 1;
167                let id_u64 = map.max_id;
168                map.inner.insert(id_u64, val);
169                K::from_u64(id_u64)
170            })
171            .collect();
172        (map, ids)
173    }
174    
175    /// 循环插入：先生成递增 Id，再通过闭包（Id → V）生成值并插入
176    /// 适用于值需要依赖自身 Id 的场景（如循环引用/关联 Id 的场景）
177    pub fn insert_cyclic<F>(&mut self, f: F) -> K
178    where
179        F: FnOnce(K) -> V,
180    {
181        self.max_id += 1;
182        let new_id = K::from_u64(self.max_id);
183        let value = f(new_id);
184        self.inner.insert(self.max_id, value);
185        new_id
186    }
187    
188    /// 根据 Id 查询值
189    pub fn get(&self, id: K) -> Option<&V> {
190        self.inner.get(&id.as_u64())
191    }
192    
193    /// 根据 Id 查询可变值
194    pub fn get_mut(&mut self, id: K) -> Option<&mut V> {
195        self.inner.get_mut(&id.as_u64())
196    }
197    
198    /// 根据 Id 删除值
199    pub fn remove(&mut self, id: K) -> Option<V> {
200        self.inner.remove(&id.as_u64())
201    }
202    
203    /// 判断是否包含指定 Id
204    pub fn contains_id(&self, id: K) -> bool {
205        self.inner.contains_key(&id.as_u64())
206    }
207    
208    /// 获取当前最大 Id（仅用于参考，删除 Id 后不会回退）
209    pub fn max_id(&self) -> K {
210        K::from_u64(self.max_id)
211    }
212    
213    /// 获取元素数量
214    pub fn len(&self) -> usize {
215        self.inner.len()
216    }
217    
218    /// 判断是否为空
219    pub fn is_empty(&self) -> bool {
220        self.inner.is_empty()
221    }
222    
223    /// 清空所有元素（保留 max_id 不变，避免 Id 重复）
224    pub fn clear(&mut self) {
225        self.inner.clear();
226    }
227}
228
229impl<K: Id, V: Serialize> Serialize for IdMap<K, V> {
230    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
231    where
232        S: Serializer,
233    {
234        // 直接序列化内部的 HashMap，就像 IdMap 不存在一样
235        self.inner.serialize(serializer)
236    }
237}
238
239impl<'de, K: Id, V: Deserialize<'de>> Deserialize<'de> for IdMap<K, V> {
240    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
241    where
242        D: Deserializer<'de>,
243    {
244        let inner: HashMap<u64, V> = HashMap::deserialize(deserializer)?;
245        
246        let max_id = inner.keys().copied().max().unwrap_or(0);
247        
248        Ok(IdMap {
249            inner,
250            max_id,
251            _marker: PhantomData,
252        })
253    }
254}
255
256use std::collections::hash_map;
257use std::iter::{Map};
258
259impl<K: Id, V> IntoIterator for IdMap<K, V> {
260    type Item = (K, V);
261    type IntoIter = Map<hash_map::IntoIter<u64, V>, fn((u64, V)) -> (K, V)>;
262    
263    fn into_iter(self) -> Self::IntoIter {
264        self.inner.into_iter().map(|(id_u64, v)| (K::from_u64(id_u64), v))
265    }
266}
267
268impl<'a, K: Id, V> IntoIterator for &'a IdMap<K, V> {
269    type Item = (K, &'a V);
270    type IntoIter = Map<hash_map::Iter<'a, u64, V>, fn((&'a u64, &'a V)) -> (K, &'a V)>;
271    
272    fn into_iter(self) -> Self::IntoIter {
273        self.inner.iter().map(|(id_u64, v)| (K::from_u64(*id_u64), v))
274    }
275}
276
277impl<'a, K: Id, V> IntoIterator for &'a mut IdMap<K, V> {
278    type Item = (K, &'a mut V);
279    type IntoIter = Map<hash_map::IterMut<'a, u64, V>, fn((&'a u64, &'a mut V)) -> (K, &'a mut V)>;
280    
281    fn into_iter(self) -> Self::IntoIter {
282        self.inner.iter_mut().map(|(id_u64, v)| (K::from_u64(*id_u64), v))
283    }
284}
285
286impl<K: Id, V> IdMap<K, V> {
287    pub fn iter<'a>(&'a self) -> Map<hash_map::Iter<'a, u64, V>, fn((&'a u64, &'a V)) -> (K, &'a V)> {
288        self.inner.iter().map(|(id_u64, v)| (K::from_u64(*id_u64), v))
289    }
290    
291    pub fn iter_mut<'a>(&'a mut self) -> Map<hash_map::IterMut<'a, u64, V>, fn((&'a u64, &'a mut V)) -> (K, &'a mut V)> {
292        self.inner.iter_mut().map(|(id_u64, v)| (K::from_u64(*id_u64), v))
293    }
294    
295    pub fn keys(&self) -> Map<hash_map::Keys<'_, u64, V>, fn(&u64) -> K> {
296        self.inner.keys().map(|id_u64| K::from_u64(*id_u64))
297    }
298    
299    pub fn values(&self) -> hash_map::Values<'_, u64, V> {
300        self.inner.values()
301    }
302    
303    pub fn values_mut(&mut self) -> hash_map::ValuesMut<'_, u64, V> {
304        self.inner.values_mut()
305    }
306}
307
308impl<K: Id, V> Index<K> for IdMap<K, V> {
309    type Output = V;
310    
311    fn index(&self, id: K) -> &Self::Output {
312        self.get(id).expect("invalid IdMap id")
313    }
314}
315
316impl<K: Id, V> IndexMut<K> for IdMap<K, V> {
317    fn index_mut(&mut self, id: K) -> &mut Self::Output {
318        self.get_mut(id).expect("invalid IdMap id")
319    }
320}
321
322// ============================ 测试用例 ============================
323#[cfg(test)]
324mod tests {
325    use super::*;
326    use serde_json;
327    
328    // 测试默认 Id + 自动递增生成
329    #[test]
330    fn test_default_id_auto_generate() {
331        let mut map = IdMap::new();
332        
333        // 插入值，自动返回递增 Id
334        let id1 = map.insert("hello");
335        let id2 = map.insert("world");
336        let id3 = map.insert("rust");
337        
338        // 验证 Id 递增（从 1 开始）
339        assert_eq!(id1, DefaultId(1));
340        assert_eq!(id2, DefaultId(2));
341        assert_eq!(id3, DefaultId(3));
342        
343        // 验证值查询
344        assert_eq!(map.get(id1), Some(&"hello"));
345        assert_eq!(map[id2], "world");
346        assert_eq!(map.max_id(), DefaultId(3));
347        
348        // 删除值后，max_id 不回退
349        map.remove(id2);
350        assert_eq!(map.max_id(), DefaultId(3));
351        let id4 = map.insert("new value");
352        assert_eq!(id4, DefaultId(4)); // 继续递增
353        
354        // 数量/空判断
355        assert_eq!(map.len(), 3);
356        map.clear();
357        assert!(map.is_empty());
358    }
359    
360    // 测试自定义 Id
361    new_id_type! {
362        struct MyId;
363    }
364    
365    #[test]
366    fn test_custom_id() {
367        let mut map = IdMap::<MyId, u32>::with_id();
368        
369        let id1 = map.insert(42);
370        let id2 = map.insert(100);
371        
372        assert_eq!(id1, MyId(1));
373        assert_eq!(id2, MyId(2));
374        assert_eq!(map.get(id1), Some(&42));
375        
376        // 删除测试
377        map.remove(id1);
378        assert!(!map.contains_id(id1));
379    }
380    
381    // 测试 Id 透明序列化
382    #[test]
383    fn test_id_serde() {
384        // 测试默认 Id
385        let id = DefaultId(123456789);
386        let json = serde_json::to_string(&id).unwrap();
387        assert_eq!(json, "123456789"); // 直接输出 u64 字符串
388        let id2: DefaultId = serde_json::from_str(&json).unwrap();
389        assert_eq!(id2, id);
390        
391        // 测试自定义 Id
392        let my_id = MyId(987654321);
393        let json = serde_json::to_string(&my_id).unwrap();
394        let my_id2: MyId = serde_json::from_str(&json).unwrap();
395        assert_eq!(my_id2, my_id);
396    }
397}
obj_alloc/id_map.rs

obj_alloc/
id_map.rs
