[ty] Add ty_server::Db trait (#21241)

Author: MichaReiser

crates/ty/src/lib.rs

@@ -450,12 +450,12 @@ impl ty_project::ProgressReporter for IndicatifReporter {
         self.bar.set_draw_target(self.printer.progress_target());
     }
 
-    fn report_checked_file(&self, db: &dyn Db, file: File, diagnostics: &[Diagnostic]) {
+    fn report_checked_file(&self, db: &ProjectDatabase, file: File, diagnostics: &[Diagnostic]) {
         self.collector.report_checked_file(db, file, diagnostics);
         self.bar.inc(1);
     }
 
-    fn report_diagnostics(&mut self, db: &dyn Db, diagnostics: Vec<Diagnostic>) {
+    fn report_diagnostics(&mut self, db: &ProjectDatabase, diagnostics: Vec<Diagnostic>) {
         self.collector.report_diagnostics(db, diagnostics);
     }
 }

crates/ty_project/src/lib.rs

@@ -124,12 +124,12 @@ pub trait ProgressReporter: Send + Sync {
     fn set_files(&mut self, files: usize);
 
     /// Report the completion of checking a given file along with its diagnostics.
-    fn report_checked_file(&self, db: &dyn Db, file: File, diagnostics: &[Diagnostic]);
+    fn report_checked_file(&self, db: &ProjectDatabase, file: File, diagnostics: &[Diagnostic]);
 
     /// Reports settings or IO related diagnostics. The diagnostics
     /// can belong to different files or no file at all.
     /// But it's never a file for which [`Self::report_checked_file`] gets called.
-    fn report_diagnostics(&mut self, db: &dyn Db, diagnostics: Vec<Diagnostic>);
+    fn report_diagnostics(&mut self, db: &ProjectDatabase, diagnostics: Vec<Diagnostic>);
 }
 
 /// Reporter that collects all diagnostics into a `Vec`.
@@ -149,7 +149,7 @@ impl CollectReporter {
 
 impl ProgressReporter for CollectReporter {
     fn set_files(&mut self, _files: usize) {}
-    fn report_checked_file(&self, _db: &dyn Db, _file: File, diagnostics: &[Diagnostic]) {
+    fn report_checked_file(&self, _db: &ProjectDatabase, _file: File, diagnostics: &[Diagnostic]) {
         if diagnostics.is_empty() {
             return;
         }
@@ -160,7 +160,7 @@ impl ProgressReporter for CollectReporter {
             .extend(diagnostics.iter().map(Clone::clone));
     }
 
-    fn report_diagnostics(&mut self, _db: &dyn Db, diagnostics: Vec<Diagnostic>) {
+    fn report_diagnostics(&mut self, _db: &ProjectDatabase, diagnostics: Vec<Diagnostic>) {
         self.0.get_mut().unwrap().extend(diagnostics);
     }
 }

crates/ty_server/src/db.rs

@@ -0,0 +1,33 @@
+use crate::NotebookDocument;
+use crate::session::index::Document;
+use crate::system::LSPSystem;
+use ruff_db::Db as _;
+use ruff_db::files::{File, FilePath};
+use ty_project::{Db as ProjectDb, ProjectDatabase};
+
+#[salsa::db]
+pub(crate) trait Db: ProjectDb {
+    /// Returns the LSP [`Document`] corresponding to `File` or
+    /// `None` if the file isn't open in the editor.
+    fn document(&self, file: File) -> Option<&Document>;
+
+    /// Returns the LSP [`NotebookDocument`] corresponding to `File` or
+    /// `None` if the file isn't open in the editor or if it isn't a notebook.
+    fn notebook_document(&self, file: File) -> Option<&NotebookDocument> {
+        self.document(file)?.as_notebook()
+    }
+}
+
+#[salsa::db]
+impl Db for ProjectDatabase {
+    fn document(&self, file: File) -> Option<&Document> {
+        self.system()
+            .as_any()
+            .downcast_ref::<LSPSystem>()
+            .and_then(|system| match file.path(self) {
+                FilePath::System(path) => system.system_path_to_document(path),
+                FilePath::SystemVirtual(path) => system.system_virtual_path_to_document(path),
+                FilePath::Vendored(_) => None,
+            })
+    }
+}

crates/ty_server/src/document/location.rs

@@ -1,9 +1,10 @@
-use crate::PositionEncoding;
-use crate::document::{FileRangeExt, ToRangeExt};
 use lsp_types::Location;
 use ruff_db::files::FileRange;
 use ty_ide::{NavigationTarget, ReferenceTarget};
-use ty_python_semantic::Db;
+
+use crate::Db;
+use crate::PositionEncoding;
+use crate::document::{FileRangeExt, ToRangeExt};
 
 pub(crate) trait ToLink {
     fn to_location(&self, db: &dyn Db, encoding: PositionEncoding) -> Option<Location>;

crates/ty_server/src/document/range.rs

@@ -1,6 +1,6 @@
 use super::PositionEncoding;
+use crate::Db;
 use crate::system::file_to_url;
-use ty_python_semantic::Db;
 
 use lsp_types as types;
 use lsp_types::{Location, Position, Url};

crates/ty_server/src/lib.rs

@@ -4,13 +4,15 @@ use anyhow::Context;
 use lsp_server::Connection;
 use ruff_db::system::{OsSystem, SystemPathBuf};
 
+use crate::db::Db;
 pub use crate::logging::{LogLevel, init_logging};
 pub use crate::server::{PartialWorkspaceProgress, PartialWorkspaceProgressParams, Server};
 pub use crate::session::{ClientOptions, DiagnosticMode};
 pub use document::{NotebookDocument, PositionEncoding, TextDocument};
 pub(crate) use session::Session;
 
 mod capabilities;
+mod db;
 mod document;
 mod logging;
 mod server;

crates/ty_server/src/server/api/diagnostics.rs

@@ -10,8 +10,9 @@ use rustc_hash::FxHashMap;
 use ruff_db::diagnostic::{Annotation, Severity, SubDiagnostic};
 use ruff_db::files::FileRange;
 use ruff_db::system::SystemPathBuf;
-use ty_project::{Db, ProjectDatabase};
+use ty_project::{Db as _, ProjectDatabase};
 
+use crate::Db;
 use crate::document::{FileRangeExt, ToRangeExt};
 use crate::session::DocumentSnapshot;
 use crate::session::client::Client;

crates/ty_server/src/server/api/notifications/did_change_watched_files.rs

@@ -8,7 +8,7 @@ use crate::system::AnySystemPath;
 use lsp_types as types;
 use lsp_types::{FileChangeType, notification as notif};
 use rustc_hash::FxHashMap;
-use ty_project::Db;
+use ty_project::Db as _;
 use ty_project::watch::{ChangeEvent, ChangedKind, CreatedKind, DeletedKind};
 
 pub(crate) struct DidChangeWatchedFiles;

crates/ty_server/src/server/api/requests/document_symbols.rs

@@ -4,8 +4,9 @@ use lsp_types::request::DocumentSymbolRequest;
 use lsp_types::{DocumentSymbol, DocumentSymbolParams, SymbolInformation, Url};
 use ruff_db::files::File;
 use ty_ide::{HierarchicalSymbols, SymbolId, SymbolInfo, document_symbols};
-use ty_project::{Db, ProjectDatabase};
+use ty_project::ProjectDatabase;
 
+use crate::Db;
 use crate::document::{PositionEncoding, ToRangeExt};
 use crate::server::api::symbols::{convert_symbol_kind, convert_to_lsp_symbol_information};
 use crate::server::api::traits::{

crates/ty_server/src/server/api/requests/execute_command.rs

@@ -9,7 +9,7 @@ use lsp_server::ErrorCode;
 use lsp_types::{self as types, request as req};
 use std::fmt::Write;
 use std::str::FromStr;
-use ty_project::Db;
+use ty_project::Db as _;
 
 pub(crate) struct ExecuteCommand;