[red-knot] Separate definitions_by_definition into separate fields (#…

…16277)

A minor cleanup that breaks up a `HashMap` of an enum into separate
`HashMap`s for each variant. (These separate fields were already how
this cache was being described in the big comment at the top of the
file!)

crates/red_knot_python_semantic/src/semantic_index/use_def.rs

@@ -293,19 +293,22 @@ pub(crate) struct UseDefMap<'db> {
     /// [`SymbolBindings`] reaching a [`ScopedUseId`].
     bindings_by_use: IndexVec<ScopedUseId, SymbolBindings>,
 
-    /// [`SymbolBindings`] or [`SymbolDeclarations`] reaching a given [`Definition`].
-    ///
     /// If the definition is a binding (only) -- `x = 1` for example -- then we need
     /// [`SymbolDeclarations`] to know whether this binding is permitted by the live declarations.
     ///
+    /// If the definition is both a declaration and a binding -- `x: int = 1` for example -- then
+    /// we don't actually need anything here, all we'll need to validate is that our own RHS is a
+    /// valid assignment to our own annotation.
+    declarations_by_binding: FxHashMap<Definition<'db>, SymbolDeclarations>,
+
     /// If the definition is a declaration (only) -- `x: int` for example -- then we need
     /// [`SymbolBindings`] to know whether this declaration is consistent with the previously
     /// inferred type.
     ///
     /// If the definition is both a declaration and a binding -- `x: int = 1` for example -- then
     /// we don't actually need anything here, all we'll need to validate is that our own RHS is a
     /// valid assignment to our own annotation.
-    definitions_by_definition: FxHashMap<Definition<'db>, SymbolDefinitions>,
+    bindings_by_declaration: FxHashMap<Definition<'db>, SymbolBindings>,
 
     /// [`SymbolState`] visible at end of scope for each symbol.
     public_symbols: IndexVec<ScopedSymbolId, SymbolState>,
@@ -343,25 +346,14 @@ impl<'db> UseDefMap<'db> {
         &self,
         declaration: Definition<'db>,
     ) -> BindingWithConstraintsIterator<'_, 'db> {
-        if let SymbolDefinitions::Bindings(bindings) = &self.definitions_by_definition[&declaration]
-        {
-            self.bindings_iterator(bindings)
-        } else {
-            unreachable!("Declaration has non-Bindings in definitions_by_definition");
-        }
+        self.bindings_iterator(&self.bindings_by_declaration[&declaration])
     }
 
-    pub(crate) fn declarations_at_binding<'map>(
-        &'map self,
+    pub(crate) fn declarations_at_binding(
+        &self,
         binding: Definition<'db>,
-    ) -> DeclarationsIterator<'map, 'db> {
-        if let SymbolDefinitions::Declarations(declarations) =
-            &self.definitions_by_definition[&binding]
-        {
-            self.declarations_iterator(declarations)
-        } else {
-            unreachable!("Binding has non-Declarations in definitions_by_definition");
-        }
+    ) -> DeclarationsIterator<'_, 'db> {
+        self.declarations_iterator(&self.declarations_by_binding[&binding])
     }
 
     pub(crate) fn public_declarations<'map>(
@@ -420,13 +412,6 @@ pub(crate) struct EagerBindingsKey {
 /// A snapshot of bindings that can be used to resolve a reference in a nested eager scope.
 type EagerBindings = IndexVec<ScopedEagerBindingsId, SymbolBindings>;
 
-/// Either live bindings or live declarations for a symbol.
-#[derive(Debug, PartialEq, Eq, salsa::Update)]
-enum SymbolDefinitions {
-    Bindings(SymbolBindings),
-    Declarations(SymbolDeclarations),
-}
-
 #[derive(Debug)]
 pub(crate) struct BindingWithConstraintsIterator<'map, 'db> {
     all_definitions: &'map IndexVec<ScopedDefinitionId, Option<Definition<'db>>>,
@@ -537,8 +522,11 @@ pub(super) struct UseDefMapBuilder<'db> {
     /// Live bindings at each so-far-recorded use.
     bindings_by_use: IndexVec<ScopedUseId, SymbolBindings>,
 
-    /// Live bindings or declarations for each so-far-recorded definition.
-    definitions_by_definition: FxHashMap<Definition<'db>, SymbolDefinitions>,
+    /// Live declarations for each so-far-recorded binding.
+    declarations_by_binding: FxHashMap<Definition<'db>, SymbolDeclarations>,
+
+    /// Live bindings for each so-far-recorded declaration.
+    bindings_by_declaration: FxHashMap<Definition<'db>, SymbolBindings>,
 
     /// Currently live bindings and declarations for each symbol.
     symbol_states: IndexVec<ScopedSymbolId, SymbolState>,
@@ -556,7 +544,8 @@ impl Default for UseDefMapBuilder<'_> {
             visibility_constraints: VisibilityConstraintsBuilder::default(),
             scope_start_visibility: ScopedVisibilityConstraintId::ALWAYS_TRUE,
             bindings_by_use: IndexVec::new(),
-            definitions_by_definition: FxHashMap::default(),
+            declarations_by_binding: FxHashMap::default(),
+            bindings_by_declaration: FxHashMap::default(),
             symbol_states: IndexVec::new(),
             eager_bindings: EagerBindings::default(),
         }
@@ -578,10 +567,8 @@ impl<'db> UseDefMapBuilder<'db> {
     pub(super) fn record_binding(&mut self, symbol: ScopedSymbolId, binding: Definition<'db>) {
         let def_id = self.all_definitions.push(Some(binding));
         let symbol_state = &mut self.symbol_states[symbol];
-        self.definitions_by_definition.insert(
-            binding,
-            SymbolDefinitions::Declarations(symbol_state.declarations().clone()),
-        );
+        self.declarations_by_binding
+            .insert(binding, symbol_state.declarations().clone());
         symbol_state.record_binding(def_id, self.scope_start_visibility);
     }
 
@@ -658,10 +645,8 @@ impl<'db> UseDefMapBuilder<'db> {
     ) {
         let def_id = self.all_definitions.push(Some(declaration));
         let symbol_state = &mut self.symbol_states[symbol];
-        self.definitions_by_definition.insert(
-            declaration,
-            SymbolDefinitions::Bindings(symbol_state.bindings().clone()),
-        );
+        self.bindings_by_declaration
+            .insert(declaration, symbol_state.bindings().clone());
         symbol_state.record_declaration(def_id);
     }
 
@@ -670,7 +655,8 @@ impl<'db> UseDefMapBuilder<'db> {
         symbol: ScopedSymbolId,
         definition: Definition<'db>,
     ) {
-        // We don't need to store anything in self.definitions_by_definition.
+        // We don't need to store anything in self.bindings_by_declaration or
+        // self.declarations_by_binding.
         let def_id = self.all_definitions.push(Some(definition));
         let symbol_state = &mut self.symbol_states[symbol];
         symbol_state.record_declaration(def_id);
@@ -770,7 +756,8 @@ impl<'db> UseDefMapBuilder<'db> {
         self.all_constraints.shrink_to_fit();
         self.symbol_states.shrink_to_fit();
         self.bindings_by_use.shrink_to_fit();
-        self.definitions_by_definition.shrink_to_fit();
+        self.declarations_by_binding.shrink_to_fit();
+        self.bindings_by_declaration.shrink_to_fit();
         self.eager_bindings.shrink_to_fit();
 
         UseDefMap {
@@ -779,7 +766,8 @@ impl<'db> UseDefMapBuilder<'db> {
             visibility_constraints: self.visibility_constraints.build(),
             bindings_by_use: self.bindings_by_use,
             public_symbols: self.symbol_states,
-            definitions_by_definition: self.definitions_by_definition,
+            declarations_by_binding: self.declarations_by_binding,
+            bindings_by_declaration: self.bindings_by_declaration,
             eager_bindings: self.eager_bindings,
         }
     }

crates/red_knot_python_semantic/src/semantic_index/use_def/symbol_state.rs

@@ -115,7 +115,7 @@ type VisibilityConstraintsIterator<'a> = std::slice::Iter<'a, ScopedVisibilityCo
 
 /// Live declarations for a single symbol at some point in control flow, with their
 /// corresponding visibility constraints.
-#[derive(Clone, Debug, Default, PartialEq, Eq)]
+#[derive(Clone, Debug, Default, PartialEq, Eq, salsa::Update)]
 pub(super) struct SymbolDeclarations {
     /// [`BitSet`]: which declarations (as [`ScopedDefinitionId`]) can reach the current location?
     ///
@@ -203,7 +203,7 @@ impl SymbolDeclarations {
 
 /// Live bindings for a single symbol at some point in control flow. Each live binding comes
 /// with a set of narrowing constraints and a visibility constraint.
-#[derive(Clone, Debug, Default, PartialEq, Eq)]
+#[derive(Clone, Debug, Default, PartialEq, Eq, salsa::Update)]
 pub(super) struct SymbolBindings {
     /// [`BitSet`]: which bindings (as [`ScopedDefinitionId`]) can reach the current location?
     ///