separate lambda parameters from function parameter definitions

Author: ibraheemdev

crates/ty_python_semantic/src/semantic_index.rs

@@ -865,7 +865,9 @@ mod tests {
     use crate::Db;
     use crate::db::tests::{TestDb, TestDbBuilder};
     use crate::semantic_index::ast_ids::{HasScopedUseId, ScopedUseId};
-    use crate::semantic_index::definition::{Definition, DefinitionKind};
+    use crate::semantic_index::definition::{
+        Definition, DefinitionKind, LambdaParameterDefinitionNodeKind, ParameterDefinitionNodeKind,
+    };
     use crate::semantic_index::place::PlaceTable;
     use crate::semantic_index::scope::{FileScopeId, Scope, ScopeKind};
     use crate::semantic_index::symbol::ScopedSymbolId;
@@ -1156,14 +1158,14 @@ def f(a: str, /, b: str, c: int = 1, *args, d: int = 2, **kwargs):
             .unwrap();
         assert!(matches!(
             args_binding.kind(&db),
-            DefinitionKind::VariadicPositionalParameter(_)
+            DefinitionKind::Parameter(ParameterDefinitionNodeKind::VariadicPositionalParameter(_))
         ));
         let kwargs_binding = use_def
             .first_public_binding(function_table.symbol_id("kwargs").expect("symbol exists"))
             .unwrap();
         assert!(matches!(
             kwargs_binding.kind(&db),
-            DefinitionKind::VariadicKeywordParameter(_)
+            DefinitionKind::Parameter(ParameterDefinitionNodeKind::VariadicKeywordParameter(_))
         ));
     }
 
@@ -1186,29 +1188,44 @@ def f(a: str, /, b: str, c: int = 1, *args, d: int = 2, **kwargs):
         let lambda_table = index.place_table(lambda_scope_id);
         assert_eq!(
             names(lambda_table),
-            vec!["a", "b", "c", "d", "args", "kwargs"],
+            vec!["a", "b", "c", "args", "d", "kwargs"],
         );
 
         let use_def = index.use_def_map(lambda_scope_id);
         for name in ["a", "b", "c", "d"] {
             let binding = use_def
                 .first_public_binding(lambda_table.symbol_id(name).expect("symbol exists"))
                 .unwrap();
-            assert!(matches!(binding.kind(&db), DefinitionKind::Parameter(_)));
+            assert!(matches!(
+                binding.kind(&db),
+                DefinitionKind::LambdaParameter(LambdaParameterDefinitionNodeKind {
+                    index: _,
+                    lambda: _,
+                    parameter: ParameterDefinitionNodeKind::Parameter(_)
+                })
+            ));
         }
         let args_binding = use_def
             .first_public_binding(lambda_table.symbol_id("args").expect("symbol exists"))
             .unwrap();
         assert!(matches!(
             args_binding.kind(&db),
-            DefinitionKind::VariadicPositionalParameter(_)
+            DefinitionKind::LambdaParameter(LambdaParameterDefinitionNodeKind {
+                index: 3,
+                lambda: _,
+                parameter: ParameterDefinitionNodeKind::VariadicPositionalParameter(_)
+            })
         ));
         let kwargs_binding = use_def
             .first_public_binding(lambda_table.symbol_id("kwargs").expect("symbol exists"))
             .unwrap();
         assert!(matches!(
             kwargs_binding.kind(&db),
-            DefinitionKind::VariadicKeywordParameter(_)
+            DefinitionKind::LambdaParameter(LambdaParameterDefinitionNodeKind {
+                index: 5,
+                lambda: _,
+                parameter: ParameterDefinitionNodeKind::VariadicKeywordParameter(_)
+            })
         ));
     }
 

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -27,8 +27,9 @@ use crate::semantic_index::definition::{
     ComprehensionDefinitionNodeRef, Definition, DefinitionCategory, DefinitionNodeKey,
     DefinitionNodeRef, Definitions, DictKeyAssignmentNodeRef, ExceptHandlerDefinitionNodeRef,
     ForStmtDefinitionNodeRef, ImportDefinitionNodeRef, ImportFromDefinitionNodeRef,
-    ImportFromSubmoduleDefinitionNodeRef, LoopHeaderDefinitionNodeRef, LoopStmtRef,
-    MatchPatternDefinitionNodeRef, StarImportDefinitionNodeRef, WithItemDefinitionNodeRef,
+    ImportFromSubmoduleDefinitionNodeRef, LambdaParameterDefinitionNodeRef,
+    LoopHeaderDefinitionNodeRef, LoopStmtRef, MatchPatternDefinitionNodeRef,
+    ParameterDefinitionNodeRef, StarImportDefinitionNodeRef, WithItemDefinitionNodeRef,
 };
 use crate::semantic_index::expression::{Expression, ExpressionKind};
 use crate::semantic_index::member::MemberExprBuilder;
@@ -1635,7 +1636,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
                 .mark_parameter();
             self.add_definition(
                 symbol.into(),
-                DefinitionNodeRef::VariadicPositionalParameter(vararg),
+                ParameterDefinitionNodeRef::VariadicPositionalParameter(vararg),
             );
         }
         if let Some(kwarg) = parameters.kwarg.as_ref() {
@@ -1645,15 +1646,90 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
                 .mark_parameter();
             self.add_definition(
                 symbol.into(),
-                DefinitionNodeRef::VariadicKeywordParameter(kwarg),
+                ParameterDefinitionNodeRef::VariadicKeywordParameter(kwarg),
             );
         }
     }
 
     fn declare_parameter(&mut self, parameter: &'ast ast::ParameterWithDefault) {
         let symbol = self.add_symbol(parameter.name().id().clone());
 
-        let definition = self.add_definition(symbol.into(), parameter);
+        let definition = self.add_definition(
+            symbol.into(),
+            ParameterDefinitionNodeRef::Parameter(parameter),
+        );
+
+        self.current_place_table_mut()
+            .symbol_mut(symbol)
+            .mark_parameter();
+
+        // Insert a mapping from the inner Parameter node to the same definition. This
+        // ensures that calling `HasType::inferred_type` on the inner parameter returns
+        // a valid type (and doesn't panic)
+        let existing_definition = self.definitions_by_node.insert(
+            (&parameter.parameter).into(),
+            Definitions::single(definition),
+        );
+        debug_assert_eq!(existing_definition, None);
+    }
+
+    fn declare_lambda_parameters(&mut self, parameters: &'ast ast::Parameters) {
+        let mut index = 0;
+        for parameter in &parameters.posonlyargs {
+            self.declare_lambda_parameter(index, parameter);
+            index += 1;
+        }
+        for parameter in &parameters.args {
+            self.declare_lambda_parameter(index, parameter);
+            index += 1;
+        }
+        if let Some(vararg) = parameters.vararg.as_ref() {
+            let symbol = self.add_symbol(vararg.name.id().clone());
+            self.current_place_table_mut()
+                .symbol_mut(symbol)
+                .mark_parameter();
+            self.add_definition(
+                symbol.into(),
+                LambdaParameterDefinitionNodeRef {
+                    index,
+                    parameter: ParameterDefinitionNodeRef::VariadicPositionalParameter(vararg),
+                },
+            );
+            index += 1;
+        }
+        for parameter in &parameters.kwonlyargs {
+            self.declare_lambda_parameter(index, parameter);
+            index += 1;
+        }
+        if let Some(kwarg) = parameters.kwarg.as_ref() {
+            let symbol = self.add_symbol(kwarg.name.id().clone());
+            self.current_place_table_mut()
+                .symbol_mut(symbol)
+                .mark_parameter();
+            self.add_definition(
+                symbol.into(),
+                LambdaParameterDefinitionNodeRef {
+                    index,
+                    parameter: ParameterDefinitionNodeRef::VariadicKeywordParameter(kwarg),
+                },
+            );
+        }
+    }
+
+    fn declare_lambda_parameter(
+        &mut self,
+        index: usize,
+        parameter: &'ast ast::ParameterWithDefault,
+    ) {
+        let symbol = self.add_symbol(parameter.name().id().clone());
+
+        let definition = self.add_definition(
+            symbol.into(),
+            LambdaParameterDefinitionNodeRef {
+                index,
+                parameter: ParameterDefinitionNodeRef::Parameter(parameter),
+            },
+        );
 
         self.current_place_table_mut()
             .symbol_mut(symbol)
@@ -3168,7 +3244,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
 
                 // Add symbols and definitions for the parameters to the lambda scope.
                 if let Some(parameters) = lambda.parameters.as_ref() {
-                    self.declare_parameters(parameters);
+                    self.declare_lambda_parameters(parameters);
                 }
 
                 self.visit_expr(lambda.body.as_ref());