obj_pool/
par.rs

1use crate::*;
2use parking_lot::{
3    MappedRwLockReadGuard, MappedRwLockWriteGuard, RwLock, RwLockReadGuard, RwLockWriteGuard,
4};
5use std::cell::Cell;
6use std::convert::TryInto;
7
8thread_local! {
9  static COUNTER: Cell<usize> = const { Cell::new(0) };
10}
11
12pub struct ParObjPool<T, const S: usize> {
13    shards: [RwLock<ObjPool<T>>; S],
14}
15
16impl<T, const S: usize> Default for ParObjPool<T, S> {
17    fn default() -> Self {
18        Self::new()
19    }
20}
21
22impl<T, const S: usize> ParObjPool<T, S> {
23    pub fn new() -> Self {
24        let mut shards = Vec::with_capacity(S);
25        shards.resize_with(S, || RwLock::new(ObjPool::<T>::new()));
26        Self {
27            shards: shards.try_into().expect("invalid array"),
28        }
29    }
30
31    pub fn insert(&self, object: T) -> ObjId {
32        let counter = COUNTER.with(|c| {
33            let v = c.get();
34            c.set(v.wrapping_add(1));
35            v
36        });
37        let shard_index = counter % S;
38        self.obj_id_to_external(self.shards[shard_index].write().insert(object), shard_index)
39    }
40
41    pub fn remove(&self, obj_id: ObjId) -> Option<T> {
42        let (shard_index, obj_id) = self.obj_id_from_external(obj_id);
43        self.shards[shard_index].write().remove(obj_id)
44    }
45
46    pub fn get(&self, obj_id: ObjId) -> Option<MappedRwLockReadGuard<'_, T>> {
47        let (shard_index, obj_id) = self.obj_id_from_external(obj_id);
48        RwLockReadGuard::try_map(self.shards[shard_index].read(), |obj_pool| {
49            obj_pool.get(obj_id)
50        })
51        .ok()
52    }
53
54    pub fn try_get(&self, obj_id: ObjId) -> Option<MappedRwLockReadGuard<'_, T>> {
55        let (shard_index, obj_id) = self.obj_id_from_external(obj_id);
56        RwLockReadGuard::try_map(self.shards[shard_index].try_read()?, |obj_pool| {
57            obj_pool.get(obj_id)
58        })
59        .ok()
60    }
61
62    pub fn get_mut(&self, obj_id: ObjId) -> Option<MappedRwLockWriteGuard<'_, T>> {
63        let (shard_index, obj_id) = self.obj_id_from_external(obj_id);
64        RwLockWriteGuard::try_map(self.shards[shard_index].write(), |obj_pool| {
65            obj_pool.get_mut(obj_id)
66        })
67        .ok()
68    }
69
70    pub fn try_get_mut(&self, obj_id: ObjId) -> Option<MappedRwLockWriteGuard<'_, T>> {
71        let (shard_index, obj_id) = self.obj_id_from_external(obj_id);
72        RwLockWriteGuard::try_map(self.shards[shard_index].try_write()?, |obj_pool| {
73            obj_pool.get_mut(obj_id)
74        })
75        .ok()
76    }
77
78    pub fn clear(&self) {
79        for shard in &self.shards {
80            let mut shard = shard.write();
81            shard.clear();
82            shard.shrink_to_fit();
83        }
84    }
85
86    pub fn shrink_to_fit(&self) {
87        for shard in &self.shards {
88            let mut shard = shard.write();
89            shard.shrink_to_fit();
90        }
91    }
92
93    pub fn capacity(&self) -> usize {
94        self.shards.iter().map(|s| s.read().capacity()).sum()
95    }
96
97    fn obj_id_to_external(&self, obj_id: ObjId, shard_index: usize) -> ObjId {
98        ObjId(NonZeroU32::new((shard_index << 26) as u32 | obj_id.get()).expect("invalid value"))
99    }
100
101    fn obj_id_from_external(&self, obj_id: ObjId) -> (usize, ObjId) {
102        let v = obj_id.get();
103        (
104            (v >> 26) as usize,
105            ObjId(NonZeroU32::new(v & 0x03FFFFFF).expect("invalid value")),
106        )
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    #[test]
115    fn test() {
116        let o = ParObjPool::<usize, 16>::new();
117        let k = o.insert(10);
118        assert_eq!(o.get(k).map(|v| *v), Some(10));
119        let k = o.insert(20);
120        assert_eq!(o.get(k).map(|v| *v), Some(20));
121        let k = o.insert(30);
122        assert_eq!(o.get(k).map(|v| *v), Some(30));
123        assert_eq!(o.try_get(k).map(|v| *v), Some(30));
124        assert_eq!(o.try_get_mut(k).map(|v| *v), Some(30));
125    }
126}
obj_pool/par.rs

obj_pool/
par.rs
