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 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    /// Insert a key which will never be mapped to. Returns the index of the entry.
80    ///
81    /// This is useful for entries which are handled externally.
82    pub fn push_unmapped_entry(&mut self, key: K, value: V) -> IdxSize {
83        let ret = self.tuples.len() as IdxSize;
84        self.tuples.push((key, value));
85        ret
86    }
87
88    /// Gets the key and value at the given index by insertion order.
89    #[inline(always)]
90    pub fn get_index(&self, idx: IdxSize) -> Option<(&K, &V)> {
91        let t = self.tuples.get(idx as usize)?;
92        Some((&t.0, &t.1))
93    }
94
95    /// Gets the key and value at the given index by insertion order.
96    ///
97    /// # Safety
98    /// The index must be less than len().
99    #[inline(always)]
100    pub unsafe fn get_index_unchecked(&self, idx: IdxSize) -> (&K, &V) {
101        let t = unsafe { self.tuples.get_unchecked(idx as usize) };
102        (&t.0, &t.1)
103    }
104
105    /// Iterates over the keys in insertion order.
106    pub fn iter_keys(&self) -> impl Iterator<Item = &K> {
107        self.tuples.iter().map(|t| &t.0)
108    }
109
110    /// Iterates over the values in insertion order.
111    pub fn iter_values(&self) -> impl Iterator<Item = &V> {
112        self.tuples.iter().map(|t| &t.1)
113    }
114}
115
116pub enum Entry<'a, K, V> {
117    Occupied(OccupiedEntry<'a, K, V>),
118    Vacant(VacantEntry<'a, K, V>),
119}
120
121pub struct OccupiedEntry<'a, K, V> {
122    entry: TOccupiedEntry<'a, IdxSize>,
123    tuples: &'a mut Vec<(K, V)>,
124}
125
126impl<'a, K, V> OccupiedEntry<'a, K, V> {
127    pub fn index(&self) -> IdxSize {
128        *self.entry.get()
129    }
130
131    pub fn into_mut(self) -> &'a mut V {
132        let idx = self.index();
133        unsafe { &mut self.tuples.get_unchecked_mut(idx as usize).1 }
134    }
135}
136
137pub struct VacantEntry<'a, K, V> {
138    key: K,
139    entry: TVacantEntry<'a, IdxSize>,
140    tuples: &'a mut Vec<(K, V)>,
141}
142
143impl<'a, K, V> VacantEntry<'a, K, V> {
144    pub fn index(&self) -> IdxSize {
145        self.tuples.len() as IdxSize
146    }
147
148    pub fn insert(self, value: V) -> &'a mut V {
149        unsafe {
150            let tuple_idx: IdxSize = self.tuples.len().try_into().unwrap();
151            self.tuples.push((self.key, value));
152            self.entry.insert(tuple_idx);
153            &mut self.tuples.last_mut().unwrap_unchecked().1
154        }
155    }
156}
polars_utils/idx_map/total_idx_map.rs

polars_utils/idx_map/
total_idx_map.rs
