polars_utils/idx_map/
total_idx_map.rs

Help
1use hashbrown::hash_table::{
2    Entry as TEntry, HashTable, OccupiedEntry as TOccupiedEntry, VacantEntry as TVacantEntry,
3};
4
5use crate::IdxSize;
6use crate::aliases::PlRandomState;
7use crate::total_ord::{BuildHasherTotalExt, TotalEq, TotalHash};
8
9/// An IndexMap where the keys are hashed and compared with TotalOrd/TotalEq.
10pub struct TotalIndexMap<K, V> {
11    table: HashTable<IdxSize>,
12    tuples: Vec<(K, V)>,
13    random_state: PlRandomState,
14}
15
16impl<K, V> Default for TotalIndexMap<K, V> {
17    fn default() -> Self {
18        Self {
19            table: HashTable::new(),
20            tuples: Vec::new(),
21            random_state: PlRandomState::default(),
22        }
23    }
24}
25
26impl<K: TotalHash + TotalEq, V> TotalIndexMap<K, V> {
27    pub fn reserve(&mut self, additional: usize) {
28        self.table.reserve(additional, |i| unsafe {
29            let tuple = self.tuples.get_unchecked(*i as usize);
30            self.random_state.tot_hash_one(&tuple.0)
31        });
32        self.tuples.reserve(additional);
33    }
34
35    pub fn len(&self) -> IdxSize {
36        self.table.len() as IdxSize
37    }
38
39    pub fn is_empty(&self) -> bool {
40        self.table.is_empty()
41    }
42
43    pub fn get(&self, key: &K) -> Option<&V> {
44        let hash = self.random_state.tot_hash_one(key);
45        let idx = self.table.find(hash, |i| unsafe {
46            let t = self.tuples.get_unchecked(*i as usize);
47            hash == self.random_state.tot_hash_one(&t.0) && key.tot_eq(&t.0)
48        })?;
49        unsafe { Some(&self.tuples.get_unchecked(*idx as usize).1) }
50    }
51
52    pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
53        let hash = self.random_state.tot_hash_one(&key);
54        let entry = self.table.entry(
55            hash,
56            |i| unsafe {
57                let t = self.tuples.get_unchecked(*i as usize);
58                hash == self.random_state.tot_hash_one(&t.0) && key.tot_eq(&t.0)
59            },
60            |i| unsafe {
61                let t = self.tuples.get_unchecked(*i as usize);
62                self.random_state.tot_hash_one(&t.0)
63            },
64        );
65
66        match entry {
67            TEntry::Occupied(o) => Entry::Occupied(OccupiedEntry {
68                entry: o,
69                tuples: &mut self.tuples,
70            }),
71            TEntry::Vacant(v) => Entry::Vacant(VacantEntry {
72                key,
73                entry: v,
74                tuples: &mut self.tuples,
75            }),
76        }
77    }
78
79    /// Gets the key and value at the given index by insertion order.
80    #[inline(always)]
81    pub fn get_index(&self, idx: IdxSize) -> Option<(&K, &V)> {
82        let t = self.tuples.get(idx as usize)?;
83        Some((&t.0, &t.1))
84    }
85
86    /// Gets the key and value at the given index by insertion order.
87    ///
88    /// # Safety
89    /// The index must be less than len().
90    #[inline(always)]
91    pub unsafe fn get_index_unchecked(&self, idx: IdxSize) -> (&K, &V) {
92        let t = unsafe { self.tuples.get_unchecked(idx as usize) };
93        (&t.0, &t.1)
94    }
95
96    /// Iterates over the keys in insertion order.
97    pub fn iter_keys(&self) -> impl Iterator<Item = &K> {
98        self.tuples.iter().map(|t| &t.0)
99    }
100
101    /// Iterates over the values in insertion order.
102    pub fn iter_values(&self) -> impl Iterator<Item = &V> {
103        self.tuples.iter().map(|t| &t.1)
104    }
105}
106
107pub enum Entry<'a, K, V> {
108    Occupied(OccupiedEntry<'a, K, V>),
109    Vacant(VacantEntry<'a, K, V>),
110}
111
112pub struct OccupiedEntry<'a, K, V> {
113    entry: TOccupiedEntry<'a, IdxSize>,
114    tuples: &'a mut Vec<(K, V)>,
115}
116
117impl<'a, K, V> OccupiedEntry<'a, K, V> {
118    pub fn index(&self) -> IdxSize {
119        *self.entry.get()
120    }
121
122    pub fn into_mut(self) -> &'a mut V {
123        let idx = self.index();
124        unsafe { &mut self.tuples.get_unchecked_mut(idx as usize).1 }
125    }
126}
127
128pub struct VacantEntry<'a, K, V> {
129    key: K,
130    entry: TVacantEntry<'a, IdxSize>,
131    tuples: &'a mut Vec<(K, V)>,
132}
133
134impl<'a, K, V> VacantEntry<'a, K, V> {
135    pub fn index(&self) -> IdxSize {
136        self.tuples.len() as IdxSize
137    }
138
139    pub fn insert(self, value: V) -> &'a mut V {
140        unsafe {
141            let tuple_idx: IdxSize = self.tuples.len().try_into().unwrap();
142            self.tuples.push((self.key, value));
143            self.entry.insert(tuple_idx);
144            &mut self.tuples.last_mut().unwrap_unchecked().1
145        }
146    }
147}
polars_utils/idx_map/total_idx_map.rs

polars_utils/idx_map/
total_idx_map.rs
