[ty] Deduplicate retained scope inference types

crates/ty_python_core/src/frozen.rs

@@ -1,3 +1,8 @@
+use std::hash::Hash;
+
+use ruff_index::Idx;
+use rustc_hash::FxHashMap;
+
 /// Compact immutable key-value entries stored in key order.
 ///
 /// Analysis builds these tables with hash maps, but after construction they only need keyed
@@ -100,6 +105,111 @@ impl<'a, K, V> IntoIterator for &'a mut FrozenMap<K, V> {
     }
 }
 
+#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, get_size2::GetSize, salsa::Update)]
+struct FrozenValueIndex(u32);
+
+impl Idx for FrozenValueIndex {
+    fn new(value: usize) -> Self {
+        assert!(u32::try_from(value).is_ok());
+        #[expect(clippy::cast_possible_truncation)]
+        Self(value as u32)
+    }
+
+    fn index(self) -> usize {
+        self.0 as usize
+    }
+}
+
+/// Compact immutable key-value entries that deduplicate repeated values.
+#[derive(Debug, Eq, PartialEq, salsa::Update, get_size2::GetSize)]
+pub struct FrozenValueMap<K, V> {
+    entries: FrozenMap<K, FrozenValueIndex>,
+    values: Box<[V]>,
+}
+
+impl<K, V> FrozenValueMap<K, V> {
+    pub fn get(&self, key: &K) -> Option<&V>
+    where
+        K: Ord,
+    {
+        self.entries
+            .get(key)
+            .map(|index| &self.values[index.index()])
+    }
+
+    pub fn iter(&self) -> impl DoubleEndedIterator<Item = (K, V)> + ExactSizeIterator + '_
+    where
+        K: Copy,
+        V: Copy,
+    {
+        self.entries
+            .iter()
+            .map(|(key, index)| (*key, self.values[index.index()]))
+    }
+
+    pub fn map_values<F>(&mut self, mut map: F)
+    where
+        K: Copy + Ord,
+        V: Copy + Eq + Hash,
+        F: FnMut(K, V) -> V,
+    {
+        *self = self
+            .iter()
+            .map(|(key, value)| (key, map(key, value)))
+            .collect();
+    }
+}
+
+impl<K, V> FromIterator<(K, V)> for FrozenValueMap<K, V>
+where
+    K: Ord,
+    V: Copy + Eq + Hash,
+{
+    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
+        let mut source_entries = iter.into_iter().collect::<Vec<_>>();
+        source_entries.sort_unstable_by(|(left, _), (right, _)| left.cmp(right));
+        source_entries.dedup_by(|(left, _), (right, _)| left == right);
+
+        let mut values = Vec::new();
+        let mut value_indices = FxHashMap::default();
+        let entries = source_entries
+            .into_iter()
+            .map(|(key, value)| {
+                let index = *value_indices.entry(value).or_insert_with(|| {
+                    let index = FrozenValueIndex::new(values.len());
+                    values.push(value);
+                    index
+                });
+                (key, index)
+            })
+            .collect::<Vec<_>>();
+
+        Self {
+            entries: FrozenMap(entries.into_boxed_slice()),
+            values: values.into_boxed_slice(),
+        }
+    }
+}
+
+impl<K, V, S> From<std::collections::HashMap<K, V, S>> for FrozenValueMap<K, V>
+where
+    K: Ord,
+    V: Copy + Eq + Hash,
+{
+    fn from(map: std::collections::HashMap<K, V, S>) -> Self {
+        map.into_iter().collect()
+    }
+}
+
+impl<K, V> Default for FrozenValueMap<K, V> {
+    fn default() -> Self {
+        Self {
+            entries: FrozenMap::default(),
+            values: Box::default(),
+        }
+    }
+}
+
 /// Compact immutable keys stored in ascending order.
 ///
 /// Analysis builds these sets with hash sets, but after construction they only need membership
@@ -141,3 +251,45 @@ impl<K> Default for FrozenSet<K> {
         Self(Box::default())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{FrozenMap, FrozenValueMap};
+
+    #[test]
+    fn frozen_value_map_deduplicates_values() {
+        let map = FrozenValueMap::from_iter([(3, [1; 4]), (1, [2; 4]), (2, [1; 4])]);
+
+        assert_eq!(map.values.len(), 2);
+        assert_eq!(map.get(&1), Some(&[2; 4]));
+        assert_eq!(map.get(&2), Some(&[1; 4]));
+        assert_eq!(
+            map.iter().collect::<Vec<_>>(),
+            vec![(1, [2; 4]), (2, [1; 4]), (3, [1; 4])]
+        );
+    }
+
+    #[test]
+    fn frozen_value_map_updates_and_rededuplicates_values() {
+        let mut map = FrozenValueMap::from_iter([(1, 10), (2, 20), (3, 30)]);
+
+        map.map_values(|_, _| 42);
+
+        assert_eq!(map.values.as_ref(), &[42]);
+        assert_eq!(
+            map.iter().collect::<Vec<_>>(),
+            vec![(1, 42), (2, 42), (3, 42)]
+        );
+    }
+
+    #[test]
+    fn frozen_value_map_uses_less_heap_for_repeated_large_values() {
+        let entries = [(1, [1; 8]), (2, [1; 8]), (3, [1; 8]), (4, [2; 8])];
+        let direct = FrozenMap::from_iter(entries);
+        let deduplicated = FrozenValueMap::from_iter(entries);
+
+        assert!(
+            ruff_memory_usage::heap_size(&deduplicated) < ruff_memory_usage::heap_size(&direct)
+        );
+    }
+}

crates/ty_python_semantic/src/types/infer.rs

@@ -51,7 +51,7 @@ use rustc_hash::FxHashMap;
 use salsa;
 use salsa::plumbing::AsId;
 use std::borrow::Cow;
-pub(super) use ty_python_core::frozen::{FrozenMap, FrozenSet};
+pub(super) use ty_python_core::frozen::{FrozenMap, FrozenSet, FrozenValueMap};
 
 use crate::types::diagnostic::TypeCheckDiagnostics;
 use crate::types::function::{FunctionDecorators, FunctionType};
@@ -750,7 +750,7 @@ impl<'db> InferenceRegion<'db> {
 #[derive(Debug, Eq, PartialEq, salsa::Update, get_size2::GetSize)]
 pub(crate) struct ScopeInference<'db> {
     /// The types of every expression in this region.
-    expressions: FrozenMap<ExpressionNodeKey, Type<'db>>,
+    expressions: FrozenValueMap<ExpressionNodeKey, Type<'db>>,
 
     /// The extra data that is only present for few inference regions.
     extra: Option<Box<ScopeInferenceExtra<'db>>>,
@@ -784,7 +784,7 @@ impl<'db> ScopeInference<'db> {
                 cycle_recovery: Some(cycle_recovery),
                 ..ScopeInferenceExtra::default()
             })),
-            expressions: FrozenMap::default(),
+            expressions: FrozenValueMap::default(),
         }
     }
 
@@ -794,10 +794,9 @@ impl<'db> ScopeInference<'db> {
         previous_inference: &ScopeInference<'db>,
         cycle: &salsa::Cycle,
     ) -> ScopeInference<'db> {
-        for (expr, ty) in &mut self.expressions {
-            let previous_ty = previous_inference.expression_type(*expr);
-            *ty = ty.cycle_normalized(db, previous_ty, cycle);
-        }
+        self.expressions.map_values(|expr, ty| {
+            ty.cycle_normalized(db, previous_inference.expression_type(expr), cycle)
+        });
 
         self
     }

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -23,9 +23,10 @@ use ty_python_core::statement::StatementInner;
 
 use super::{
     DefinitionInference, DefinitionInferenceExtra, DefinitionTypes, ExpressionInference,
-    ExpressionInferenceExtra, FrozenMap, FrozenSet, FunctionDecoratorInference, InferenceRegion,
-    ScopeInference, ScopeInferenceExtra, infer_deferred_types, infer_definition_types,
-    infer_expression_types, infer_same_file_expression_type, infer_unpack_types,
+    ExpressionInferenceExtra, FrozenMap, FrozenSet, FrozenValueMap, FunctionDecoratorInference,
+    InferenceRegion, ScopeInference, ScopeInferenceExtra, infer_deferred_types,
+    infer_definition_types, infer_expression_types, infer_same_file_expression_type,
+    infer_unpack_types,
 };
 use crate::diagnostic::format_enumeration;
 use crate::place::{
@@ -534,8 +535,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
     }
 
     fn extend_scope(&mut self, inference: &ScopeInference<'db>) {
-        self.expressions
-            .extend(inference.expressions.iter().copied());
+        self.expressions.extend(inference.expressions.iter());
 
         if let Some(extra) = &inference.extra {
             self.context.extend(&extra.diagnostics);
@@ -10648,7 +10648,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         });
 
         ScopeInference {
-            expressions: FrozenMap::from(expressions),
+            expressions: FrozenValueMap::from(expressions),
             extra,
         }
     }