rkyv/de/pooling/
alloc.rs

1use core::{error::Error, fmt, hash::BuildHasherDefault};
2
3use hashbrown::hash_map::{Entry, HashMap};
4use rancor::{fail, Source};
5
6use crate::{
7    de::pooling::{ErasedPtr, Pooling, PoolingState},
8    hash::FxHasher64,
9};
10
11#[derive(Debug)]
12struct SharedPointer {
13    ptr: ErasedPtr,
14    drop: unsafe fn(ErasedPtr),
15}
16
17impl Drop for SharedPointer {
18    fn drop(&mut self) {
19        unsafe {
20            (self.drop)(self.ptr);
21        }
22    }
23}
24
25/// A shared pointer strategy that pools together deserializations of the same
26/// shared pointer.
27#[derive(Default)]
28pub struct Pool {
29    shared_pointers:
30        HashMap<usize, Option<SharedPointer>, BuildHasherDefault<FxHasher64>>,
31}
32
33impl Pool {
34    /// Creates a new shared pointer unifier.
35    #[inline]
36    pub fn new() -> Self {
37        Self::default()
38    }
39
40    /// Creates a new shared pointer unifier with initial capacity.
41    #[inline]
42    pub fn with_capacity(capacity: usize) -> Self {
43        Self {
44            shared_pointers: HashMap::with_capacity_and_hasher(
45                capacity,
46                Default::default(),
47            ),
48        }
49    }
50}
51
52impl fmt::Debug for Pool {
53    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
54        f.debug_map().entries(self.shared_pointers.iter()).finish()
55    }
56}
57
58#[derive(Debug)]
59struct NotStarted;
60
61impl fmt::Display for NotStarted {
62    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
63        write!(f, "shared pointer was not started pooling")
64    }
65}
66
67impl Error for NotStarted {}
68
69#[derive(Debug)]
70struct AlreadyFinished;
71
72impl fmt::Display for AlreadyFinished {
73    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
74        write!(f, "shared pointer was already finished pooling")
75    }
76}
77
78impl Error for AlreadyFinished {}
79
80impl<E: Source> Pooling<E> for Pool {
81    fn start_pooling(&mut self, address: usize) -> PoolingState {
82        match self.shared_pointers.entry(address) {
83            Entry::Vacant(vacant) => {
84                vacant.insert(None);
85                PoolingState::Started
86            }
87            Entry::Occupied(occupied) => {
88                if let Some(shared) = occupied.get() {
89                    PoolingState::Finished(shared.ptr)
90                } else {
91                    PoolingState::Pending
92                }
93            }
94        }
95    }
96
97    unsafe fn finish_pooling(
98        &mut self,
99        address: usize,
100        ptr: ErasedPtr,
101        drop: unsafe fn(ErasedPtr),
102    ) -> Result<(), E> {
103        match self.shared_pointers.entry(address) {
104            Entry::Vacant(_) => fail!(NotStarted),
105            Entry::Occupied(mut occupied) => {
106                let inner = occupied.get_mut();
107                if inner.is_some() {
108                    fail!(AlreadyFinished)
109                } else {
110                    *inner = Some(SharedPointer { ptr, drop });
111                    Ok(())
112                }
113            }
114        }
115    }
116}
rkyv/de/pooling/alloc.rs

rkyv/de/pooling/
alloc.rs
