Use &'db Name in type inference

astral-sh · Jun 29, 2024 · 954f54b · 954f54b
1 parent df99bef
commit 954f54b
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 29 deletions.
diff --git a/crates/red_knot_python_semantic/src/semantic_index.rs b/crates/red_knot_python_semantic/src/semantic_index.rs
@@ -77,7 +77,7 @@ pub struct SemanticIndex<'db> {
     symbol_tables: IndexVec<FileScopeId, Arc<SymbolTable<'db>>>,
 
     /// List of all scopes in this file.
-    scopes: IndexVec<FileScopeId, Scope>,
+    scopes: IndexVec<FileScopeId, Scope<'db>>,
 
     /// Maps expressions to their corresponding scope.
     /// We can't use [`ExpressionId`] here, because the challenge is how to get from
@@ -159,7 +159,7 @@ impl<'db> SemanticIndex<'db> {
 /// ID that uniquely identifies an expression inside a [`Scope`].
 
 pub struct AncestorsIter<'a> {
-    scopes: &'a IndexSlice<FileScopeId, Scope>,
+    scopes: &'a IndexSlice<FileScopeId, Scope<'a>>,
     next_id: Option<FileScopeId>,
 }
 
@@ -173,7 +173,7 @@ impl<'a> AncestorsIter<'a> {
 }
 
 impl<'a> Iterator for AncestorsIter<'a> {
-    type Item = (FileScopeId, &'a Scope);
+    type Item = (FileScopeId, &'a Scope<'a>);
 
     fn next(&mut self) -> Option<Self::Item> {
         let current_id = self.next_id?;
@@ -188,7 +188,7 @@ impl FusedIterator for AncestorsIter<'_> {}
 
 pub struct DescendentsIter<'a> {
     next_id: FileScopeId,
-    descendents: std::slice::Iter<'a, Scope>,
+    descendents: std::slice::Iter<'a, Scope<'a>>,
 }
 
 impl<'a> DescendentsIter<'a> {
@@ -204,7 +204,7 @@ impl<'a> DescendentsIter<'a> {
 }
 
 impl<'a> Iterator for DescendentsIter<'a> {
-    type Item = (FileScopeId, &'a Scope);
+    type Item = (FileScopeId, &'a Scope<'a>);
 
     fn next(&mut self) -> Option<Self::Item> {
         let descendent = self.descendents.next()?;
@@ -240,7 +240,7 @@ impl<'a> ChildrenIter<'a> {
 }
 
 impl<'a> Iterator for ChildrenIter<'a> {
-    type Item = (FileScopeId, &'a Scope);
+    type Item = (FileScopeId, &'a Scope<'a>);
 
     fn next(&mut self) -> Option<Self::Item> {
         self.descendents

diff --git a/crates/red_knot_python_semantic/src/semantic_index/builder.rs b/crates/red_knot_python_semantic/src/semantic_index/builder.rs
@@ -28,7 +28,7 @@ pub(super) struct SemanticIndexBuilder<'a> {
     current_definition: Option<Definition>,
 
     // Semantic Index fields
-    scopes: IndexVec<FileScopeId, Scope>,
+    scopes: IndexVec<FileScopeId, Scope<'a>>,
     symbol_tables: IndexVec<FileScopeId, SymbolTableBuilder<'a>>,
     ast_ids: IndexVec<FileScopeId, AstIdsBuilder>,
     expression_scopes: FxHashMap<NodeKey, FileScopeId>,
@@ -51,7 +51,7 @@ impl<'a> SemanticIndexBuilder<'a> {
 
         builder.push_scope_with_parent(
             NodeWithScopeId::Module,
-            &Name::new_static("<module>"),
+            Cow::Owned(Name::new_static("<module>")),
             None,
             None,
             None,
@@ -70,26 +70,32 @@ impl<'a> SemanticIndexBuilder<'a> {
     fn push_scope(
         &mut self,
         node: NodeWithScopeId,
-        name: &Name,
+        name: &'a Name,
         defining_symbol: Option<FileSymbolId>,
         definition: Option<Definition>,
     ) {
         let parent = self.current_scope();
-        self.push_scope_with_parent(node, name, defining_symbol, definition, Some(parent));
+        self.push_scope_with_parent(
+            node,
+            Cow::Borrowed(name),
+            defining_symbol,
+            definition,
+            Some(parent),
+        );
     }
 
     fn push_scope_with_parent(
         &mut self,
         node: NodeWithScopeId,
-        name: &Name,
+        name: Cow<'a, Name>,
         defining_symbol: Option<FileSymbolId>,
         definition: Option<Definition>,
         parent: Option<FileScopeId>,
     ) {
         let children_start = self.scopes.next_index() + 1;
 
         let scope = Scope {
-            name: name.clone(),
+            name,
             parent,
             defining_symbol,
             definition,
@@ -143,7 +149,7 @@ impl<'a> SemanticIndexBuilder<'a> {
 
     fn with_type_params(
         &mut self,
-        name: &Name,
+        name: &'a Name,
         with_params: &WithTypeParams<'a>,
         defining_symbol: FileSymbolId,
         nested: impl FnOnce(&mut Self) -> FileScopeId,

diff --git a/crates/red_knot_python_semantic/src/semantic_index/symbol.rs b/crates/red_knot_python_semantic/src/semantic_index/symbol.rs
@@ -215,16 +215,16 @@ impl FileScopeId {
 }
 
 #[derive(Debug, Eq, PartialEq)]
-pub struct Scope {
-    pub(super) name: Name,
+pub struct Scope<'db> {
+    pub(super) name: Cow<'db, Name>,
     pub(super) parent: Option<FileScopeId>,
     pub(super) definition: Option<Definition>,
     pub(super) defining_symbol: Option<FileSymbolId>,
     pub(super) kind: ScopeKind,
     pub(super) descendents: Range<FileScopeId>,
 }
 
-impl Scope {
+impl Scope<'_> {
     pub fn name(&self) -> &Name {
         &self.name
     }

diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs
@@ -254,7 +254,7 @@ impl ScopedTypeId for ScopedFunctionTypeId {
 #[derive(Debug, Eq, PartialEq, Clone)]
 pub struct FunctionType<'a> {
     /// name of the function at definition
-    name: Name,
+    name: &'a Name,
     /// types of all decorators on this function
     decorators: Vec<Type<'a>>,
 }
@@ -315,7 +315,7 @@ impl<'db> TypeId<'db, ScopedClassTypeId> {
 #[derive(Debug, Eq, PartialEq, Clone)]
 pub struct ClassType<'db> {
     /// Name of the class at definition
-    name: Name,
+    name: &'db Name,
 
     /// Types of all class bases
     bases: Vec<Type<'db>>,

diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs
@@ -123,7 +123,7 @@ impl<'db> TypeInferenceBuilder<'db> {
     }
 
     /// Infers the types of a `module`.
-    pub(super) fn infer_module(&mut self, module: &ast::ModModule) {
+    pub(super) fn infer_module(&mut self, module: &'db ast::ModModule) {
         self.infer_body(&module.body);
     }
 
@@ -133,7 +133,7 @@ impl<'db> TypeInferenceBuilder<'db> {
         }
     }
 
-    pub(super) fn infer_class_body(&mut self, class: &ast::StmtClassDef) {
+    pub(super) fn infer_class_body(&mut self, class: &'db ast::StmtClassDef) {
         self.infer_body(&class.body);
     }
 
@@ -143,17 +143,17 @@ impl<'db> TypeInferenceBuilder<'db> {
         }
     }
 
-    pub(super) fn infer_function_body(&mut self, function: &ast::StmtFunctionDef) {
+    pub(super) fn infer_function_body(&mut self, function: &'db ast::StmtFunctionDef) {
         self.infer_body(&function.body);
     }
 
-    fn infer_body(&mut self, suite: &[ast::Stmt]) {
+    fn infer_body(&mut self, suite: &'db [ast::Stmt]) {
         for statement in suite {
             self.infer_statement(statement);
         }
     }
 
-    fn infer_statement(&mut self, statement: &ast::Stmt) {
+    fn infer_statement(&mut self, statement: &'db ast::Stmt) {
         match statement {
             ast::Stmt::FunctionDef(function) => self.infer_function_definition_statement(function),
             ast::Stmt::ClassDef(class) => self.infer_class_definition_statement(class),
@@ -173,7 +173,7 @@ impl<'db> TypeInferenceBuilder<'db> {
         }
     }
 
-    fn infer_function_definition_statement(&mut self, function: &ast::StmtFunctionDef) {
+    fn infer_function_definition_statement(&mut self, function: &'db ast::StmtFunctionDef) {
         let ast::StmtFunctionDef {
             range: _,
             is_async: _,
@@ -198,7 +198,7 @@ impl<'db> TypeInferenceBuilder<'db> {
         }
 
         let function_ty = self.function_ty(FunctionType {
-            name: name.id.clone(),
+            name: &name.id,
             decorators: decorator_tys,
         });
 
@@ -214,7 +214,7 @@ impl<'db> TypeInferenceBuilder<'db> {
             .insert(Definition::FunctionDef(function_id), function_ty);
     }
 
-    fn infer_class_definition_statement(&mut self, class: &ast::StmtClassDef) {
+    fn infer_class_definition_statement(&mut self, class: &'db ast::StmtClassDef) {
         let ast::StmtClassDef {
             range: _,
             name,
@@ -247,7 +247,7 @@ impl<'db> TypeInferenceBuilder<'db> {
         assert_eq!(class_body_scope.kind(), ScopeKind::Class);
 
         let class_ty = self.class_ty(ClassType {
-            name: name.id.clone(),
+            name: &name.id,
             bases,
             body_scope: class_body_scope_id.to_scope_id(self.db, self.file_id),
         });
@@ -256,7 +256,7 @@ impl<'db> TypeInferenceBuilder<'db> {
             .insert(Definition::ClassDef(class_id), class_ty);
     }
 
-    fn infer_if_statement(&mut self, if_statement: &ast::StmtIf) {
+    fn infer_if_statement(&mut self, if_statement: &'db ast::StmtIf) {
         let ast::StmtIf {
             range: _,
             test,
@@ -328,7 +328,7 @@ impl<'db> TypeInferenceBuilder<'db> {
         );
     }
 
-    fn infer_for_statement(&mut self, for_statement: &ast::StmtFor) {
+    fn infer_for_statement(&mut self, for_statement: &'db ast::StmtFor) {
         let ast::StmtFor {
             range: _,
             target,