ty_python_core/
frozen.rs

1use std::hash::Hash;
2
3use ruff_index::{FrozenIndexVec, IndexVec, newtype_index};
4use rustc_hash::FxHashMap;
5
6/// Compact immutable key-value entries stored in key order.
7///
8/// Analysis builds these tables with hash maps, but after construction they only need keyed
9/// lookup. A sorted slice avoids retaining hash-table capacity for every indexed file.
10#[derive(Debug, Eq, PartialEq, get_size2::GetSize, salsa::SalsaValue)]
11pub struct FrozenMap<K, V>(Box<[(K, V)]>);
12
13impl<K, V> FrozenMap<K, V> {
14    pub fn iter(&self) -> std::slice::Iter<'_, (K, V)> {
15        self.0.iter()
16    }
17
18    pub fn iter_mut(
19        &mut self,
20    ) -> impl DoubleEndedIterator<Item = (&K, &mut V)> + ExactSizeIterator + std::iter::FusedIterator
21    {
22        self.into_iter()
23    }
24
25    pub(crate) fn keys(&self) -> impl DoubleEndedIterator<Item = &K> + ExactSizeIterator {
26        self.0.iter().map(|(key, _)| key)
27    }
28
29    pub fn values(&self) -> impl DoubleEndedIterator<Item = &V> + ExactSizeIterator {
30        self.0.iter().map(|(_, value)| value)
31    }
32}
33
34impl<K: Ord, V> FromIterator<(K, V)> for FrozenMap<K, V> {
35    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
36        let entries = sort_and_deduplicate(iter.into_iter().collect());
37        Self(entries.into_boxed_slice())
38    }
39}
40
41impl<K, V> From<std::collections::BTreeMap<K, V>> for FrozenMap<K, V> {
42    fn from(map: std::collections::BTreeMap<K, V>) -> Self {
43        Self(map.into_iter().collect())
44    }
45}
46
47impl<K: Ord, V, S> From<std::collections::HashMap<K, V, S>> for FrozenMap<K, V> {
48    fn from(map: std::collections::HashMap<K, V, S>) -> Self {
49        Self::from_entries(map.into_iter().collect())
50    }
51}
52
53impl<K: Ord, V> FrozenMap<K, V> {
54    /// Creates a frozen map from entries with unique keys.
55    pub(crate) fn from_entries(mut entries: Vec<(K, V)>) -> Self {
56        entries.sort_unstable_by(|(left, _), (right, _)| left.cmp(right));
57        debug_assert!(
58            entries
59                .windows(2)
60                .all(|entries| entries[0].0 != entries[1].0),
61            "frozen map keys must be unique",
62        );
63        Self(entries.into_boxed_slice())
64    }
65
66    pub fn get(&self, key: &K) -> Option<&V> {
67        self.0
68            .binary_search_by(|(candidate, _)| candidate.cmp(key))
69            .ok()
70            .map(|index| &self.0[index].1)
71    }
72}
73
74impl<K, V> Default for FrozenMap<K, V> {
75    fn default() -> Self {
76        Self(Box::default())
77    }
78}
79
80impl<K: Ord, V> std::ops::Index<&K> for FrozenMap<K, V> {
81    type Output = V;
82
83    #[track_caller]
84    fn index(&self, index: &K) -> &Self::Output {
85        self.get(index).expect("key not found")
86    }
87}
88
89impl<K, V> IntoIterator for FrozenMap<K, V> {
90    type Item = (K, V);
91    type IntoIter = std::vec::IntoIter<(K, V)>;
92
93    fn into_iter(self) -> Self::IntoIter {
94        self.0.into_vec().into_iter()
95    }
96}
97
98impl<'a, K, V> IntoIterator for &'a FrozenMap<K, V> {
99    type Item = &'a (K, V);
100    type IntoIter = std::slice::Iter<'a, (K, V)>;
101
102    fn into_iter(self) -> Self::IntoIter {
103        self.0.iter()
104    }
105}
106
107impl<'a, K, V> IntoIterator for &'a mut FrozenMap<K, V> {
108    type Item = (&'a K, &'a mut V);
109    type IntoIter =
110        std::iter::Map<std::slice::IterMut<'a, (K, V)>, fn(&'a mut (K, V)) -> (&'a K, &'a mut V)>;
111
112    fn into_iter(self) -> Self::IntoIter {
113        self.0.iter_mut().map(|(key, value)| (&*key, value))
114    }
115}
116
117#[newtype_index]
118#[derive(get_size2::GetSize, salsa::SalsaValue)]
119struct FrozenValueIndex;
120
121/// Sorts entries by key and removes duplicate keys, retaining the last value for each key.
122///
123/// Stable sorting preserves the input order of equal-key entries, allowing the deduplication
124/// pass to provide last-entry-wins semantics like standard map collection.
125fn sort_and_deduplicate<K: Ord, V>(mut entries: Vec<(K, V)>) -> Vec<(K, V)> {
126    entries.sort_by(|(left, _), (right, _)| left.cmp(right));
127    entries.dedup_by(|(later_key, later_value), (earlier_key, earlier_value)| {
128        if later_key == earlier_key {
129            // `dedup_by` removes the later entry, so move its value to the retained entry.
130            std::mem::swap(later_value, earlier_value);
131            true
132        } else {
133            false
134        }
135    });
136
137    entries
138}
139
140fn index_values<K, V>(
141    entries: impl IntoIterator<Item = (K, V)>,
142) -> (Vec<(K, FrozenValueIndex)>, IndexVec<FrozenValueIndex, V>)
143where
144    V: Copy + Eq + Hash,
145{
146    let mut values = IndexVec::new();
147    let mut value_indices = FxHashMap::default();
148    let entries = entries
149        .into_iter()
150        .map(|(key, value)| {
151            let index = *value_indices
152                .entry(value)
153                .or_insert_with(|| values.push(value));
154            (key, index)
155        })
156        .collect();
157
158    (entries, values)
159}
160
161/// Compact immutable key-value entries that deduplicate repeated values.
162#[derive(Debug, Eq, PartialEq, get_size2::GetSize, salsa::SalsaValue)]
163pub struct FrozenValueMap<K, V> {
164    entries: FrozenMap<K, FrozenValueIndex>,
165    values: FrozenIndexVec<FrozenValueIndex, V>,
166}
167
168impl<K, V> FrozenValueMap<K, V> {
169    pub fn get(&self, key: &K) -> Option<&V>
170    where
171        K: Ord,
172    {
173        self.entries.get(key).map(|index| &self.values[*index])
174    }
175
176    pub fn iter(&self) -> impl DoubleEndedIterator<Item = (K, V)> + ExactSizeIterator + '_
177    where
178        K: Copy,
179        V: Copy,
180    {
181        self.entries
182            .iter()
183            .map(|(key, index)| (*key, self.values[*index]))
184    }
185
186    pub fn map_values<F>(&mut self, mut map: F)
187    where
188        K: Copy + Ord,
189        V: Copy + Eq + Hash,
190        F: FnMut(K, V) -> V,
191    {
192        let (entries, values) =
193            index_values(self.iter().map(|(key, value)| (key, map(key, value))));
194        *self = Self {
195            entries: FrozenMap(entries.into_boxed_slice()),
196            values: values.into(),
197        };
198    }
199}
200
201impl<K, V> FromIterator<(K, V)> for FrozenValueMap<K, V>
202where
203    K: Ord,
204    V: Copy + Eq + Hash,
205{
206    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
207        let source_entries = sort_and_deduplicate(iter.into_iter().collect());
208
209        let (entries, values) = index_values(source_entries);
210
211        Self {
212            entries: FrozenMap(entries.into_boxed_slice()),
213            values: values.into(),
214        }
215    }
216}
217
218impl<K, V, S> From<std::collections::HashMap<K, V, S>> for FrozenValueMap<K, V>
219where
220    K: Ord,
221    V: Copy + Eq + Hash,
222{
223    fn from(map: std::collections::HashMap<K, V, S>) -> Self {
224        let (mut entries, values) = index_values(map);
225        entries.sort_unstable_by(|(left, _), (right, _)| left.cmp(right));
226
227        Self {
228            entries: FrozenMap(entries.into_boxed_slice()),
229            values: values.into(),
230        }
231    }
232}
233
234impl<K, V> Default for FrozenValueMap<K, V> {
235    fn default() -> Self {
236        Self {
237            entries: FrozenMap::default(),
238            values: IndexVec::new().into(),
239        }
240    }
241}
242
243/// Compact immutable keys stored in ascending order.
244///
245/// Analysis builds these sets with hash sets, but after construction they only need membership
246/// tests and iteration. A sorted slice avoids retaining hash-table capacity.
247#[derive(Debug, Eq, PartialEq, get_size2::GetSize)]
248pub struct FrozenSet<K>(Box<[K]>);
249
250impl<K: Ord, S> From<std::collections::HashSet<K, S>> for FrozenSet<K> {
251    fn from(set: std::collections::HashSet<K, S>) -> Self {
252        let mut entries = set.into_iter().collect::<Vec<_>>();
253        entries.sort_unstable();
254        Self(entries.into_boxed_slice())
255    }
256}
257
258impl<K: Ord> FrozenSet<K> {
259    pub fn contains(&self, key: &K) -> bool {
260        self.0.binary_search(key).is_ok()
261    }
262}
263
264impl<K> FrozenSet<K> {
265    pub fn iter(&self) -> std::slice::Iter<'_, K> {
266        self.0.iter()
267    }
268}
269
270impl<'a, K> IntoIterator for &'a FrozenSet<K> {
271    type Item = &'a K;
272    type IntoIter = std::slice::Iter<'a, K>;
273
274    fn into_iter(self) -> Self::IntoIter {
275        self.0.iter()
276    }
277}
278
279impl<K> Default for FrozenSet<K> {
280    fn default() -> Self {
281        Self(Box::default())
282    }
283}
284
285#[cfg(test)]
286mod tests {
287    use super::{FrozenMap, FrozenValueMap};
288
289    #[test]
290    fn frozen_map_preserves_last_value_for_duplicate_keys() {
291        let map = FrozenMap::from_iter([(2, "two"), (1, "first"), (1, "last")]);
292
293        assert_eq!(
294            map.iter().copied().collect::<Vec<_>>(),
295            vec![(1, "last"), (2, "two")]
296        );
297    }
298
299    #[test]
300    fn frozen_map_iterates_with_mutable_values() {
301        let mut map = FrozenMap::from_iter([(2, 20), (1, 10)]);
302
303        for (key, value) in &mut map {
304            *value += *key;
305        }
306
307        assert_eq!(
308            map.iter().copied().collect::<Vec<_>>(),
309            vec![(1, 11), (2, 22)]
310        );
311    }
312
313    #[test]
314    fn frozen_value_map_deduplicates_values() {
315        let map = FrozenValueMap::from_iter([(3, [1; 4]), (1, [2; 4]), (2, [1; 4])]);
316
317        assert_eq!(map.values.len(), 2);
318        assert_eq!(map.get(&1), Some(&[2; 4]));
319        assert_eq!(map.get(&2), Some(&[1; 4]));
320        assert_eq!(
321            map.iter().collect::<Vec<_>>(),
322            vec![(1, [2; 4]), (2, [1; 4]), (3, [1; 4])]
323        );
324    }
325
326    #[test]
327    fn frozen_value_map_preserves_last_value_for_duplicate_keys() {
328        let map = FrozenValueMap::from_iter([(2, 20), (1, 10), (1, 11)]);
329
330        assert_eq!(map.iter().collect::<Vec<_>>(), vec![(1, 11), (2, 20)]);
331    }
332
333    #[test]
334    fn frozen_value_map_updates_and_rededuplicates_values() {
335        let mut map = FrozenValueMap::from_iter([(1, 10), (2, 20), (3, 30)]);
336
337        map.map_values(|_, _| 42);
338
339        assert_eq!(&map.values.raw, &[42]);
340        assert_eq!(
341            map.iter().collect::<Vec<_>>(),
342            vec![(1, 42), (2, 42), (3, 42)]
343        );
344    }
345
346    #[test]
347    fn frozen_value_map_uses_less_heap_for_repeated_large_values() {
348        let entries = [(1, [1; 8]), (2, [1; 8]), (3, [1; 8]), (4, [2; 8])];
349        let direct = FrozenMap::from_iter(entries);
350        let deduplicated = FrozenValueMap::from_iter(entries);
351
352        assert!(
353            ruff_memory_usage::heap_size(&deduplicated) < ruff_memory_usage::heap_size(&direct)
354        );
355    }
356}
ty_python_core/frozen.rs

ty_python_core/
frozen.rs
