[red-knot] Simplify visibility constraint handling for *-import definitions (#17486)

Author: AlexWaygood

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -355,15 +355,14 @@ impl<'db> SemanticIndexBuilder<'db> {
         self.current_use_def_map_mut().merge(state);
     }
 
-    /// Return a 2-element tuple, where the first element is the [`ScopedSymbolId`] of the
-    /// symbol added, and the second element is a boolean indicating whether the symbol was *newly*
-    /// added or not
-    fn add_symbol(&mut self, name: Name) -> (ScopedSymbolId, bool) {
+    /// Add a symbol to the symbol table and the use-def map.
+    /// Return the [`ScopedSymbolId`] that uniquely identifies the symbol in both.
+    fn add_symbol(&mut self, name: Name) -> ScopedSymbolId {
         let (symbol_id, added) = self.current_symbol_table().add_symbol(name);
         if added {
             self.current_use_def_map_mut().add_symbol(symbol_id);
         }
-        (symbol_id, added)
+        symbol_id
     }
 
     fn add_attribute(&mut self, name: Name) -> ScopedSymbolId {
@@ -797,7 +796,7 @@ impl<'db> SemanticIndexBuilder<'db> {
                         ..
                     }) => (name, &None, default),
                 };
-                let (symbol, _) = self.add_symbol(name.id.clone());
+                let symbol = self.add_symbol(name.id.clone());
                 // TODO create Definition for PEP 695 typevars
                 // note that the "bound" on the typevar is a totally different thing than whether
                 // or not a name is "bound" by a typevar declaration; the latter is always true.
@@ -895,20 +894,20 @@ impl<'db> SemanticIndexBuilder<'db> {
             self.declare_parameter(parameter);
         }
         if let Some(vararg) = parameters.vararg.as_ref() {
-            let (symbol, _) = self.add_symbol(vararg.name.id().clone());
+            let symbol = self.add_symbol(vararg.name.id().clone());
             self.add_definition(
                 symbol,
                 DefinitionNodeRef::VariadicPositionalParameter(vararg),
             );
         }
         if let Some(kwarg) = parameters.kwarg.as_ref() {
-            let (symbol, _) = self.add_symbol(kwarg.name.id().clone());
+            let symbol = self.add_symbol(kwarg.name.id().clone());
             self.add_definition(symbol, DefinitionNodeRef::VariadicKeywordParameter(kwarg));
         }
     }
 
     fn declare_parameter(&mut self, parameter: &'db ast::ParameterWithDefault) {
-        let (symbol, _) = self.add_symbol(parameter.name().id().clone());
+        let symbol = self.add_symbol(parameter.name().id().clone());
 
         let definition = self.add_definition(symbol, parameter);
 
@@ -1139,7 +1138,7 @@ where
                 // The symbol for the function name itself has to be evaluated
                 // at the end to match the runtime evaluation of parameter defaults
                 // and return-type annotations.
-                let (symbol, _) = self.add_symbol(name.id.clone());
+                let symbol = self.add_symbol(name.id.clone());
 
                 // Record a use of the function name in the scope that it is defined in, so that it
                 // can be used to find previously defined functions with the same name. This is
@@ -1174,11 +1173,11 @@ where
                 );
 
                 // In Python runtime semantics, a class is registered after its scope is evaluated.
-                let (symbol, _) = self.add_symbol(class.name.id.clone());
+                let symbol = self.add_symbol(class.name.id.clone());
                 self.add_definition(symbol, class);
             }
             ast::Stmt::TypeAlias(type_alias) => {
-                let (symbol, _) = self.add_symbol(
+                let symbol = self.add_symbol(
                     type_alias
                         .name
                         .as_name_expr()
@@ -1215,7 +1214,7 @@ where
                         (Name::new(alias.name.id.split('.').next().unwrap()), false)
                     };
 
-                    let (symbol, _) = self.add_symbol(symbol_name);
+                    let symbol = self.add_symbol(symbol_name);
                     self.add_definition(
                         symbol,
                         ImportDefinitionNodeRef {
@@ -1286,7 +1285,7 @@ where
                         //
                         // For more details, see the doc-comment on `StarImportPlaceholderPredicate`.
                         for export in exported_names(self.db, referenced_module) {
-                            let (symbol_id, newly_added) = self.add_symbol(export.clone());
+                            let symbol_id = self.add_symbol(export.clone());
                             let node_ref = StarImportDefinitionNodeRef { node, symbol_id };
                             let star_import = StarImportPlaceholderPredicate::new(
                                 self.db,
@@ -1295,28 +1294,15 @@ where
                                 referenced_module,
                             );
 
-                            // Fast path for if there were no previous definitions
-                            // of the symbol defined through the `*` import:
-                            // we can apply the visibility constraint to *only* the added definition,
-                            // rather than all definitions
-                            if newly_added {
-                                self.push_additional_definition(symbol_id, node_ref);
-                                self.current_use_def_map_mut()
-                                    .record_and_negate_star_import_visibility_constraint(
-                                        star_import,
-                                        symbol_id,
-                                    );
-                            } else {
-                                let pre_definition = self.flow_snapshot();
-                                self.push_additional_definition(symbol_id, node_ref);
-                                let constraint_id =
-                                    self.record_visibility_constraint(star_import.into());
-                                let post_definition = self.flow_snapshot();
-                                self.flow_restore(pre_definition.clone());
-                                self.record_negated_visibility_constraint(constraint_id);
-                                self.flow_merge(post_definition);
-                                self.simplify_visibility_constraints(pre_definition);
-                            }
+                            let pre_definition =
+                                self.current_use_def_map().single_symbol_snapshot(symbol_id);
+                            self.push_additional_definition(symbol_id, node_ref);
+                            self.current_use_def_map_mut()
+                                .record_and_negate_star_import_visibility_constraint(
+                                    star_import,
+                                    symbol_id,
+                                    pre_definition,
+                                );
                         }
 
                         continue;
@@ -1336,7 +1322,7 @@ where
                     self.has_future_annotations |= alias.name.id == "annotations"
                         && node.module.as_deref() == Some("__future__");
 
-                    let (symbol, _) = self.add_symbol(symbol_name.clone());
+                    let symbol = self.add_symbol(symbol_name.clone());
 
                     self.add_definition(
                         symbol,
@@ -1754,7 +1740,7 @@ where
                         // which is invalid syntax. However, it's still pretty obvious here that the user
                         // *wanted* `e` to be bound, so we should still create a definition here nonetheless.
                         if let Some(symbol_name) = symbol_name {
-                            let (symbol, _) = self.add_symbol(symbol_name.id.clone());
+                            let symbol = self.add_symbol(symbol_name.id.clone());
 
                             self.add_definition(
                                 symbol,
@@ -1841,7 +1827,7 @@ where
                     (ast::ExprContext::Del, _) => (false, true),
                     (ast::ExprContext::Invalid, _) => (false, false),
                 };
-                let (symbol, _) = self.add_symbol(id.clone());
+                let symbol = self.add_symbol(id.clone());
 
                 if is_use {
                     self.mark_symbol_used(symbol);
@@ -2245,7 +2231,7 @@ where
             range: _,
         }) = pattern
         {
-            let (symbol, _) = self.add_symbol(name.id().clone());
+            let symbol = self.add_symbol(name.id().clone());
             let state = self.current_match_case.as_ref().unwrap();
             self.add_definition(
                 symbol,
@@ -2266,7 +2252,7 @@ where
             rest: Some(name), ..
         }) = pattern
         {
-            let (symbol, _) = self.add_symbol(name.id().clone());
+            let symbol = self.add_symbol(name.id().clone());
             let state = self.current_match_case.as_ref().unwrap();
             self.add_definition(
                 symbol,

crates/red_knot_python_semantic/src/semantic_index/use_def.rs

@@ -775,7 +775,16 @@ impl<'db> UseDefMapBuilder<'db> {
             .add_and_constraint(self.scope_start_visibility, constraint);
     }
 
-    /// This method exists solely as a fast path for handling `*`-import visibility constraints.
+    /// Snapshot the state of a single symbol at the current point in control flow.
+    ///
+    /// This is only used for `*`-import visibility constraints, which are handled differently
+    /// to most other visibility constraints. See the doc-comment for
+    /// [`Self::record_and_negate_star_import_visibility_constraint`] for more details.
+    pub(super) fn single_symbol_snapshot(&self, symbol: ScopedSymbolId) -> SymbolState {
+        self.symbol_states[symbol].clone()
+    }
+
+    /// This method exists solely for handling `*`-import visibility constraints.
     ///
     /// The reason why we add visibility constraints for [`Definition`]s created by `*` imports
     /// is laid out in the doc-comment for [`StarImportPlaceholderPredicate`]. But treating these
@@ -784,12 +793,11 @@ impl<'db> UseDefMapBuilder<'db> {
     /// dominates. (Although `*` imports are not common generally, they are used in several
     /// important places by typeshed.)
     ///
-    /// To solve these regressions, it was observed that we could add a fast path for `*`-import
-    /// definitions which added a new symbol to the global scope (as opposed to `*`-import definitions
-    /// that provided redefinitions for *pre-existing* global-scope symbols). The fast path does a
-    /// number of things differently to our normal handling of visibility constraints:
+    /// To solve these regressions, it was observed that we could do significantly less work for
+    /// `*`-import definitions. We do a number of things differently here to our normal handling of
+    /// visibility constraints:
     ///
-    /// - It only applies and negates the visibility constraints to a single symbol, rather than to
+    /// - We only apply and negate the visibility constraints to a single symbol, rather than to
     ///   all symbols. This is possible here because, unlike most definitions, we know in advance that
     ///   exactly one definition occurs inside the "if-true" predicate branch, and we know exactly
     ///   which definition it is.
@@ -800,9 +808,9 @@ impl<'db> UseDefMapBuilder<'db> {
     ///   the visibility constraints is only important for symbols that did not have any new
     ///   definitions inside either the "if-predicate-true" branch or the "if-predicate-false" branch.
     ///
-    /// - It avoids multiple expensive calls to [`Self::snapshot`]. This is possible because we know
-    ///   the symbol is newly added, so we know the prior state of the symbol was
-    ///   [`SymbolState::undefined`].
+    /// - We only snapshot the state for a single symbol prior to the definition, rather than doing
+    ///   expensive calls to [`Self::snapshot`]. Again, this is possible because we know
+    ///   that only a single definition occurs inside the "if-predicate-true" predicate branch.
     ///
     /// - Normally we take care to check whether an "if-predicate-true" branch or an
     ///   "if-predicate-false" branch contains a terminal statement: these can affect the visibility
@@ -815,17 +823,17 @@ impl<'db> UseDefMapBuilder<'db> {
         &mut self,
         star_import: StarImportPlaceholderPredicate<'db>,
         symbol: ScopedSymbolId,
+        pre_definition_state: SymbolState,
     ) {
         let predicate_id = self.add_predicate(star_import.into());
         let visibility_id = self.visibility_constraints.add_atom(predicate_id);
         let negated_visibility_id = self
             .visibility_constraints
             .add_not_constraint(visibility_id);
 
-        let mut post_definition_state = std::mem::replace(
-            &mut self.symbol_states[symbol],
-            SymbolState::undefined(self.scope_start_visibility),
-        );
+        let mut post_definition_state =
+            std::mem::replace(&mut self.symbol_states[symbol], pre_definition_state);
+
         post_definition_state
             .record_visibility_constraint(&mut self.visibility_constraints, visibility_id);
 

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

@@ -314,7 +314,7 @@ impl SymbolBindings {
 }
 
 #[derive(Clone, Debug, PartialEq, Eq)]
-pub(super) struct SymbolState {
+pub(in crate::semantic_index) struct SymbolState {
     declarations: SymbolDeclarations,
     bindings: SymbolBindings,
 }