red-knot[salsa part 2]: Setup semantic DB and Jar (#11837)

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>

Cargo.lock

@@ -2329,6 +2329,7 @@ version = "0.0.0"
 dependencies = [
  "bitflags 2.5.0",
  "is-macro",
+ "ruff_db",
  "ruff_index",
  "ruff_python_ast",
  "ruff_python_parser",
@@ -2336,6 +2337,8 @@ dependencies = [
  "ruff_source_file",
  "ruff_text_size",
  "rustc-hash",
+ "salsa-2022",
+ "tracing",
 ]
 
 [[package]]

Cargo.toml

@@ -14,6 +14,7 @@ license = "MIT"
 [workspace.dependencies]
 ruff = { path = "crates/ruff" }
 ruff_cache = { path = "crates/ruff_cache" }
+ruff_db = { path = "crates/ruff_db" }
 ruff_diagnostics = { path = "crates/ruff_diagnostics" }
 ruff_formatter = { path = "crates/ruff_formatter" }
 ruff_index = { path = "crates/ruff_index" }
@@ -81,7 +82,7 @@ libcst = { version = "1.1.0", default-features = false }
 log = { version = "0.4.17" }
 lsp-server = { version = "0.7.6" }
 lsp-types = { git = "https://github.com/astral-sh/lsp-types.git", rev = "3512a9f", features = [
-    "proposed",
+  "proposed",
 ] }
 matchit = { version = "0.8.1" }
 memchr = { version = "2.7.1" }
@@ -111,7 +112,7 @@ serde-wasm-bindgen = { version = "0.6.4" }
 serde_json = { version = "1.0.113" }
 serde_test = { version = "1.0.152" }
 serde_with = { version = "3.6.0", default-features = false, features = [
-    "macros",
+  "macros",
 ] }
 shellexpand = { version = "3.0.0" }
 similar = { version = "2.4.0", features = ["inline"] }
@@ -138,10 +139,10 @@ unicode-normalization = { version = "0.1.23" }
 ureq = { version = "2.9.6" }
 url = { version = "2.5.0" }
 uuid = { version = "1.6.1", features = [
-    "v4",
-    "fast-rng",
-    "macro-diagnostics",
-    "js",
+  "v4",
+  "fast-rng",
+  "macro-diagnostics",
+  "js",
 ] }
 walkdir = { version = "2.3.2" }
 wasm-bindgen = { version = "0.2.92" }

crates/ruff_db/src/lib.rs

@@ -50,12 +50,14 @@ mod tests {
     use crate::file_system::{FileSystem, MemoryFileSystem};
     use crate::vfs::{VendoredPathBuf, Vfs};
     use crate::{Db, Jar};
+    use salsa::DebugWithDb;
+    use std::sync::Arc;
 
     /// Database that can be used for testing.
     ///
     /// Uses an in memory filesystem and it stubs out the vendored files by default.
     #[salsa::db(Jar)]
-    pub struct TestDb {
+    pub(crate) struct TestDb {
         storage: salsa::Storage<Self>,
         vfs: Vfs,
         file_system: MemoryFileSystem,
@@ -63,8 +65,7 @@ mod tests {
     }
 
     impl TestDb {
-        #[allow(unused)]
-        pub fn new() -> Self {
+        pub(crate) fn new() -> Self {
             let mut vfs = Vfs::default();
             vfs.stub_vendored::<VendoredPathBuf, String>([]);
 
@@ -77,20 +78,37 @@ mod tests {
         }
 
         #[allow(unused)]
-        pub fn file_system(&self) -> &MemoryFileSystem {
+        pub(crate) fn file_system(&self) -> &MemoryFileSystem {
             &self.file_system
         }
 
+        /// Empties the internal store of salsa events that have been emitted,
+        /// and returns them as a `Vec` (equivalent to [`std::mem::take`]).
+        ///
+        /// ## Panics
+        /// If there are pending database snapshots.
         #[allow(unused)]
-        pub fn events(&self) -> std::sync::Arc<std::sync::Mutex<Vec<salsa::Event>>> {
-            self.events.clone()
+        pub(crate) fn take_salsa_events(&mut self) -> Vec<salsa::Event> {
+            let inner = Arc::get_mut(&mut self.events)
+                .expect("expected no pending salsa database snapshots.");
+
+            std::mem::take(inner.get_mut().unwrap())
         }
 
-        pub fn file_system_mut(&mut self) -> &mut MemoryFileSystem {
+        /// Clears the emitted salsa events.
+        ///
+        /// ## Panics
+        /// If there are pending database snapshots.
+        #[allow(unused)]
+        pub(crate) fn clear_salsa_events(&mut self) {
+            self.take_salsa_events();
+        }
+
+        pub(crate) fn file_system_mut(&mut self) -> &mut MemoryFileSystem {
             &mut self.file_system
         }
 
-        pub fn vfs_mut(&mut self) -> &mut Vfs {
+        pub(crate) fn vfs_mut(&mut self) -> &mut Vfs {
             &mut self.vfs
         }
     }
@@ -107,7 +125,7 @@ mod tests {
 
     impl salsa::Database for TestDb {
         fn salsa_event(&self, event: salsa::Event) {
-            tracing::trace!("event: {:?}", event);
+            tracing::trace!("event: {:?}", event.debug(self));
             let mut events = self.events.lock().unwrap();
             events.push(event);
         }

crates/ruff_db/src/source.rs

@@ -0,0 +1,128 @@
+use std::ops::Deref;
+use std::sync::Arc;
+
+use ruff_source_file::LineIndex;
+
+use crate::vfs::VfsFile;
+use crate::Db;
+
+/// Reads the content of file.
+#[salsa::tracked]
+pub fn source_text(db: &dyn Db, file: VfsFile) -> SourceText {
+    let content = file.read(db);
+
+    SourceText {
+        inner: Arc::from(content),
+    }
+}
+
+/// Computes the [`LineIndex`] for `file`.
+#[salsa::tracked]
+pub fn line_index(db: &dyn Db, file: VfsFile) -> LineIndex {
+    let source = source_text(db, file);
+
+    LineIndex::from_source_text(&source)
+}
+
+/// The source text of a [`VfsFile`].
+///
+/// Cheap cloneable in `O(1)`.
+#[derive(Clone, Eq, PartialEq)]
+pub struct SourceText {
+    inner: Arc<str>,
+}
+
+impl SourceText {
+    pub fn as_str(&self) -> &str {
+        &self.inner
+    }
+}
+
+impl Deref for SourceText {
+    type Target = str;
+
+    fn deref(&self) -> &str {
+        self.as_str()
+    }
+}
+
+impl std::fmt::Debug for SourceText {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_tuple("SourceText").field(&self.inner).finish()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use filetime::FileTime;
+    use salsa::EventKind;
+
+    use ruff_source_file::OneIndexed;
+    use ruff_text_size::TextSize;
+
+    use crate::file_system::FileSystemPath;
+    use crate::source::{line_index, source_text};
+    use crate::tests::TestDb;
+    use crate::Db;
+
+    #[test]
+    fn re_runs_query_when_file_revision_changes() {
+        let mut db = TestDb::new();
+        let path = FileSystemPath::new("test.py");
+
+        db.file_system_mut().write_file(path, "x = 10".to_string());
+
+        let file = db.file(path);
+
+        assert_eq!(&*source_text(&db, file), "x = 10");
+
+        db.file_system_mut().write_file(path, "x = 20".to_string());
+        file.set_revision(&mut db).to(FileTime::now().into());
+
+        assert_eq!(&*source_text(&db, file), "x = 20");
+    }
+
+    #[test]
+    fn text_is_cached_if_revision_is_unchanged() {
+        let mut db = TestDb::new();
+        let path = FileSystemPath::new("test.py");
+
+        db.file_system_mut().write_file(path, "x = 10".to_string());
+
+        let file = db.file(path);
+
+        assert_eq!(&*source_text(&db, file), "x = 10");
+
+        // Change the file permission only
+        file.set_permissions(&mut db).to(Some(0o777));
+
+        db.events().lock().unwrap().clear();
+        assert_eq!(&*source_text(&db, file), "x = 10");
+
+        let events = db.events();
+        let events = events.lock().unwrap();
+
+        assert!(!events
+            .iter()
+            .any(|event| matches!(event.kind, EventKind::WillExecute { .. })));
+    }
+
+    #[test]
+    fn line_index_for_source() {
+        let mut db = TestDb::new();
+        let path = FileSystemPath::new("test.py");
+
+        db.file_system_mut()
+            .write_file(path, "x = 10\ny = 20".to_string());
+
+        let file = db.file(path);
+        let index = line_index(&db, file);
+        let text = source_text(&db, file);
+
+        assert_eq!(index.line_count(), 2);
+        assert_eq!(
+            index.line_start(OneIndexed::from_zero_indexed(0), &text),
+            TextSize::new(0)
+        );
+    }
+}

crates/ruff_python_semantic/Cargo.toml

@@ -14,6 +14,7 @@ license = { workspace = true }
 doctest = false
 
 [dependencies]
+ruff_db = { workspace = true, optional = true }
 ruff_index = { workspace = true }
 ruff_python_ast = { workspace = true }
 ruff_python_stdlib = { workspace = true }
@@ -22,6 +23,8 @@ ruff_text_size = { workspace = true }
 
 bitflags = { workspace = true }
 is-macro = { workspace = true }
+salsa = { workspace = true, optional = true }
+tracing = { workspace = true, optional = true }
 rustc-hash = { workspace = true }
 
 [dev-dependencies]
@@ -29,3 +32,6 @@ ruff_python_parser = { workspace = true }
 
 [lints]
 workspace = true
+
+[features]
+red_knot = ["dep:ruff_db", "dep:salsa", "dep:tracing"]

crates/ruff_python_semantic/src/db.rs

@@ -0,0 +1,91 @@
+use ruff_db::{Db as SourceDb, Upcast};
+use salsa::DbWithJar;
+
+// Salsa doesn't support a struct without fields, so allow the clippy lint for now.
+#[allow(clippy::empty_structs_with_brackets)]
+#[salsa::jar(db=Db)]
+pub struct Jar();
+
+/// Database giving access to semantic information about a Python program.
+pub trait Db: SourceDb + DbWithJar<Jar> + Upcast<dyn SourceDb> {}
+
+#[cfg(test)]
+mod tests {
+    use super::{Db, Jar};
+    use ruff_db::file_system::{FileSystem, MemoryFileSystem};
+    use ruff_db::vfs::Vfs;
+    use ruff_db::{Db as SourceDb, Jar as SourceJar, Upcast};
+    use salsa::DebugWithDb;
+
+    #[salsa::db(Jar, SourceJar)]
+    pub(crate) struct TestDb {
+        storage: salsa::Storage<Self>,
+        vfs: Vfs,
+        file_system: MemoryFileSystem,
+        events: std::sync::Arc<std::sync::Mutex<Vec<salsa::Event>>>,
+    }
+
+    impl TestDb {
+        #[allow(unused)]
+        pub(crate) fn new() -> Self {
+            Self {
+                storage: salsa::Storage::default(),
+                file_system: MemoryFileSystem::default(),
+                events: std::sync::Arc::default(),
+                vfs: Vfs::with_stubbed_vendored(),
+            }
+        }
+
+        #[allow(unused)]
+        pub(crate) fn memory_file_system(&self) -> &MemoryFileSystem {
+            &self.file_system
+        }
+
+        #[allow(unused)]
+        pub(crate) fn memory_file_system_mut(&mut self) -> &mut MemoryFileSystem {
+            &mut self.file_system
+        }
+
+        #[allow(unused)]
+        pub(crate) fn vfs_mut(&mut self) -> &mut Vfs {
+            &mut self.vfs
+        }
+    }
+
+    impl SourceDb for TestDb {
+        fn file_system(&self) -> &dyn FileSystem {
+            &self.file_system
+        }
+
+        fn vfs(&self) -> &Vfs {
+            &self.vfs
+        }
+    }
+
+    impl Upcast<dyn SourceDb> for TestDb {
+        fn upcast(&self) -> &(dyn SourceDb + 'static) {
+            self
+        }
+    }
+
+    impl Db for TestDb {}
+
+    impl salsa::Database for TestDb {
+        fn salsa_event(&self, event: salsa::Event) {
+            tracing::trace!("event: {:?}", event.debug(self));
+            let mut events = self.events.lock().unwrap();
+            events.push(event);
+        }
+    }
+
+    impl salsa::ParallelDatabase for TestDb {
+        fn snapshot(&self) -> salsa::Snapshot<Self> {
+            salsa::Snapshot::new(Self {
+                storage: self.storage.snapshot(),
+                vfs: self.vfs.snapshot(),
+                file_system: self.file_system.snapshot(),
+                events: self.events.clone(),
+            })
+        }
+    }
+}

crates/ruff_python_semantic/src/lib.rs

@@ -2,6 +2,8 @@ pub mod analyze;
 mod binding;
 mod branches;
 mod context;
+#[cfg(feature = "red_knot")]
+mod db;
 mod definition;
 mod globals;
 mod model;
@@ -20,3 +22,6 @@ pub use nodes::*;
 pub use reference::*;
 pub use scope::*;
 pub use star_import::*;
+
+#[cfg(feature = "red_knot")]
+pub use db::{Db, Jar};