Skip to main content

obj_alloc/
id_map.rs

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