Add prefetch cache for VFS to optimize file reads

- Introduced a `PrefetchCache` structure to store pre-read file contents and canonical paths, enhancing the speed of initial reads during snapshot building.
- Implemented parallel file reading using `rayon` and `walkdir` to populate the cache efficiently.
- Updated the `VfsInner` struct to include the prefetch cache and modified read methods to utilize it, reducing I/O operations.
- Added functions to set and clear the prefetch cache in the `Vfs` API, improving memory management during live operations.
- Enhanced tests to validate the functionality of the prefetch cache, ensuring correct behavior during file reads and canonicalization.

These changes aim to significantly improve the performance of the VFS by minimizing file access times during initial operations.

crates/memofs/src/lib.rs

@@ -22,6 +22,7 @@ mod noop_backend;
 mod snapshot;
 mod std_backend;
 
+use std::collections::HashMap;
 use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex, MutexGuard};
 use std::{io, str};
@@ -31,6 +32,16 @@ pub use noop_backend::NoopBackend;
 pub use snapshot::VfsSnapshot;
 pub use std_backend::{CriticalErrorHandler, StdBackend, WatcherCriticalError};
 
+/// Pre-read file contents and canonical paths for fast startup.
+///
+/// Populated in parallel (via rayon + walkdir) before snapshot building,
+/// then consumed by VFS reads during `snapshot_from_vfs()`. Cleared after
+/// the initial tree build to avoid stale data during live operation.
+pub struct PrefetchCache {
+    pub files: HashMap<PathBuf, Vec<u8>>,
+    pub canonical: HashMap<PathBuf, PathBuf>,
+}
+
 mod sealed {
     use super::*;
 
@@ -145,27 +156,39 @@ pub enum VfsEvent {
 struct VfsInner {
     backend: Box<dyn VfsBackend>,
     watch_enabled: bool,
+    prefetch_cache: Option<PrefetchCache>,
 }
 
 impl VfsInner {
-    fn read<P: AsRef<Path>>(&mut self, path: P) -> io::Result<Arc<Vec<u8>>> {
-        let path = path.as_ref();
+    /// Read raw bytes from the prefetch cache or the backend.
+    /// Removes the entry from the cache on hit to free memory.
+    fn read_raw(&mut self, path: &Path) -> io::Result<Vec<u8>> {
+        if let Some(cache) = &mut self.prefetch_cache {
+            if let Some(contents) = cache.files.remove(path) {
+                if self.watch_enabled {
+                    self.backend.watch(path)?;
+                }
+                return Ok(contents);
+            }
+        }
+
         let contents = self.backend.read(path)?;
 
         if self.watch_enabled {
             self.backend.watch(path)?;
         }
 
-        Ok(Arc::new(contents))
+        Ok(contents)
+    }
+
+    fn read<P: AsRef<Path>>(&mut self, path: P) -> io::Result<Arc<Vec<u8>>> {
+        let path = path.as_ref();
+        Ok(Arc::new(self.read_raw(path)?))
     }
 
     fn read_to_string<P: AsRef<Path>>(&mut self, path: P) -> io::Result<Arc<String>> {
         let path = path.as_ref();
-        let contents = self.backend.read(path)?;
-
-        if self.watch_enabled {
-            self.backend.watch(path)?;
-        }
+        let contents = self.read_raw(path)?;
 
         let contents_str = str::from_utf8(&contents).map_err(|_| {
             io::Error::new(
@@ -228,6 +251,13 @@ impl VfsInner {
 
     fn canonicalize<P: AsRef<Path>>(&mut self, path: P) -> io::Result<PathBuf> {
         let path = path.as_ref();
+
+        if let Some(cache) = &mut self.prefetch_cache {
+            if let Some(canonical) = cache.canonical.remove(path) {
+                return Ok(canonical);
+            }
+        }
+
         self.backend.canonicalize(path)
     }
 
@@ -300,6 +330,7 @@ impl Vfs {
         let lock = VfsInner {
             backend: Box::new(backend),
             watch_enabled: true,
+            prefetch_cache: None,
         };
 
         Self {
@@ -307,6 +338,24 @@ impl Vfs {
         }
     }
 
+    /// Load a prefetch cache for fast initial reads.
+    ///
+    /// File reads and canonicalize calls will check the cache before hitting
+    /// the backend. Call [`clear_prefetch_cache`] after the initial snapshot
+    /// build to free memory and ensure live operations get fresh data.
+    pub fn set_prefetch_cache(&self, cache: PrefetchCache) {
+        let mut inner = self.inner.lock().unwrap();
+        inner.prefetch_cache = Some(cache);
+    }
+
+    /// Drop the prefetch cache, freeing memory.
+    ///
+    /// After this call, all reads go through the backend as normal.
+    pub fn clear_prefetch_cache(&self) {
+        let mut inner = self.inner.lock().unwrap();
+        inner.prefetch_cache = None;
+    }
+
     /// Manually lock the Vfs, useful for large batches of operations.
     pub fn lock(&self) -> VfsLock<'_> {
         VfsLock {
@@ -617,7 +666,8 @@ impl VfsLock<'_> {
 
 #[cfg(test)]
 mod test {
-    use crate::{InMemoryFs, StdBackend, Vfs, VfsSnapshot};
+    use crate::{InMemoryFs, PrefetchCache, StdBackend, Vfs, VfsSnapshot};
+    use std::collections::HashMap;
     use std::io;
     use std::path::PathBuf;
 
@@ -693,4 +743,137 @@ mod test {
         let err = vfs.canonicalize(&file_path).unwrap_err();
         assert_eq!(err.kind(), io::ErrorKind::NotFound);
     }
+
+    fn make_prefetch(files: Vec<(&str, &[u8])>, canonical: Vec<(&str, &str)>) -> PrefetchCache {
+        PrefetchCache {
+            files: files
+                .into_iter()
+                .map(|(k, v)| (PathBuf::from(k), v.to_vec()))
+                .collect(),
+            canonical: canonical
+                .into_iter()
+                .map(|(k, v)| (PathBuf::from(k), PathBuf::from(v)))
+                .collect(),
+        }
+    }
+
+    #[test]
+    fn prefetch_cache_read_hit() {
+        let imfs = InMemoryFs::new();
+        let vfs = Vfs::new(imfs);
+        vfs.set_prefetch_cache(make_prefetch(vec![("test", b"cached")], vec![]));
+
+        assert_eq!(vfs.read("test").unwrap().as_slice(), b"cached");
+    }
+
+    #[test]
+    fn prefetch_cache_read_depletion() {
+        let mut imfs = InMemoryFs::new();
+        imfs.load_snapshot("test", VfsSnapshot::file("backend"))
+            .unwrap();
+        let vfs = Vfs::new(imfs);
+        vfs.set_prefetch_cache(make_prefetch(vec![("test", b"cached")], vec![]));
+
+        assert_eq!(vfs.read("test").unwrap().as_slice(), b"cached");
+        assert_eq!(vfs.read("test").unwrap().as_slice(), b"backend");
+    }
+
+    #[test]
+    fn prefetch_cache_read_miss_falls_through() {
+        let mut imfs = InMemoryFs::new();
+        imfs.load_snapshot("other", VfsSnapshot::file("backend"))
+            .unwrap();
+        let vfs = Vfs::new(imfs);
+        vfs.set_prefetch_cache(make_prefetch(vec![("test", b"cached")], vec![]));
+
+        assert_eq!(vfs.read("other").unwrap().as_slice(), b"backend");
+    }
+
+    #[test]
+    fn prefetch_cache_read_to_string_hit() {
+        let imfs = InMemoryFs::new();
+        let vfs = Vfs::new(imfs);
+        vfs.set_prefetch_cache(make_prefetch(vec![("test", b"hello")], vec![]));
+
+        assert_eq!(vfs.read_to_string("test").unwrap().as_str(), "hello");
+    }
+
+    #[test]
+    fn prefetch_cache_lf_normalized_hit() {
+        let imfs = InMemoryFs::new();
+        let vfs = Vfs::new(imfs);
+        vfs.set_prefetch_cache(make_prefetch(vec![("test", b"line1\r\nline2\r\n")], vec![]));
+
+        assert_eq!(
+            vfs.read_to_string_lf_normalized("test").unwrap().as_str(),
+            "line1\nline2\n"
+        );
+    }
+
+    #[test]
+    fn prefetch_cache_canonicalize_hit() {
+        let imfs = InMemoryFs::new();
+        let vfs = Vfs::new(imfs);
+        vfs.set_prefetch_cache(make_prefetch(vec![], vec![("src/foo", "/abs/src/foo")]));
+
+        assert_eq!(
+            vfs.canonicalize("src/foo").unwrap(),
+            PathBuf::from("/abs/src/foo")
+        );
+    }
+
+    #[test]
+    fn prefetch_cache_canonicalize_miss_falls_through() {
+        let mut imfs = InMemoryFs::new();
+        imfs.load_snapshot("/real/file.txt", VfsSnapshot::file("x"))
+            .unwrap();
+        let vfs = Vfs::new(imfs);
+        vfs.set_prefetch_cache(make_prefetch(vec![], vec![("other", "/abs/other")]));
+
+        assert_eq!(
+            vfs.canonicalize("/real/file.txt").unwrap(),
+            PathBuf::from("/real/file.txt")
+        );
+    }
+
+    #[test]
+    fn prefetch_cache_clear_restores_backend() {
+        let mut imfs = InMemoryFs::new();
+        imfs.load_snapshot("test", VfsSnapshot::file("backend"))
+            .unwrap();
+        let vfs = Vfs::new(imfs);
+        vfs.set_prefetch_cache(make_prefetch(vec![("test", b"cached")], vec![]));
+        vfs.clear_prefetch_cache();
+
+        assert_eq!(vfs.read("test").unwrap().as_slice(), b"backend");
+    }
+
+    #[test]
+    fn prefetch_cache_no_cache_same_as_none() {
+        let mut imfs = InMemoryFs::new();
+        imfs.load_snapshot("test", VfsSnapshot::file("data"))
+            .unwrap();
+        let vfs = Vfs::new(imfs);
+
+        assert_eq!(vfs.read("test").unwrap().as_slice(), b"data");
+    }
+
+    #[test]
+    fn prefetch_cache_watch_registered_on_cache_hit() {
+        let contents = "hello world".to_string();
+        let dir = tempfile::tempdir().unwrap();
+        let file_path = dir.path().join("file.txt");
+        fs_err::write(&file_path, &contents).unwrap();
+
+        let vfs = Vfs::new(StdBackend::new_for_testing());
+        let mut cache_files = HashMap::new();
+        cache_files.insert(file_path.clone(), contents.as_bytes().to_vec());
+        vfs.set_prefetch_cache(PrefetchCache {
+            files: cache_files,
+            canonical: HashMap::new(),
+        });
+
+        let result = vfs.read(&file_path).unwrap();
+        assert_eq!(result.as_slice(), contents.as_bytes());
+    }
 }