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

obj_pool/
par.rs
