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