Skip to main content

osom_lib_hash_tables/bytell/hash_table/
bytell_mutable.rs

1use core::{borrow::Borrow, hash::Hash, marker::PhantomData};
2
3use osom_lib_macros::debug_check_or_release_hint;
4use osom_lib_primitives::{kvp::KVP, power_of_two::PowerOfTwo32};
5
6use crate::{
7    bytell::{
8        configuration::BytellConfig,
9        hash_table::{BytellHashTable, block_layout::BlockLayoutHolder, control_byte::ControlByte, entry::Entry},
10    },
11    errors::HashTableError,
12    helpers::{ptr_to_mut, ptr_to_ref},
13    traits::MutableHashTable,
14};
15
16struct BytellMutableIter<'a, TKey: 'a, TValue: 'a, TConfig> {
17    data: *mut u8,
18    last_element_index: u32,
19    blocks_count: PowerOfTwo32,
20    _marker: PhantomData<(&'a TKey, &'a mut TValue, TConfig)>,
21}
22
23impl<'a, TKey: 'a, TValue: 'a, TConfig> BytellMutableIter<'a, TKey, TValue, TConfig>
24where
25    TKey: Eq + Hash,
26    TConfig: BytellConfig,
27{
28    pub const fn from_hash_table(
29        table: &'a BytellHashTable<TKey, TValue, TConfig>,
30    ) -> BytellMutableIter<'a, TKey, TValue, TConfig> {
31        Self {
32            data: table.data,
33            last_element_index: 0,
34            blocks_count: table.blocks_count,
35            _marker: PhantomData,
36        }
37    }
38}
39
40impl<'a, TKey: 'a, TValue: 'a, TConfig> Iterator for BytellMutableIter<'a, TKey, TValue, TConfig> {
41    type Item = KVP<&'a TKey, &'a mut TValue>;
42
43    fn next(&mut self) -> Option<Self::Item> {
44        let elements_count = BlockLayoutHolder::<TKey, TValue>::LAYOUT.block_capacity().value();
45        let capacity = self.blocks_count.value() * elements_count;
46        let mut el_idx = self.last_element_index;
47
48        loop {
49            unsafe {
50                if el_idx == capacity {
51                    self.last_element_index = el_idx;
52                    return None;
53                }
54
55                let entry = Entry::<TKey, TValue>::new(self.data, self.blocks_count, el_idx as usize);
56                el_idx += 1;
57
58                if !(*entry.control_byte()).contains_data() {
59                    continue;
60                }
61
62                let kvp = ptr_to_mut!(entry.kvp());
63                self.last_element_index = el_idx;
64                return Some(kvp.as_mut_kvp());
65            }
66        }
67    }
68}
69
70impl<TKey, TValue, TConfig> MutableHashTable<TKey, TValue> for BytellHashTable<TKey, TValue, TConfig>
71where
72    TKey: Eq + Hash,
73    TConfig: BytellConfig,
74{
75    fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut TValue>
76    where
77        TKey: Borrow<Q>,
78        Q: Eq + Hash + ?Sized,
79    {
80        self.get_key_value_mut(key).map(|kvp| kvp.value)
81    }
82
83    fn get_key_value_mut<Q>(&mut self, key: &Q) -> Option<KVP<&TKey, &mut TValue>>
84    where
85        TKey: Borrow<Q>,
86        Q: Eq + Hash + ?Sized,
87    {
88        unsafe {
89            let result = self.get_key_value_raw(key)?;
90            let (key_ptr, value_ptr) = KVP::unpack_ptr(result);
91            Some(KVP {
92                key: key_ptr.as_ref_unchecked(),
93                value: value_ptr.as_mut_unchecked(),
94            })
95        }
96    }
97
98    fn remove_entry<Q>(&mut self, key: &Q) -> Option<KVP<TKey, TValue>>
99    where
100        TKey: Borrow<Q>,
101        Q: Eq + Hash + ?Sized,
102    {
103        if self.blocks_count.value() == 0 {
104            return None;
105        }
106
107        unsafe {
108            let mut entry = self.get_entry_by_key(key);
109            let control_byte = ptr_to_ref!(entry.control_byte());
110            if !control_byte.is_direct_hit() {
111                return None;
112            }
113            let mut entry_parent = None;
114
115            loop {
116                if ptr_to_ref!(entry.kvp()).key.borrow() == key {
117                    break;
118                }
119
120                if let Some(next) = entry.next_link() {
121                    entry_parent = Some(entry);
122                    entry = next;
123                    continue;
124                }
125
126                return None;
127            }
128
129            // Swap found entry with the tail
130            {
131                let mut tail_parent = None;
132                let tail = {
133                    let mut current_entry = entry.clone();
134                    loop {
135                        if let Some(next) = current_entry.next_link() {
136                            tail_parent = Some(current_entry);
137                            current_entry = next;
138                        } else {
139                            break current_entry;
140                        }
141                    }
142                };
143
144                if tail != entry {
145                    let tail_kvp = ptr_to_mut!(tail.kvp());
146                    let entry_kvp = ptr_to_mut!(entry.kvp());
147                    core::mem::swap(tail_kvp, entry_kvp);
148                    entry_parent = tail_parent;
149                    entry = tail;
150                }
151            }
152
153            if let Some(acutal_entry_parent) = entry_parent {
154                ptr_to_mut!(acutal_entry_parent.control_byte()).set_distance_index(0);
155            }
156
157            let kvp = entry.kvp().read();
158            *ptr_to_mut!(entry.control_byte()) = ControlByte::EMPTY;
159            self.elements_count -= 1;
160            Some(kvp)
161        }
162    }
163
164    fn try_insert_or_update_with<'a, FAdd, FUpdate>(
165        &mut self,
166        key: TKey,
167        adder: FAdd,
168        updater: FUpdate,
169    ) -> Result<&mut TValue, HashTableError>
170    where
171        FAdd: FnOnce() -> TValue,
172        FUpdate: FnOnce(&mut TValue),
173    {
174        if self.blocks_count.value() == 0 {
175            self.grow()?;
176        }
177
178        let mut counter = 0;
179
180        #[allow(clippy::redundant_else)]
181        'start: loop {
182            counter += 1;
183            assert!(counter < 30, "Tried to insert over 30 times. Giving up.");
184
185            unsafe {
186                let entry = self.get_entry_by_key(&key);
187                let control_byte = *entry.control_byte();
188
189                debug_check_or_release_hint!(
190                    control_byte != ControlByte::RESERVED,
191                    "Got reserved here? Should not happen."
192                );
193
194                if control_byte.is_direct_hit() {
195                    let mut entry = entry;
196                    loop {
197                        let kvp = ptr_to_mut!(entry.kvp());
198                        if kvp.key == key {
199                            updater(&mut kvp.value);
200                            return Ok(&mut kvp.value);
201                        }
202
203                        if let Some(next) = entry.next_link() {
204                            entry = next;
205                            continue;
206                        }
207
208                        break;
209                    }
210
211                    if self.should_grow() {
212                        self.grow()?;
213                        continue 'start;
214                    }
215
216                    let Some((free_entry, free_distance_index)) = self.search_for_free_entry(&entry) else {
217                        self.grow()?;
218                        continue 'start;
219                    };
220
221                    *free_entry.control_byte() = ControlByte::NEW_TAIL;
222                    ptr_to_mut!(entry.control_byte()).set_distance_index(free_distance_index);
223                    let new_kvp = KVP { key, value: adder() };
224                    free_entry.kvp().write(new_kvp);
225                    self.elements_count += 1;
226
227                    return Ok(&mut ptr_to_mut!(free_entry.kvp()).value);
228                } else {
229                    if self.should_grow() {
230                        self.grow()?;
231                        continue 'start;
232                    }
233
234                    if control_byte != ControlByte::EMPTY {
235                        // It is neither direct hit nor empty. Which means we hit storage entry.
236                        // We need to move around the linked list this entry belongs to.
237
238                        let mut current_parent = self.find_parent_for_storage_entry(&entry);
239                        let mut current_entry = entry.clone();
240
241                        let mut first_iteration = true;
242                        loop {
243                            let Some((free_entry, free_distance_index)) = self.search_for_free_entry(&current_parent)
244                            else {
245                                self.grow()?;
246                                continue 'start;
247                            };
248
249                            *ptr_to_mut!(free_entry.control_byte()) = ControlByte::NEW_TAIL;
250                            free_entry.kvp().write(current_entry.kvp().read());
251                            ptr_to_mut!(current_parent.control_byte()).set_distance_index(free_distance_index);
252
253                            let new_control = if first_iteration {
254                                first_iteration = false;
255                                ControlByte::RESERVED
256                            } else {
257                                ControlByte::EMPTY
258                            };
259
260                            let next_link = current_entry.next_link();
261                            *current_entry.control_byte() = new_control;
262
263                            let Some(next) = next_link else {
264                                break;
265                            };
266
267                            current_parent = free_entry;
268                            current_entry = next;
269                        }
270                    }
271
272                    self.elements_count += 1;
273                    *entry.control_byte() = ControlByte::NEW_DIRECT_HIT;
274                    let kvp = KVP { key, value: adder() };
275                    entry.kvp().write(kvp);
276                    return Ok(&mut ptr_to_mut!(entry.kvp()).value);
277                }
278            }
279        }
280    }
281
282    #[inline(always)]
283    fn iter_mut<'a>(&'a mut self) -> impl Iterator<Item = KVP<&'a TKey, &'a mut TValue>>
284    where
285        TKey: 'a,
286        TValue: 'a,
287        Self: 'a,
288    {
289        BytellMutableIter::from_hash_table(self)
290    }
291}