Review comments

crates/red_knot/src/db.rs

@@ -11,7 +11,7 @@ use crate::lint::{Diagnostics, LintSemanticStorage, LintSyntaxStorage};
 use crate::module::{Module, ModuleData, ModuleName, ModuleResolver, ModuleSearchPath};
 use crate::parse::{Parsed, ParsedStorage};
 use crate::source::{Source, SourceStorage};
-use crate::symbols::{Definition, SymbolId, SymbolTable, SymbolTablesStorage};
+use crate::symbols::{Definition, GlobalSymbolId, SymbolTable, SymbolTablesStorage};
 use crate::types::{Type, TypeStore};
 
 mod jars;
@@ -121,16 +121,21 @@ pub trait SemanticDb: SourceDb {
 
     fn symbol_table(&self, file_id: FileId) -> QueryResult<Arc<SymbolTable>>;
 
+    fn resolve_global_symbol(
+        &self,
+        module: ModuleName,
+        name: &str,
+    ) -> QueryResult<Option<GlobalSymbolId>>;
+
     fn type_store(&self) -> QueryResult<&TypeStore>;
 
     fn infer_definition_type(
         &self,
-        file_id: FileId,
-        symbol_id: SymbolId,
+        symbol: GlobalSymbolId,
         definition: Definition,
     ) -> QueryResult<Type>;
 
-    fn infer_symbol_type(&self, file_id: FileId, symbol_id: SymbolId) -> QueryResult<Type>;
+    fn infer_symbol_type(&self, symbol: GlobalSymbolId) -> QueryResult<Type>;
 
     // mutations
 
@@ -183,7 +188,9 @@ pub(crate) mod tests {
     };
     use crate::parse::{parse, Parsed};
     use crate::source::{source_text, Source};
-    use crate::symbols::{symbol_table, Definition, SymbolId, SymbolTable};
+    use crate::symbols::{
+        resolve_global_symbol, symbol_table, Definition, GlobalSymbolId, SymbolTable,
+    };
     use crate::types::{infer_definition_type, infer_symbol_type, type_store, Type, TypeStore};
 
     use super::{SemanticDb, SemanticJar};
@@ -265,21 +272,28 @@ pub(crate) mod tests {
             symbol_table(self, file_id)
         }
 
+        fn resolve_global_symbol(
+            &self,
+            module: ModuleName,
+            name: &str,
+        ) -> QueryResult<Option<GlobalSymbolId>> {
+            resolve_global_symbol(self, module, name)
+        }
+
         fn type_store(&self) -> QueryResult<&TypeStore> {
             type_store(self)
         }
 
-        fn infer_symbol_type(&self, file_id: FileId, symbol_id: SymbolId) -> QueryResult<Type> {
-            infer_symbol_type(self, file_id, symbol_id)
+        fn infer_symbol_type(&self, symbol: GlobalSymbolId) -> QueryResult<Type> {
+            infer_symbol_type(self, symbol)
         }
 
         fn infer_definition_type(
             &self,
-            file_id: FileId,
-            symbol_id: SymbolId,
+            symbol: GlobalSymbolId,
             definition: Definition,
         ) -> QueryResult<Type> {
-            infer_definition_type(self, file_id, symbol_id, definition)
+            infer_definition_type(self, symbol, definition)
         }
 
         fn add_module(&mut self, path: &Path) -> Option<(Module, Vec<Arc<ModuleData>>)> {

crates/red_knot/src/lint.rs

@@ -12,7 +12,7 @@ use crate::files::FileId;
 use crate::module::ModuleName;
 use crate::parse::Parsed;
 use crate::source::Source;
-use crate::symbols::{Definition, SymbolId, SymbolTable};
+use crate::symbols::{Definition, GlobalSymbolId, SymbolId, SymbolTable};
 use crate::types::Type;
 
 #[tracing::instrument(level = "debug", skip(db))]
@@ -143,16 +143,14 @@ fn lint_unresolved_imports(context: &SemanticLintContext) -> QueryResult<()> {
 }
 
 fn lint_bad_overrides(context: &SemanticLintContext) -> QueryResult<()> {
-    // TODO possibly we should have a special marker on the real typing module (from typeshed) so
-    // if you have your own "typing" module in your project, we don't consider it THE typing module
-    let Some(typing_module) = context.db.resolve_module(ModuleName::new("typing"))? else {
-        debug_assert!(false, "typing module should always exist");
-        return Ok(());
-    };
-    let typing_file = context.db.module_to_file(typing_module)?;
-    let typing_table = context.db.symbol_table(typing_file)?;
-    let Some(typing_override) = typing_table.root_symbol_id_by_name("override") else {
-        debug_assert!(false, "typing.override should exist");
+    // TODO we should have a special marker on the real typing module (from typeshed) so if you
+    // have your own "typing" module in your project, we don't consider it THE typing module (and
+    // same for other stdlib modules that our lint rules care about)
+    let Some(typing_override) = context
+        .db
+        .resolve_global_symbol(ModuleName::new("typing"), "override")?
+    else {
+        // TODO once we bundle typeshed, this should be unreachable!()
         return Ok(());
     };
 
@@ -162,28 +160,21 @@ fn lint_bad_overrides(context: &SemanticLintContext) -> QueryResult<()> {
         if !matches!(definition, Definition::FunctionDef(_)) {
             continue;
         }
-        let ty = context
-            .db
-            .infer_definition_type(context.file_id, symbol, definition.clone())?;
+        let ty = context.db.infer_definition_type(
+            GlobalSymbolId {
+                file_id: context.file_id,
+                symbol_id: symbol,
+            },
+            definition.clone(),
+        )?;
         let Type::Function(func) = ty else {
-            debug_assert!(false, "type of a FunctionDef should always be a Function");
-            continue;
+            unreachable!("type of a FunctionDef should always be a Function");
         };
         let Some(class) = func.get_containing_class(context.db)? else {
             // not a method of a class
             continue;
         };
-        let mut override_decorated = false;
-        for deco_ty in func.function(context.db)?.decorators() {
-            let Type::Function(deco_func) = deco_ty else {
-                continue;
-            };
-            if deco_func.file() == typing_file && deco_func.symbol(context.db)? == typing_override {
-                override_decorated = true;
-                break;
-            }
-        }
-        if override_decorated {
+        if func.has_decorator(context.db, typing_override)? {
             let method_name = func.name(context.db)?;
             if class
                 .get_super_class_member(context.db, &method_name)?
@@ -229,7 +220,10 @@ impl<'a> SemanticLintContext<'a> {
     }
 
     pub fn infer_symbol_type(&self, symbol_id: SymbolId) -> QueryResult<Type> {
-        self.db.infer_symbol_type(self.file_id, symbol_id)
+        self.db.infer_symbol_type(GlobalSymbolId {
+            file_id: self.file_id,
+            symbol_id,
+        })
     }
 
     pub fn push_diagnostic(&self, diagnostic: String) {

crates/red_knot/src/program/mod.rs

@@ -16,7 +16,7 @@ use crate::module::{
 };
 use crate::parse::{parse, Parsed};
 use crate::source::{source_text, Source};
-use crate::symbols::{symbol_table, SymbolId, SymbolTable};
+use crate::symbols::{resolve_global_symbol, symbol_table, GlobalSymbolId, SymbolTable};
 use crate::types::{infer_definition_type, infer_symbol_type, type_store, Type};
 use crate::Workspace;
 
@@ -114,21 +114,28 @@ impl SemanticDb for Program {
         symbol_table(self, file_id)
     }
 
+    fn resolve_global_symbol(
+        &self,
+        module: ModuleName,
+        name: &str,
+    ) -> QueryResult<Option<crate::symbols::GlobalSymbolId>> {
+        resolve_global_symbol(self, module, name)
+    }
+
     fn type_store(&self) -> QueryResult<&crate::types::TypeStore> {
         type_store(self)
     }
 
-    fn infer_symbol_type(&self, file_id: FileId, symbol_id: SymbolId) -> QueryResult<Type> {
-        infer_symbol_type(self, file_id, symbol_id)
+    fn infer_symbol_type(&self, symbol: GlobalSymbolId) -> QueryResult<Type> {
+        infer_symbol_type(self, symbol)
     }
 
     fn infer_definition_type(
         &self,
-        file_id: FileId,
-        symbol_id: SymbolId,
+        symbol: GlobalSymbolId,
         definition: crate::symbols::Definition,
     ) -> QueryResult<Type> {
-        infer_definition_type(self, file_id, symbol_id, definition)
+        infer_definition_type(self, symbol, definition)
     }
 
     // Mutations

crates/red_knot/src/symbols.rs

@@ -21,7 +21,6 @@ use crate::files::FileId;
 use crate::module::ModuleName;
 use crate::Name;
 
-#[allow(unreachable_pub)]
 #[tracing::instrument(level = "debug", skip(db))]
 pub fn symbol_table<Db>(db: &Db, file_id: FileId) -> QueryResult<Arc<SymbolTable>>
 where
@@ -35,6 +34,35 @@ where
     })
 }
 
+#[tracing::instrument(level = "debug", skip(db))]
+pub fn resolve_global_symbol<Db>(
+    db: &Db,
+    module: ModuleName,
+    name: &str,
+) -> QueryResult<Option<GlobalSymbolId>>
+where
+    Db: SemanticDb + HasJar<SemanticJar>,
+{
+    let Some(typing_module) = db.resolve_module(module)? else {
+        return Ok(None);
+    };
+    let typing_file = db.module_to_file(typing_module)?;
+    let typing_table = db.symbol_table(typing_file)?;
+    let Some(typing_override) = typing_table.root_symbol_id_by_name(name) else {
+        return Ok(None);
+    };
+    Ok(Some(GlobalSymbolId {
+        file_id: typing_file,
+        symbol_id: typing_override,
+    }))
+}
+
+#[derive(Copy, Clone, Debug, PartialEq)]
+pub struct GlobalSymbolId {
+    pub(crate) file_id: FileId,
+    pub(crate) symbol_id: SymbolId,
+}
+
 type Map<K, V> = hashbrown::HashMap<K, V, ()>;
 
 #[newtype_index]
@@ -69,7 +97,9 @@ pub(crate) struct Scope {
     kind: ScopeKind,
     parent: Option<ScopeId>,
     children: Vec<ScopeId>,
+    /// the definition (e.g. class or function) that created this scope
     definition: Option<Definition>,
+    /// the symbol (e.g. class or function) that owns this scope
     defining_symbol: Option<SymbolId>,
     /// symbol IDs, hashed by symbol name
     symbols_by_name: Map<SymbolId, ()>,
@@ -149,6 +179,7 @@ impl Symbol {
 // TODO storing TypedNodeKey for definitions means we have to search to find them again in the AST;
 // this is at best O(log n). If looking up definitions is a bottleneck we should look for
 // alternatives here.
+// TODO intern Definitions in SymbolTable and reference using IDs?
 #[derive(Clone, Debug)]
 pub enum Definition {
     // For the import cases, we don't need reference to any arbitrary AST subtrees (annotations,
@@ -330,19 +361,20 @@ impl SymbolTable {
         &self.scopes_by_id[self.scope_id_of_symbol(symbol_id)]
     }
 
-    pub(crate) fn iter_parent_scopes(
+    pub(crate) fn parent_scopes(
         &self,
         scope_id: ScopeId,
-    ) -> ScopeIterator<ParentScopeIdsIterator> {
+    ) -> ScopeIterator<impl Iterator<Item = ScopeId> + '_> {
         ScopeIterator {
             table: self,
-            ids: ParentScopeIdsIterator {
-                table: self,
-                scope_id: Some(scope_id),
-            },
+            ids: std::iter::successors(Some(scope_id), |scope| self.scopes_by_id[*scope].parent),
         }
     }
 
+    pub(crate) fn parent_scope(&self, scope_id: ScopeId) -> Option<ScopeId> {
+        self.scopes_by_id[scope_id].parent
+    }
+
     pub(crate) fn scope_id_for_node(&self, node_key: &NodeKey) -> ScopeId {
         self.scopes_by_node[node_key]
     }
@@ -461,32 +493,13 @@ where
     }
 }
 
-pub(crate) struct ParentScopeIdsIterator<'a> {
-    table: &'a SymbolTable,
-    scope_id: Option<ScopeId>,
-}
-
-impl<'a> Iterator for ParentScopeIdsIterator<'a> {
-    type Item = ScopeId;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        if let Some(scope_id) = self.scope_id {
-            let parent = self.table.scopes_by_id[scope_id].parent;
-            self.scope_id = parent;
-            parent
-        } else {
-            None
-        }
-    }
-}
-
-impl<'a> FusedIterator for ParentScopeIdsIterator<'a> {}
-
+// TODO maybe get rid of this and just do all data access via methods on ScopeId?
 pub(crate) struct ScopeIterator<'a, I> {
     table: &'a SymbolTable,
     ids: I,
 }
 
+/// iterate (`ScopeId`, `Scope`) pairs for given `ScopeId` iterator
 impl<'a, I> Iterator for ScopeIterator<'a, I>
 where
     I: Iterator<Item = ScopeId>,