[red-knot] Simplify virtual file support (#13043)

## Summary

This PR simplifies the virtual file support in the red knot core,
specifically:

* Update `File::add_virtual_file` method to `File::virtual_file` which
will always create a new virtual file and override the existing entry in
the lookup table
* Add `VirtualFile` which is a wrapper around `File` and provides
methods to increment the file revision / close the virtual file
* Add a new `File::try_virtual_file` to lookup the `VirtualFile` from
`Files`
* Add `File::sync_virtual_path` which takes in the `SystemVirtualPath`,
looks up the `VirtualFile` for it and calls the `sync` method to
increment the file revision
* Removes the `virtual_path_metadata` method on `System` trait

## Test Plan

- [x] Make sure the existing red knot tests pass
- [x] Updated code works well with the LSP

crates/red_knot_server/src/system.rs

@@ -144,19 +144,6 @@ impl System for LSPSystem {
         }
     }
 
-    fn virtual_path_metadata(&self, path: &SystemVirtualPath) -> Result<Metadata> {
-        // Virtual paths only exists in the LSP system, so we don't need to check the OS system.
-        let document = self
-            .system_virtual_path_to_document_ref(path)?
-            .ok_or_else(|| virtual_path_not_found(path))?;
-
-        Ok(Metadata::new(
-            document_revision(&document),
-            None,
-            FileType::File,
-        ))
-    }
-
     fn read_virtual_path_to_string(&self, path: &SystemVirtualPath) -> Result<String> {
         let document = self
             .system_virtual_path_to_document_ref(path)?

crates/red_knot_wasm/src/lib.rs

@@ -234,13 +234,6 @@ impl System for WasmSystem {
         Notebook::from_source_code(&content)
     }
 
-    fn virtual_path_metadata(
-        &self,
-        _path: &SystemVirtualPath,
-    ) -> ruff_db::system::Result<Metadata> {
-        Err(not_found())
-    }
-
     fn read_virtual_path_to_string(
         &self,
         _path: &SystemVirtualPath,

crates/ruff_db/src/files.rs

@@ -12,7 +12,7 @@ use ruff_notebook::{Notebook, NotebookError};
 use crate::file_revision::FileRevision;
 use crate::files::file_root::FileRoots;
 use crate::files::private::FileStatus;
-use crate::system::{Metadata, SystemPath, SystemPathBuf, SystemVirtualPath, SystemVirtualPathBuf};
+use crate::system::{SystemPath, SystemPathBuf, SystemVirtualPath, SystemVirtualPathBuf};
 use crate::vendored::{VendoredPath, VendoredPathBuf};
 use crate::{vendored, Db, FxDashMap};
 
@@ -60,8 +60,8 @@ struct FilesInner {
     /// so that queries that depend on the existence of a file are re-executed when the file is created.
     system_by_path: FxDashMap<SystemPathBuf, File>,
 
-    /// Lookup table that maps [`SystemVirtualPathBuf`]s to salsa interned [`File`] instances.
-    system_virtual_by_path: FxDashMap<SystemVirtualPathBuf, File>,
+    /// Lookup table that maps [`SystemVirtualPathBuf`]s to [`VirtualFile`] instances.
+    system_virtual_by_path: FxDashMap<SystemVirtualPathBuf, VirtualFile>,
 
     /// Lookup table that maps vendored files to the salsa [`File`] ingredients.
     vendored_by_path: FxDashMap<VendoredPathBuf, File>,
@@ -147,31 +147,31 @@ impl Files {
         Ok(file)
     }
 
-    /// Looks up a virtual file by its `path`.
+    /// Create a new virtual file at the given path and store it for future lookups.
     ///
-    /// For a non-existing file, creates a new salsa [`File`] ingredient and stores it for future lookups.
-    ///
-    /// The operations fails if the system failed to provide a metadata for the path.
-    pub fn add_virtual_file(&self, db: &dyn Db, path: &SystemVirtualPath) -> Option<File> {
-        let file = match self.inner.system_virtual_by_path.entry(path.to_path_buf()) {
-            Entry::Occupied(entry) => *entry.get(),
-            Entry::Vacant(entry) => {
-                let metadata = db.system().virtual_path_metadata(path).ok()?;
-
-                tracing::trace!("Adding virtual file '{}'", path);
-
-                let file = File::builder(FilePath::SystemVirtual(path.to_path_buf()))
-                    .revision(metadata.revision())
-                    .permissions(metadata.permissions())
-                    .new(db);
-
-                entry.insert(file);
-
-                file
-            }
-        };
+    /// This will always create a new file, overwriting any existing file at `path` in the internal
+    /// storage.
+    pub fn virtual_file(&self, db: &dyn Db, path: &SystemVirtualPath) -> VirtualFile {
+        tracing::trace!("Adding virtual file {}", path);
+        let virtual_file = VirtualFile(
+            File::builder(FilePath::SystemVirtual(path.to_path_buf()))
+                .status(FileStatus::Exists)
+                .revision(FileRevision::zero())
+                .permissions(None)
+                .new(db),
+        );
+        self.inner
+            .system_virtual_by_path
+            .insert(path.to_path_buf(), virtual_file);
+        virtual_file
+    }
 
-        Some(file)
+    /// Tries to look up a virtual file by its path. Returns `None` if no such file exists yet.
+    pub fn try_virtual_file(&self, path: &SystemVirtualPath) -> Option<VirtualFile> {
+        self.inner
+            .system_virtual_by_path
+            .get(&path.to_path_buf())
+            .map(|entry| *entry.value())
     }
 
     /// Looks up the closest  root for `path`. Returns `None` if `path` isn't enclosed by any source root.
@@ -354,6 +354,13 @@ impl File {
         Self::sync_system_path(db, &absolute, None);
     }
 
+    /// Increments the revision for the virtual file at `path`.
+    pub fn sync_virtual_path(db: &mut dyn Db, path: &SystemVirtualPath) {
+        if let Some(virtual_file) = db.files().try_virtual_file(path) {
+            virtual_file.sync(db);
+        }
+    }
+
     /// Syncs the [`File`]'s state with the state of the file on the system.
     pub fn sync(self, db: &mut dyn Db) {
         let path = self.path(db).clone();
@@ -366,29 +373,20 @@ impl File {
             FilePath::Vendored(_) => {
                 // Readonly, can never be out of date.
             }
-            FilePath::SystemVirtual(system_virtual) => {
-                Self::sync_system_virtual_path(db, &system_virtual, self);
+            FilePath::SystemVirtual(_) => {
+                VirtualFile(self).sync(db);
             }
         }
     }
 
+    /// Private method providing the implementation for [`Self::sync_path`] and [`Self::sync`] for
+    /// system paths.
     fn sync_system_path(db: &mut dyn Db, path: &SystemPath, file: Option<File>) {
         let Some(file) = file.or_else(|| db.files().try_system(db, path)) else {
             return;
         };
-        let metadata = db.system().path_metadata(path);
-        Self::sync_impl(db, metadata, file);
-    }
-
-    fn sync_system_virtual_path(db: &mut dyn Db, path: &SystemVirtualPath, file: File) {
-        let metadata = db.system().virtual_path_metadata(path);
-        Self::sync_impl(db, metadata, file);
-    }
 
-    /// Private method providing the implementation for [`Self::sync_system_path`] and
-    /// [`Self::sync_system_virtual_path`].
-    fn sync_impl(db: &mut dyn Db, metadata: crate::system::Result<Metadata>, file: File) {
-        let (status, revision, permission) = match metadata {
+        let (status, revision, permission) = match db.system().path_metadata(path) {
             Ok(metadata) if metadata.file_type().is_file() => (
                 FileStatus::Exists,
                 metadata.revision(),
@@ -422,6 +420,35 @@ impl File {
     }
 }
 
+/// A virtual file that doesn't exist on the file system.
+///
+/// This is a wrapper around a [`File`] that provides additional methods to interact with a virtual
+/// file.
+#[derive(Copy, Clone)]
+pub struct VirtualFile(File);
+
+impl VirtualFile {
+    /// Returns the underlying [`File`].
+    pub fn file(&self) -> File {
+        self.0
+    }
+
+    /// Increments the revision of the underlying [`File`].
+    fn sync(&self, db: &mut dyn Db) {
+        let file = self.0;
+        tracing::debug!("Updating the revision of '{}'", file.path(db));
+        let current_revision = file.revision(db);
+        file.set_revision(db)
+            .to(FileRevision::new(current_revision.as_u128() + 1));
+    }
+
+    /// Closes the virtual file.
+    pub fn close(&self, db: &mut dyn Db) {
+        tracing::debug!("Closing virtual file '{}'", self.0.path(db));
+        self.0.set_status(db).to(FileStatus::NotFound);
+    }
+}
+
 // The types in here need to be public because they're salsa ingredients but we
 // don't want them to be publicly accessible. That's why we put them into a private module.
 mod private {

crates/ruff_db/src/parsed.rs

@@ -121,9 +121,9 @@ mod tests {
 
         db.write_virtual_file(path, "x = 10");
 
-        let file = db.files().add_virtual_file(&db, path).unwrap();
+        let virtual_file = db.files().virtual_file(&db, path);
 
-        let parsed = parsed_module(&db, file);
+        let parsed = parsed_module(&db, virtual_file.file());
 
         assert!(parsed.is_valid());
 
@@ -137,9 +137,9 @@ mod tests {
 
         db.write_virtual_file(path, "%timeit a = b");
 
-        let file = db.files().add_virtual_file(&db, path).unwrap();
+        let virtual_file = db.files().virtual_file(&db, path);
 
-        let parsed = parsed_module(&db, file);
+        let parsed = parsed_module(&db, virtual_file.file());
 
         assert!(parsed.is_valid());
 

crates/ruff_db/src/system.rs

@@ -63,9 +63,6 @@ pub trait System: Debug {
     /// representation fall-back to deserializing the notebook from a string.
     fn read_to_notebook(&self, path: &SystemPath) -> std::result::Result<Notebook, NotebookError>;
 
-    /// Reads the metadata of the virtual file at `path`.
-    fn virtual_path_metadata(&self, path: &SystemVirtualPath) -> Result<Metadata>;
-
     /// Reads the content of the virtual file at `path` into a [`String`].
     fn read_virtual_path_to_string(&self, path: &SystemVirtualPath) -> Result<String>;
 

crates/ruff_db/src/system/memory_fs.rs

@@ -136,22 +136,6 @@ impl MemoryFileSystem {
         ruff_notebook::Notebook::from_source_code(&content)
     }
 
-    pub(crate) fn virtual_path_metadata(
-        &self,
-        path: impl AsRef<SystemVirtualPath>,
-    ) -> Result<Metadata> {
-        let virtual_files = self.inner.virtual_files.read().unwrap();
-        let file = virtual_files
-            .get(&path.as_ref().to_path_buf())
-            .ok_or_else(not_found)?;
-
-        Ok(Metadata {
-            revision: file.last_modified.into(),
-            permissions: Some(MemoryFileSystem::PERMISSION),
-            file_type: FileType::File,
-        })
-    }
-
     pub(crate) fn read_virtual_path_to_string(
         &self,
         path: impl AsRef<SystemVirtualPath>,

crates/ruff_db/src/system/os.rs

@@ -77,10 +77,6 @@ impl System for OsSystem {
         Notebook::from_path(path.as_std_path())
     }
 
-    fn virtual_path_metadata(&self, _path: &SystemVirtualPath) -> Result<Metadata> {
-        Err(not_found())
-    }
-
     fn read_virtual_path_to_string(&self, _path: &SystemVirtualPath) -> Result<String> {
         Err(not_found())
     }

crates/ruff_db/src/system/test.rs

@@ -69,13 +69,6 @@ impl System for TestSystem {
         }
     }
 
-    fn virtual_path_metadata(&self, path: &SystemVirtualPath) -> Result<Metadata> {
-        match &self.inner {
-            TestSystemInner::Stub(fs) => fs.virtual_path_metadata(path),
-            TestSystemInner::System(system) => system.virtual_path_metadata(path),
-        }
-    }
-
     fn read_virtual_path_to_string(&self, path: &SystemVirtualPath) -> Result<String> {
         match &self.inner {
             TestSystemInner::Stub(fs) => fs.read_virtual_path_to_string(path),