intern_mint/
pool.rs

1use std::{ops::Deref, sync::LazyLock};
2
3use hashbrown::{HashTable, hash_table::Entry};
4use parking_lot::{Mutex, MutexGuard};
5use triomphe::Arc;
6
7type LockedShard = HashTable<Arc<[u8]>>;
8type Shard = Mutex<LockedShard>;
9
10#[derive(Debug, Default, Clone, Copy)]
11pub struct MemoryUsage {
12    pub len: usize,
13    pub capacity: usize,
14}
15
16pub(crate) struct ShardedSet {
17    pub(crate) shift: usize,
18    pub(crate) hash_builder: ahash::RandomState,
19    pub(crate) shards: Box<[Shard]>,
20}
21
22impl ShardedSet {
23    fn get_hash_and_shard(&self, value: &[u8]) -> (u64, MutexGuard<LockedShard>) {
24        // hash before locking
25        let hash = self.hash_builder.hash_one(value);
26        // copied from https://github.com/xacrimon/dashmap/blob/366ce7e7872866a06de66eb95002fa6cf2c117a7/src/lib.rs#L419
27        let idx = ((hash << 7) >> self.shift) as usize;
28        let shard = self.shards[idx].lock();
29        (hash, shard)
30    }
31
32    fn hasher(&self, value: &Arc<[u8]>) -> u64 {
33        self.hash_builder.hash_one(value.deref())
34    }
35
36    pub(crate) fn get_from_existing_ref(&self, value: &[u8]) -> Option<Arc<[u8]>> {
37        let (hash, shard) = self.get_hash_and_shard(value);
38        shard
39            .find(hash, |o| std::ptr::addr_eq(o.as_ptr(), value.as_ptr()))
40            .cloned()
41    }
42
43    pub(crate) fn get_or_insert(&self, value: &[u8]) -> Arc<[u8]> {
44        let (hash, mut shard) = self.get_hash_and_shard(value);
45
46        shard
47            .entry(hash, |o| o.deref() == value, |o| self.hasher(o))
48            .or_insert_with(|| Arc::from(value))
49            .get()
50            .clone()
51    }
52
53    /// Only try to remove values from the pool when the reference count is two
54    /// one for the given [value] and another for the reference in the pool
55    pub(crate) fn remove_if_needed(&self, value: &Arc<[u8]>) {
56        const MINIMUM_STRONG_COUNT: usize = 2;
57
58        if Arc::strong_count(value) > MINIMUM_STRONG_COUNT {
59            return;
60        }
61
62        let (hash, mut shard) = self.get_hash_and_shard(value);
63
64        let entry = shard.entry(
65            hash,
66            |o| std::ptr::addr_eq(o.as_ptr(), value.as_ptr()),
67            |o| self.hasher(o),
68        );
69
70        if let Entry::Occupied(entry) = entry
71            && Arc::strong_count(entry.get()) <= MINIMUM_STRONG_COUNT
72        {
73            entry.remove();
74        }
75    }
76
77    pub(crate) fn is_empty(&self) -> bool {
78        self.len() == 0
79    }
80
81    pub(crate) fn len(&self) -> usize {
82        self.shards.iter().map(|o| o.lock().len()).sum()
83    }
84
85    pub(crate) fn capacity(&self) -> usize {
86        self.shards.iter().map(|o| o.lock().capacity()).sum()
87    }
88
89    pub(crate) fn get_memory_usage(&self) -> MemoryUsage {
90        self.shards
91            .iter()
92            .map(|o| {
93                let o = o.lock();
94                MemoryUsage {
95                    len: o.len(),
96                    capacity: o.capacity(),
97                }
98            })
99            .reduce(|acc, o| MemoryUsage {
100                len: acc.len + o.len,
101                capacity: acc.capacity + o.capacity,
102            })
103            .unwrap_or_default()
104    }
105
106    pub(crate) fn shrink_to_fit(&self) {
107        for shard in self.shards.iter() {
108            shard.lock().shrink_to_fit(|o| self.hasher(o));
109        }
110    }
111}
112
113impl Default for ShardedSet {
114    fn default() -> Self {
115        // copied from https://github.com/xacrimon/dashmap/blob/366ce7e7872866a06de66eb95002fa6cf2c117a7/src/lib.rs#L63
116        static DEFAULT_SHARDS_COUNT: LazyLock<usize> = LazyLock::new(|| {
117            (std::thread::available_parallelism().map_or(1, usize::from) * 4).next_power_of_two()
118        });
119
120        // copied from https://github.com/xacrimon/dashmap/blob/366ce7e7872866a06de66eb95002fa6cf2c117a7/src/lib.rs#L269
121        let shift =
122            (std::mem::size_of::<usize>() * 8) - DEFAULT_SHARDS_COUNT.trailing_zeros() as usize;
123
124        Self {
125            shift,
126            hash_builder: Default::default(),
127            shards: (0..*DEFAULT_SHARDS_COUNT)
128                .map(|_| Default::default())
129                .collect(),
130        }
131    }
132}
133
134pub(crate) static POOL: LazyLock<ShardedSet> = LazyLock::new(Default::default);
135
136pub fn is_empty() -> bool {
137    POOL.is_empty()
138}
139
140pub fn len() -> usize {
141    POOL.len()
142}
143
144pub fn capacity() -> usize {
145    POOL.capacity()
146}
147
148pub fn get_memory_usage() -> MemoryUsage {
149    POOL.get_memory_usage()
150}
151
152pub fn shrink_to_fit() {
153    POOL.shrink_to_fit();
154}
intern_mint/pool.rs

intern_mint/
pool.rs
