[red-knot] Remove <Db: SemanticDb> contraints in favor of dynamic d…

…ispatch (#11339)
astral-sh · May 8, 2024 · 4541337 · 4541337
1 parent 8e9ddee
commit 4541337
Show file tree

Hide file tree

Showing 11 changed files with 222 additions and 332 deletions.
diff --git a/crates/red_knot/src/db.rs b/crates/red_knot/src/db.rs
@@ -1,4 +1,3 @@
-use std::path::Path;
 use std::sync::Arc;
 
 pub use jars::{HasJar, HasJars};
@@ -7,12 +6,12 @@ pub use runtime::DbRuntime;
 pub use storage::JarsStorage;
 
 use crate::files::FileId;
-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::{SymbolId, SymbolTable, SymbolTablesStorage};
-use crate::types::{Type, TypeStore};
+use crate::lint::{LintSemanticStorage, LintSyntaxStorage};
+use crate::module::ModuleResolver;
+use crate::parse::ParsedStorage;
+use crate::source::SourceStorage;
+use crate::symbols::SymbolTablesStorage;
+use crate::types::TypeStore;
 
 mod jars;
 mod query;
@@ -61,6 +60,8 @@ pub trait ParallelDatabase: Database + Send {
     fn snapshot(&self) -> Snapshot<Self>;
 }
 
+pub trait DbWithJar<Jar>: Database + HasJar<Jar> {}
+
 /// Readonly snapshot of a database.
 ///
 /// ## Dead locks
@@ -96,45 +97,24 @@ where
     }
 }
 
+pub trait Upcast<T: ?Sized> {
+    fn upcast(&self) -> &T;
+}
+
 // Red knot specific databases code.
 
-pub trait SourceDb: Database {
+pub trait SourceDb: DbWithJar<SourceJar> {
     // queries
     fn file_id(&self, path: &std::path::Path) -> FileId;
 
     fn file_path(&self, file_id: FileId) -> Arc<std::path::Path>;
-
-    fn source(&self, file_id: FileId) -> QueryResult<Source>;
-
-    fn parse(&self, file_id: FileId) -> QueryResult<Parsed>;
-}
-
-pub trait SemanticDb: SourceDb {
-    // queries
-    fn resolve_module(&self, name: ModuleName) -> QueryResult<Option<Module>>;
-
-    fn file_to_module(&self, file_id: FileId) -> QueryResult<Option<Module>>;
-
-    fn path_to_module(&self, path: &Path) -> QueryResult<Option<Module>>;
-
-    fn symbol_table(&self, file_id: FileId) -> QueryResult<Arc<SymbolTable>>;
-
-    fn infer_symbol_type(&self, file_id: FileId, symbol_id: SymbolId) -> QueryResult<Type>;
-
-    // mutations
-
-    fn add_module(&mut self, path: &Path) -> Option<(Module, Vec<Arc<ModuleData>>)>;
-
-    fn set_module_search_paths(&mut self, paths: Vec<ModuleSearchPath>);
 }
 
-pub trait LintDb: SemanticDb {
-    fn lint_syntax(&self, file_id: FileId) -> QueryResult<Diagnostics>;
+pub trait SemanticDb: SourceDb + DbWithJar<SemanticJar> + Upcast<dyn SourceDb> {}
 
-    fn lint_semantic(&self, file_id: FileId) -> QueryResult<Diagnostics>;
-}
+pub trait LintDb: SemanticDb + DbWithJar<LintJar> + Upcast<dyn SemanticDb> {}
 
-pub trait Db: LintDb {}
+pub trait Db: LintDb + Upcast<dyn LintDb> {}
 
 #[derive(Debug, Default)]
 pub struct SourceJar {
@@ -161,19 +141,10 @@ pub(crate) mod tests {
     use std::sync::Arc;
 
     use crate::db::{
-        Database, DbRuntime, HasJar, HasJars, JarsStorage, LintDb, LintJar, QueryResult, SourceDb,
-        SourceJar,
+        Database, DbRuntime, DbWithJar, HasJar, HasJars, JarsStorage, LintDb, LintJar, QueryResult,
+        SourceDb, SourceJar, Upcast,
     };
     use crate::files::{FileId, Files};
-    use crate::lint::{lint_semantic, lint_syntax, Diagnostics};
-    use crate::module::{
-        add_module, file_to_module, path_to_module, resolve_module, set_module_search_paths,
-        Module, ModuleData, ModuleName, ModuleSearchPath,
-    };
-    use crate::parse::{parse, Parsed};
-    use crate::source::{source_text, Source};
-    use crate::symbols::{symbol_table, SymbolId, SymbolTable};
-    use crate::types::{infer_symbol_type, Type};
 
     use super::{SemanticDb, SemanticJar};
 
@@ -223,56 +194,36 @@ pub(crate) mod tests {
         fn file_path(&self, file_id: FileId) -> Arc<Path> {
             self.files.path(file_id)
         }
-
-        fn source(&self, file_id: FileId) -> QueryResult<Source> {
-            source_text(self, file_id)
-        }
-
-        fn parse(&self, file_id: FileId) -> QueryResult<Parsed> {
-            parse(self, file_id)
-        }
     }
 
-    impl SemanticDb for TestDb {
-        fn resolve_module(&self, name: ModuleName) -> QueryResult<Option<Module>> {
-            resolve_module(self, name)
-        }
-
-        fn file_to_module(&self, file_id: FileId) -> QueryResult<Option<Module>> {
-            file_to_module(self, file_id)
-        }
+    impl DbWithJar<SourceJar> for TestDb {}
 
-        fn path_to_module(&self, path: &Path) -> QueryResult<Option<Module>> {
-            path_to_module(self, path)
+    impl Upcast<dyn SourceDb> for TestDb {
+        fn upcast(&self) -> &(dyn SourceDb + 'static) {
+            self
         }
+    }
 
-        fn symbol_table(&self, file_id: FileId) -> QueryResult<Arc<SymbolTable>> {
-            symbol_table(self, file_id)
-        }
+    impl SemanticDb for TestDb {}
 
-        fn infer_symbol_type(&self, file_id: FileId, symbol_id: SymbolId) -> QueryResult<Type> {
-            infer_symbol_type(self, file_id, symbol_id)
-        }
+    impl DbWithJar<SemanticJar> for TestDb {}
 
-        fn add_module(&mut self, path: &Path) -> Option<(Module, Vec<Arc<ModuleData>>)> {
-            add_module(self, path)
-        }
-
-        fn set_module_search_paths(&mut self, paths: Vec<ModuleSearchPath>) {
-            set_module_search_paths(self, paths);
+    impl Upcast<dyn SemanticDb> for TestDb {
+        fn upcast(&self) -> &(dyn SemanticDb + 'static) {
+            self
         }
     }
 
-    impl LintDb for TestDb {
-        fn lint_syntax(&self, file_id: FileId) -> QueryResult<Diagnostics> {
-            lint_syntax(self, file_id)
-        }
+    impl LintDb for TestDb {}
 
-        fn lint_semantic(&self, file_id: FileId) -> QueryResult<Diagnostics> {
-            lint_semantic(self, file_id)
+    impl Upcast<dyn LintDb> for TestDb {
+        fn upcast(&self) -> &(dyn LintDb + 'static) {
+            self
         }
     }
 
+    impl DbWithJar<LintJar> for TestDb {}
+
     impl HasJars for TestDb {
         type Jars = (SourceJar, SemanticJar, LintJar);
 

diff --git a/crates/red_knot/src/lint.rs b/crates/red_knot/src/lint.rs
@@ -7,19 +7,17 @@ use ruff_python_ast::visitor::Visitor;
 use ruff_python_ast::{ModModule, StringLiteral};
 
 use crate::cache::KeyValueCache;
-use crate::db::{HasJar, LintDb, LintJar, QueryResult, SemanticDb};
+use crate::db::{LintDb, LintJar, QueryResult};
 use crate::files::FileId;
-use crate::parse::Parsed;
-use crate::source::Source;
-use crate::symbols::{Definition, SymbolId, SymbolTable};
-use crate::types::Type;
+use crate::parse::{parse, Parsed};
+use crate::source::{source_text, Source};
+use crate::symbols::{symbol_table, Definition, SymbolId, SymbolTable};
+use crate::types::{infer_symbol_type, Type};
 
 #[tracing::instrument(level = "debug", skip(db))]
-pub(crate) fn lint_syntax<Db>(db: &Db, file_id: FileId) -> QueryResult<Diagnostics>
-where
-    Db: LintDb + HasJar<LintJar>,
-{
-    let storage = &db.jar()?.lint_syntax;
+pub(crate) fn lint_syntax(db: &dyn LintDb, file_id: FileId) -> QueryResult<Diagnostics> {
+    let lint_jar: &LintJar = db.jar()?;
+    let storage = &lint_jar.lint_syntax;
 
     #[allow(clippy::print_stdout)]
     if std::env::var("RED_KNOT_SLOW_LINT").is_ok() {
@@ -33,10 +31,10 @@ where
     storage.get(&file_id, |file_id| {
         let mut diagnostics = Vec::new();
 
-        let source = db.source(*file_id)?;
+        let source = source_text(db.upcast(), *file_id)?;
         lint_lines(source.text(), &mut diagnostics);
 
-        let parsed = db.parse(*file_id)?;
+        let parsed = parse(db.upcast(), *file_id)?;
 
         if parsed.errors().is_empty() {
             let ast = parsed.ast();
@@ -73,16 +71,14 @@ fn lint_lines(source: &str, diagnostics: &mut Vec<String>) {
 }
 
 #[tracing::instrument(level = "debug", skip(db))]
-pub(crate) fn lint_semantic<Db>(db: &Db, file_id: FileId) -> QueryResult<Diagnostics>
-where
-    Db: LintDb + HasJar<LintJar>,
-{
-    let storage = &db.jar()?.lint_semantic;
+pub(crate) fn lint_semantic(db: &dyn LintDb, file_id: FileId) -> QueryResult<Diagnostics> {
+    let lint_jar: &LintJar = db.jar()?;
+    let storage = &lint_jar.lint_semantic;
 
     storage.get(&file_id, |file_id| {
-        let source = db.source(*file_id)?;
-        let parsed = db.parse(*file_id)?;
-        let symbols = db.symbol_table(*file_id)?;
+        let source = source_text(db.upcast(), *file_id)?;
+        let parsed = parse(db.upcast(), *file_id)?;
+        let symbols = symbol_table(db.upcast(), *file_id)?;
 
         let context = SemanticLintContext {
             file_id: *file_id,
@@ -145,7 +141,7 @@ pub struct SemanticLintContext<'a> {
     source: Source,
     parsed: Parsed,
     symbols: Arc<SymbolTable>,
-    db: &'a dyn SemanticDb,
+    db: &'a dyn LintDb,
     diagnostics: RefCell<Vec<String>>,
 }
 
@@ -167,7 +163,7 @@ 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)
+        infer_symbol_type(self.db.upcast(), self.file_id, symbol_id)
     }
 
     pub fn push_diagnostic(&self, diagnostic: String) {

diff --git a/crates/red_knot/src/main.rs b/crates/red_knot/src/main.rs
@@ -11,8 +11,8 @@ use tracing_subscriber::layer::{Context, Filter, SubscriberExt};
 use tracing_subscriber::{Layer, Registry};
 use tracing_tree::time::Uptime;
 
-use red_knot::db::{HasJar, ParallelDatabase, QueryError, SemanticDb, SourceDb, SourceJar};
-use red_knot::module::{ModuleSearchPath, ModuleSearchPathKind};
+use red_knot::db::{HasJar, ParallelDatabase, QueryError, SourceDb, SourceJar};
+use red_knot::module::{set_module_search_paths, ModuleSearchPath, ModuleSearchPathKind};
 use red_knot::program::check::ExecutionMode;
 use red_knot::program::{FileWatcherChange, Program};
 use red_knot::watch::FileWatcher;
@@ -49,7 +49,7 @@ fn main() -> anyhow::Result<()> {
         ModuleSearchPathKind::FirstParty,
     );
     let mut program = Program::new(workspace);
-    program.set_module_search_paths(vec![workspace_search_path]);
+    set_module_search_paths(&mut program, vec![workspace_search_path]);
 
     let entry_id = program.file_id(entry_point);
     program.workspace_mut().open_file(entry_id);