pb_atomic_hash_map/
lib.rs

1use std::{alloc::{alloc, dealloc, Layout}, borrow::Borrow, hash::{Hash, Hasher as _}, marker::PhantomData, mem::MaybeUninit, ops::Deref, ptr::NonNull, sync::{atomic::{AtomicUsize, Ordering}, Arc}};
2
3use pb_atomic_linked_list::{prelude::AtomicLinkedList as _, AtomicLinkedList};
4
5
6fn hash<K: Hash>(key: &K) -> usize {
7    let mut hasher = std::collections::hash_map::DefaultHasher::new();
8    key.hash(&mut hasher);
9    let hashed_key = hasher.finish() as usize;
10    hashed_key
11}
12
13pub struct Entry<K: Hash,V> {
14    pub key: K,
15    pub value: V
16}
17
18impl<K: Hash, V> Entry<K,V> {
19    pub fn hash(&self) -> usize {
20        hash(&self.key)
21    }
22}
23
24struct Bucket<K: Hash,V>(AtomicLinkedList<Entry<K,V>>);
25
26impl<K: Hash, V> Bucket<K, V> {
27    pub fn insert(&mut self, key: K, value: V) -> bool {
28        unsafe {
29            let exists = self.get_raw_entry_ptr(&key).is_some();
30            self.0.insert(Entry{key, value});
31            !exists
32        }
33    }
34
35    unsafe fn get_raw_entry_ptr(&self, key: &K) -> Option<NonNull<Entry<K,V>>> {
36        let hashed_key = hash(key);
37        self.0.iter()
38        .map(|rf| std::ptr::from_ref(rf) as *mut Entry<K,V>)
39        .map(|ptr| NonNull::new_unchecked(ptr))
40        .filter(|ptr| ptr.as_ref().hash() == hashed_key)
41        .last()
42    }
43}
44
45impl<K: Hash, V> Bucket<K, V> {
46    fn new() -> Self {
47        Self(AtomicLinkedList::new())
48    }
49}
50
51/// A hash table
52struct Table<K: Hash,V> {
53    buckets: NonNull<[Bucket<K,V>]>,
54    buckets_layout: Layout,
55    capacity: usize,
56    length: AtomicUsize
57}
58
59unsafe impl<K: Hash,V> Sync for Table<K,V> {}
60unsafe impl<K: Hash,V> Send for Table<K,V> {}
61
62impl<K: Hash, V> Drop for Table<K,V> {
63    fn drop(&mut self) {
64        unsafe {
65            for bucket in self.buckets.as_mut().iter_mut() {
66                std::ptr::from_mut(bucket).drop_in_place();
67            }
68
69            dealloc(self.buckets.cast().as_ptr(), self.buckets_layout);
70        }
71    }
72}
73
74impl<K: Hash,V> Table<K,V> {
75    pub fn new(capacity: usize) -> Self {       
76        let buckets_layout = Layout::array::<Bucket<K,V>>(capacity).unwrap();
77
78        unsafe {
79            let raw_buckets_base_ptr = alloc(buckets_layout).cast::<MaybeUninit<Bucket<K,V>>>();
80            let raw_buckets_slice_ptr= std::ptr::slice_from_raw_parts_mut(
81                raw_buckets_base_ptr,
82                capacity
83            );
84
85            if let Some(mut_buckets_slice_ptr) = raw_buckets_slice_ptr.as_mut() {
86                // Initialise all buckets.
87                for bucket in mut_buckets_slice_ptr {
88                    bucket.write(Bucket::new());
89                }
90            }
91
92            let buckets = NonNull::new(std::mem::transmute(raw_buckets_slice_ptr)).unwrap();
93   
94            let length = AtomicUsize::new(0);
95
96            Self {
97                buckets, 
98                buckets_layout,
99                capacity, 
100                length
101            }
102        }
103        
104    }
105    
106
107    pub fn len(&self) -> usize {
108        return self.length.load(Ordering::Relaxed)
109    }
110
111    /// Insert a new value in the bucket.
112    pub fn insert(&self, key: K, value: V) {
113        unsafe {
114            let bucket = self.get_bucket_ptr(&key).as_mut();
115            if bucket.insert(key, value) {
116                self.length.fetch_add(1, Ordering::Relaxed);
117            }
118        }
119    }
120
121    /// Retrieve the bucket which might contain the value behind the key.
122    unsafe fn get_bucket_ptr(&self, key: &K) -> NonNull<Bucket<K,V>> {
123        let mut hasher = std::collections::hash_map::DefaultHasher::new();
124        key.hash(&mut hasher);
125        let hashed_key = hasher.finish() as usize;
126        let bucket_key = hashed_key % self.capacity;
127        let bucket_ptr = std::ptr::from_mut(
128            self
129                .buckets
130                .as_ptr()
131                .as_mut()
132                .unwrap()
133                .get_mut(bucket_key)
134                .unwrap()
135        );
136
137        return NonNull::new(bucket_ptr).unwrap()
138    }
139
140
141    unsafe fn get_raw_entry_ptr(&self, key: &K) -> Option<NonNull<Entry<K,V>>> {
142        self
143            .get_bucket_ptr(key)
144            .as_ref()
145            .get_raw_entry_ptr(key)
146    }
147}
148
149pub struct ValueIter<'a, K: Hash + 'a, V: 'a>(EntryIter<'a, K, V>);
150
151impl<'a, K: Hash + 'a, V: 'a> Iterator for ValueIter<'a, K, V> {
152    type Item = &'a V;
153
154    fn next(&mut self) -> Option<Self::Item> {
155        self.0.next().map(|entry| &entry.value)
156    }
157}
158
159
160pub struct Iter<'a, K: Hash + 'a, V: 'a>(EntryIter<'a, K, V>);
161
162impl<'a, K: Hash + 'a, V: 'a> Iterator for Iter<'a, K, V> {
163    type Item = (&'a K, &'a V);
164
165    fn next(&mut self) -> Option<Self::Item> {
166        self.0.next().map(|entry| (&entry.key, &entry.value))
167    }
168}
169
170struct EntryIter<'a, K: Hash + 'a, V: 'a> {
171    _phantom: PhantomData<&'a ()>,
172    buckets: NonNull<[Bucket<K, V>]>,
173    current_bucket_iter: Option<pb_atomic_linked_list::Iter<'a, Entry<K, V>>>,
174    current_bucket_key: usize
175}
176
177impl<'a, K: Hash + 'a, V: 'a> Iterator for EntryIter<'a, K, V> {
178    type Item = &'a Entry<K, V>;
179
180    fn next(&mut self) -> Option<Self::Item> {
181        unsafe {
182            if let Some(iter) = &mut self.current_bucket_iter {
183                if let Some(entry) = iter.next() {
184                    return Some(entry)
185                }
186                else {
187                    self.current_bucket_key += 1;
188                    if self.current_bucket_key >= self.buckets.len() {
189                        return None
190                    }                    
191                    let bucket = self.buckets.as_ref().get(self.current_bucket_key).unwrap();
192                    self.current_bucket_iter = Some(bucket.0.iter());
193                    self.next()
194                }
195
196            } else {
197                return None
198            }
199        }
200    }
201}
202
203#[derive(Debug, PartialEq, Eq)]
204pub struct ValueRef<'a, T>(&'a T);
205
206impl<'a, T: Clone> ValueRef<'a, T> {
207    pub fn to_owned(value: Self) -> T {
208        value.0.clone()
209    }
210}
211
212impl<'a, T> Deref for ValueRef<'a, T> {
213    type Target = T;
214    
215    fn deref(&self) -> &Self::Target {
216        self.0
217    }    
218}
219
220pub struct AtomicHashMap<K: Hash, V>(Arc<Table<K, V>>);
221
222impl<K: Hash, V> Clone for AtomicHashMap<K, V> {
223    fn clone(&self) -> Self {
224        Self(self.0.clone())
225    }
226}
227
228impl<K: Hash, V> AtomicHashMap<K, V> {
229    /// Creates a new atomic hash map
230    pub fn new(capacity: usize) -> Self {
231        Self(
232            Arc::new(
233                Table::new(capacity)
234            )
235        )
236    }
237
238    fn iter_entries(&self) -> EntryIter<'_, K, V> {
239        unsafe {
240            EntryIter {
241                _phantom: PhantomData,
242                buckets: self.0.buckets,
243                current_bucket_iter: self.0.buckets.as_ref().get(0).map(|bucket| bucket.0.iter()),
244                current_bucket_key: 0  
245            }
246        }
247    }
248
249    /// Iterate over all values in the hash map
250    pub fn iter_values(&self) -> ValueIter<'_, K, V> {
251        ValueIter(self.iter_entries())
252    }
253
254    /// Iterate over all key/value pair in the hash map 
255    pub fn iter(&self) -> Iter<'_, K, V> {
256        Iter(self.iter_entries())
257    }
258
259    /// Returns the length 
260    pub fn len(&self) -> usize {
261        self.0.len()
262    }
263
264    /// Insert a new value 
265    /// 
266    /// If a value is already stored in the hash map,
267    /// this will not overwrite its content, the new version will
268    /// appended in the linked list.
269    pub fn insert(&mut self, key: K, value: V) {
270        self.0.insert(key, value);
271    }
272
273    /// Borrow the value behind the key
274    /// 
275    /// Always returns the last version in the linked list.
276    pub fn borrow<'a, Q: Borrow<K>>(&'a self, key: Q) -> Option<ValueRef<'a, V>> {
277        unsafe {
278            self
279            .get_raw_value_ptr(key.borrow())
280            .map(|value| value.as_ref())
281            .map(ValueRef)
282        }
283    }
284
285    /// Get the raw pointer to the value.
286    /// 
287    /// # Unsafe
288    /// Because of obvious reasons
289    pub unsafe fn get_raw_value_ptr(&self, key: &K) -> Option<NonNull<V>> {
290        self.0
291        .get_raw_entry_ptr(key)
292        .map(|mut entry| 
293            NonNull::new(
294                std::ptr::from_mut(
295                    &mut entry.as_mut().value
296                )
297            ).unwrap()
298        )
299    }
300}
301
302#[cfg(test)]
303mod tests {
304    use std::{collections::HashSet, thread};
305
306    use crate::{AtomicHashMap, ValueRef};
307
308    #[test]
309    fn test_borrow_unexisting_value() {
310        let mut map = AtomicHashMap::<u32, u32>::new(100);
311        map.insert(20, 30); 
312        assert_eq!(map.borrow(&10), None);
313    }
314
315    #[test]
316    fn test_borrow_existing_value() {
317        let mut map = AtomicHashMap::<u32, u32>::new(100);
318        map.insert(10, 20); 
319        assert_eq!(map.borrow(&10).map(ValueRef::to_owned), Some(20));
320    }
321
322    #[test]
323    fn test_iter() {
324        let mut map = AtomicHashMap::<u32, u32>::new(100);
325        let mut expected_entries = HashSet::<(u32, u32)>::default();
326        for i in 0..1_000 {
327            map.insert(i, i);
328            expected_entries.insert((i,i));
329        }
330
331        let got = map.iter().map(|(k, v)| (*k, *v)).collect::<HashSet::<_>>();
332        assert_eq!(got, expected_entries)
333    }
334
335    #[test]
336    fn test_multiple_threads() {
337        let map = AtomicHashMap::<u32, u32>::new(100);
338        let mut map1 = map.clone();
339        let mut map2 = map.clone();
340
341        let mut expected_entries = HashSet::<(u32, u32)>::default();
342        for i in 0..=20_000 {
343            expected_entries.insert((i, i));
344        }
345
346
347        let j1 = thread::spawn(move || {
348            for i in 0..=10_000 {
349                map1.insert(i, i);
350            }
351        });
352
353        let j2 = thread::spawn(move || {
354            for i in 10_001..=20_000 {
355                map2.insert(i, i);
356            }
357        });
358
359        j1.join().unwrap();
360        j2.join().unwrap();
361
362        let got = map.iter().map(|(k, v)| (*k, *v)).collect::<HashSet::<_>>();
363        assert_eq!(got, expected_entries)
364    }
365}
pb_atomic_hash_map/lib.rs

pb_atomic_hash_map/
lib.rs
