polars_utils/idx_map/
total_idx_map.rs

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.tuples.len() as IdxSize
37    }
38
39    pub fn is_empty(&self) -> bool {
40        self.tuples.is_empty()
41    }
42
43    pub fn get(&self, key: &K) -> Option<&V> {
44        let idx = self.get_index_of(key)?;
45        unsafe { Some(&self.tuples.get_unchecked(idx as usize).1) }
46    }
47
48    /// Gets the index by insertion order of the given key.
49    pub fn get_index_of(&self, key: &K) -> Option<IdxSize> {
50        let hash = self.random_state.tot_hash_one(key);
51        self.table
52            .find(hash, |i| unsafe {
53                let t = self.tuples.get_unchecked(*i as usize);
54                hash == self.random_state.tot_hash_one(&t.0) && key.tot_eq(&t.0)
55            })
56            .copied()
57    }
58
59    pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
60        let hash = self.random_state.tot_hash_one(&key);
61        let entry = self.table.entry(
62            hash,
63            |i| unsafe {
64                let t = self.tuples.get_unchecked(*i as usize);
65                hash == self.random_state.tot_hash_one(&t.0) && key.tot_eq(&t.0)
66            },
67            |i| unsafe {
68                let t = self.tuples.get_unchecked(*i as usize);
69                self.random_state.tot_hash_one(&t.0)
70            },
71        );
72
73        match entry {
74            TEntry::Occupied(o) => Entry::Occupied(OccupiedEntry {
75                entry: o,
76                tuples: &mut self.tuples,
77            }),
78            TEntry::Vacant(v) => Entry::Vacant(VacantEntry {
79                key,
80                entry: v,
81                tuples: &mut self.tuples,
82            }),
83        }
84    }
85
86    /// Insert a key which will never be mapped to. Returns the index of the entry.
87    ///
88    /// This is useful for entries which are handled externally.
89    pub fn push_unmapped_entry(&mut self, key: K, value: V) -> IdxSize {
90        let ret = self.tuples.len() as IdxSize;
91        self.tuples.push((key, value));
92        ret
93    }
94
95    /// Gets the key and value at the given index by insertion order.
96    #[inline(always)]
97    pub fn get_index(&self, idx: IdxSize) -> Option<(&K, &V)> {
98        let t = self.tuples.get(idx as usize)?;
99        Some((&t.0, &t.1))
100    }
101
102    /// Gets the key and value at the given index by insertion order.
103    ///
104    /// # Safety
105    /// The index must be less than len().
106    #[inline(always)]
107    pub unsafe fn get_index_unchecked(&self, idx: IdxSize) -> (&K, &V) {
108        let t = unsafe { self.tuples.get_unchecked(idx as usize) };
109        (&t.0, &t.1)
110    }
111
112    /// Iterates over the keys in insertion order.
113    pub fn iter_keys(&self) -> impl Iterator<Item = &K> {
114        self.tuples.iter().map(|t| &t.0)
115    }
116
117    /// Iterates over the values in insertion order.
118    pub fn iter_values(&self) -> impl Iterator<Item = &V> {
119        self.tuples.iter().map(|t| &t.1)
120    }
121}
122
123pub enum Entry<'a, K, V> {
124    Occupied(OccupiedEntry<'a, K, V>),
125    Vacant(VacantEntry<'a, K, V>),
126}
127
128pub struct OccupiedEntry<'a, K, V> {
129    entry: TOccupiedEntry<'a, IdxSize>,
130    tuples: &'a mut Vec<(K, V)>,
131}
132
133impl<'a, K, V> OccupiedEntry<'a, K, V> {
134    pub fn index(&self) -> IdxSize {
135        *self.entry.get()
136    }
137
138    pub fn into_mut(self) -> &'a mut V {
139        let idx = self.index();
140        unsafe { &mut self.tuples.get_unchecked_mut(idx as usize).1 }
141    }
142}
143
144pub struct VacantEntry<'a, K, V> {
145    key: K,
146    entry: TVacantEntry<'a, IdxSize>,
147    tuples: &'a mut Vec<(K, V)>,
148}
149
150impl<'a, K, V> VacantEntry<'a, K, V> {
151    pub fn index(&self) -> IdxSize {
152        self.tuples.len() as IdxSize
153    }
154
155    pub fn insert(self, value: V) -> &'a mut V {
156        unsafe {
157            let tuple_idx: IdxSize = self.tuples.len().try_into().unwrap();
158            self.tuples.push((self.key, value));
159            self.entry.insert(tuple_idx);
160            &mut self.tuples.last_mut().unwrap_unchecked().1
161        }
162    }
163}
polars_utils/idx_map/total_idx_map.rs

polars_utils/idx_map/
total_idx_map.rs
