Skip to main content

osom_lib_hash_tables/abseil/hash_table/
abseil_mutable.rs

1use core::borrow::Borrow;
2use core::hash::Hash;
3
4use osom_lib_primitives::{kvp::KVP, length::Length};
5
6use crate::abseil::hash_table::abseil_block::CONTROL_BYTE_TOMBSTONE;
7use crate::abseil::hash_table::abseil_unsafe_iter::AbseilUnsafeMutIter;
8use crate::abseil::hash_table::platform::{PlatformImpl, PlatformOps as _};
9use crate::abseil::utils::probe_block_indexes;
10use crate::errors::HashTableError;
11use crate::helpers::{ptr_to_mut, ptr_to_ref};
12use crate::traits::MutableHashTable;
13
14use crate::abseil::{configuration::AbseilConfig, hash_table::AbseilHashTable};
15
16impl<TKey, TValue, TConfig> MutableHashTable<TKey, TValue> for AbseilHashTable<TKey, TValue, TConfig>
17where
18    TKey: Eq + Hash,
19    TConfig: AbseilConfig,
20{
21    fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut TValue>
22    where
23        TKey: Borrow<Q>,
24        Q: Eq + Hash + ?Sized,
25    {
26        self.get_key_value_mut(key).map(|kvp| kvp.value)
27    }
28
29    fn get_key_value_mut<Q>(&mut self, key: &Q) -> Option<KVP<&TKey, &mut TValue>>
30    where
31        TKey: Borrow<Q>,
32        Q: Eq + Hash + ?Sized,
33    {
34        let result = self.get_key_value_raw(key)?;
35        unsafe {
36            let (key_ptr, value_ptr) = KVP::unpack_ptr(result);
37            Some(KVP {
38                key: key_ptr.as_ref_unchecked(),
39                value: value_ptr.as_mut_unchecked(),
40            })
41        }
42    }
43
44    fn remove_entry<Q>(&mut self, key: &Q) -> Option<KVP<TKey, TValue>>
45    where
46        TKey: Borrow<Q>,
47        Q: Eq + Hash + ?Sized,
48    {
49        let blocks_count = self.blocks_count();
50        let (h1, h2) = self.config.calculate_partial_hashes(key);
51
52        for group_index in probe_block_indexes(h1, blocks_count) {
53            let block = self.get_block_by_index(group_index);
54            let control_bytes = ptr_to_mut!(block.control_block_ptr());
55            let mut scan_data = PlatformImpl::matching_block_scan(control_bytes, h2);
56            for matching_index in scan_data.matching_indexes {
57                let kvp_ptr = block.key_value_pair_at_index(matching_index);
58                let kvp = ptr_to_ref!(kvp_ptr);
59                if kvp.key.borrow() == key {
60                    unsafe { core::hint::assert_unchecked(matching_index < control_bytes.len()) };
61                    control_bytes[matching_index] = CONTROL_BYTE_TOMBSTONE;
62                    let kvp = unsafe { kvp_ptr.read() };
63                    unsafe {
64                        self.elements_count = Length::new_unchecked(self.elements_count.as_u32().unchecked_sub(1));
65                    }
66                    return Some(kvp);
67                }
68            }
69
70            if scan_data.empty_buckets.next().is_some() {
71                return None;
72            }
73        }
74
75        None
76    }
77
78    fn try_insert_or_update_with<FAdd, FUpdate>(
79        &mut self,
80        key: TKey,
81        adder: FAdd,
82        updater: FUpdate,
83    ) -> Result<&mut TValue, HashTableError>
84    where
85        FAdd: FnOnce() -> TValue,
86        FUpdate: FnOnce(&mut TValue),
87    {
88        if self.remaining_capacity == Length::ZERO {
89            self.grow_for_size(self.elements_count + 1)?;
90        }
91
92        let (h1, h2) = self.config.calculate_partial_hashes(&key);
93        let blocks_count = self.blocks_count();
94
95        // Single pass: search for an existing key while simultaneously tracking
96        // the first tombstone and watching for an empty slot (which terminates
97        // the probe chain and proves the key is absent).
98        let mut first_tombstone: Option<(usize, usize)> = None; // (group_index, slot_index)
99
100        for group_index in probe_block_indexes(h1, blocks_count) {
101            let block = self.get_block_by_index(group_index);
102            let control_bytes = ptr_to_mut!(block.control_block_ptr());
103            let mut scan_result = PlatformImpl::full_block_scan(control_bytes, h2);
104
105            for matching_index in scan_result.matching_indexes {
106                let kvp_ptr = block.key_value_pair_at_index(matching_index);
107                let kvp = ptr_to_mut!(kvp_ptr);
108                if kvp.key == key {
109                    updater(&mut kvp.value);
110                    return Ok(&mut kvp.value);
111                }
112            }
113
114            if first_tombstone.is_none()
115                && let Some(tombstone_idx) = scan_result.tombstones.next()
116            {
117                first_tombstone = Some((group_index, tombstone_idx));
118            }
119
120            if let Some(empty_idx) = scan_result.empty_buckets.next() {
121                // Empty slot proves the key is absent. Prefer tombstone (reuse deleted slot)
122                // over the empty slot when available.
123                let (target_group, target_slot) = if let Some((tg, ts)) = first_tombstone {
124                    (tg, ts)
125                } else {
126                    self.remaining_capacity =
127                        unsafe { Length::new_unchecked(self.remaining_capacity.as_u32().unchecked_sub(1)) };
128                    (group_index, empty_idx)
129                };
130
131                let target_block = self.get_block_by_index(target_group);
132                let target_ctrl = ptr_to_mut!(target_block.control_block_ptr());
133                let target_kvp_ptr = target_block.key_value_pair_at_index(target_slot);
134
135                unsafe {
136                    *target_ctrl.get_unchecked_mut(target_slot) = h2;
137                    target_kvp_ptr.write(KVP { key, value: adder() });
138                };
139
140                unsafe {
141                    self.elements_count = Length::new_unchecked(self.elements_count.as_u32().unchecked_add(1));
142                }
143
144                let kvp = ptr_to_mut!(target_kvp_ptr);
145
146                // Safety: the pointer is valid for the lifetime of self.
147                return Ok(&mut kvp.value);
148            }
149        }
150
151        // With correct remaining_capacity management there is always at least one
152        // empty slot in the table, so this path is unreachable.
153        unreachable!("no empty slot found despite remaining_capacity > 0")
154    }
155
156    #[inline(always)]
157    fn iter_mut<'a>(&'a mut self) -> impl Iterator<Item = KVP<&'a TKey, &'a mut TValue>> + 'a
158    where
159        TKey: 'a,
160        TValue: 'a,
161        Self: 'a,
162    {
163        AbseilUnsafeMutIter::from_hash_table(self).map(|kvp| ptr_to_mut!(kvp).as_mut_kvp())
164    }
165}