state: load exact module-source overrides before normal sources

pyrefly/benches/commit.rs

@@ -161,7 +161,7 @@ fn checked_state() -> (State, Vec<Handle>) {
     let mut transaction = state.new_committable_transaction(Require::Exports, None);
     transaction.as_mut().set_memory(
         (0..MODULES)
-            .map(|index| (memory_path(index), Some(Arc::new(contents(index)))))
+            .map(|index| (module_path(index), Some(Arc::new(contents(index)))))
             .collect(),
     );
     transaction.as_mut().run(&handles, Require::Errors, None);
@@ -212,7 +212,7 @@ fn prepare_all<'a>(
                 let mut text = source(index);
                 text.push_str(&format!("\n_touch{nonce}: int = {nonce}\n"));
                 (
-                    memory_path(index),
+                    module_path(index),
                     Some(Arc::new(FileContents::from_source(text))),
                 )
             })

pyrefly/benches/micro.rs

@@ -94,7 +94,7 @@ static SHARED_STATE: LazyLock<State> = LazyLock::new(|| {
     );
     let mut t = state.new_committable_transaction(Require::Errors, None);
     t.as_mut().set_memory(vec![(
-        PathBuf::from("_bench_init.py"),
+        ModulePath::memory(PathBuf::from("_bench_init.py")),
         Some(Arc::new(FileContents::from_source(String::new()))),
     )]);
     t.as_mut().run(&[h], Require::Errors, None);
@@ -130,7 +130,7 @@ fn check_module_error_counts(
         sys_info,
     );
     let mut t = SHARED_STATE.transaction();
-    t.set_memory(vec![(PathBuf::from(path), Some(code))]);
+    t.set_memory(vec![(ModulePath::memory(PathBuf::from(path)), Some(code))]);
     t.run(&[h.dupe()], require, None);
     let errors = t.get_errors([&h]);
     let collected = errors.collect_errors();

pyrefly/lib/commands/check.rs

@@ -1877,7 +1877,7 @@ impl CheckArgs {
 
         // Add the snippet source to the transaction's memory
         transaction.as_mut().set_memory(vec![(
-            PathBuf::from(module_path.as_path()),
+            module_path,
             Some(Arc::new(FileContents::from_source(code))),
         )]);
 

pyrefly/lib/embed.rs

@@ -60,7 +60,7 @@ pub struct Checker {
     config: ArcId<ConfigFile>,
     stdlib: Vec<(ModuleName, ModulePath)>,
     /// Loaded into memory by the first [`Checker::check`].
-    pending_stdlib: Vec<(PathBuf, Option<Arc<FileContents>>)>,
+    pending_stdlib: Vec<(ModulePath, Option<Arc<FileContents>>)>,
 }
 
 impl Checker {
@@ -114,7 +114,7 @@ impl Checker {
                 .iter()
                 .map(|(_, path, source)| {
                     (
-                        path.as_path().to_path_buf(),
+                        path.dupe(),
                         Some(Arc::new(FileContents::from_source(source.clone()))),
                     )
                 })
@@ -146,7 +146,7 @@ impl Checker {
             .into_iter()
             .chain(files.iter().map(|(name, source)| {
                 (
-                    memory_path(name).as_path().to_path_buf(),
+                    memory_path(name),
                     Some(Arc::new(FileContents::from_source((*source).to_owned()))),
                 )
             }))

pyrefly/lib/error/suppress.rs

@@ -975,7 +975,7 @@ mod tests {
         );
         let mut transaction = state.new_transaction(Require::Exports, None);
         transaction.set_memory(vec![(
-            get_path(&tdir),
+            ModulePath::memory(get_path(&tdir)),
             Some(Arc::new(FileContents::from_source(contents.to_owned()))),
         )]);
         transaction.run(&[handle.dupe()], Require::Everything, None);

pyrefly/lib/lsp/non_wasm/server.rs

@@ -3293,7 +3293,7 @@ impl Server {
                 (
                     make_open_handle(&self.state, path),
                     (
-                        path.clone(),
+                        ModulePath::memory(path.clone()),
                         Some(Arc::new(file.contents.to_file_contents())),
                     ),
                 )
@@ -4764,7 +4764,9 @@ impl Server {
                 let mut transaction = server
                     .state
                     .new_committable_transaction(Require::Exports, None);
-                transaction.as_mut().set_memory(vec![(path.clone(), None)]);
+                transaction
+                    .as_mut()
+                    .set_memory(vec![(ModulePath::memory(path.clone()), None)]);
                 let _ = server.validate_in_memory_for_transaction(
                     transaction.as_mut(),
                     telemetry_event,

pyrefly/lib/playground.rs

@@ -326,10 +326,10 @@ impl Playground {
 
             source_db.insert(module_name, memory_path.dupe());
 
-            let handle = Handle::new(module_name, memory_path, self.sys_info.dupe());
+            let handle = Handle::new(module_name, memory_path.dupe(), self.sys_info.dupe());
             self.handles.insert(filename.clone(), handle);
             file_contents.push((
-                module_path,
+                memory_path,
                 Some(Arc::new(FileContents::from_source(content.clone()))),
             ));
         }
@@ -371,7 +371,7 @@ impl Playground {
 
     pub fn update_single_file(&mut self, filename: String, content: String) {
         if let Some(_handle) = self.handles.get(&filename) {
-            let module_path = PathBuf::from(&filename);
+            let module_path = ModulePath::memory(PathBuf::from(&filename));
             let file_content = vec![(
                 module_path,
                 Some(Arc::new(FileContents::from_source(content))),

pyrefly/lib/query.rs

@@ -1592,7 +1592,7 @@ impl Query {
         // First, make sure that the types are well-formed and importable, return `Err` if not
         let code = format!("{imports}\n{snippet}\n");
         t.set_memory(vec![(
-            path.clone(),
+            ModulePath::memory(path.clone()),
             Some(Arc::new(FileContents::from_source(code))),
         )]);
         t.run(&[handle.dupe()], Require::Everything, None);

pyrefly/lib/state/errors.rs

@@ -1106,7 +1106,7 @@ mod tests {
         );
         let mut transaction = state.new_transaction(Require::Exports, None);
         transaction.set_memory(vec![(
-            get_path(&tdir),
+            ModulePath::memory(get_path(&tdir)),
             Some(Arc::new(FileContents::from_source(contents.to_owned()))),
         )]);
         transaction.run(&[handle.dupe()], Require::Everything, None);

pyrefly/lib/state/load.rs

@@ -12,7 +12,6 @@ use std::time::Instant;
 
 use anyhow::anyhow;
 use dupe::Dupe;
-use dupe::OptionDupedExt;
 use pyrefly_python::module::Module;
 use pyrefly_python::module_name::ModuleName;
 use pyrefly_python::module_path::ModulePath;
@@ -96,15 +95,13 @@ impl Load {
         memory_lookup: &MemoryFilesLookup,
         timing: Option<&TransactionTimingCounters>,
     ) -> (FileContents, Option<anyhow::Error>) {
+        if let Some(contents) = memory_lookup.get(path) {
+            return (contents.as_ref().dupe(), None);
+        }
         let res = match path.details() {
             ModulePathDetails::FileSystem(path) => Self::load_from_filesystem(path, timing),
             ModulePathDetails::Namespace(_) => Ok(FileContents::from_source("".to_owned())),
-            ModulePathDetails::Memory(path) => memory_lookup
-                .get(path)
-                .duped()
-                .as_deref()
-                .duped()
-                .ok_or_else(|| anyhow!("memory path not found")),
+            ModulePathDetails::Memory(_) => Err(anyhow!("memory path not found")),
             ModulePathDetails::BundledTypeshed(path) => typeshed().and_then(|x| {
                 x.load(path)
                     .ok_or_else(|| anyhow!("bundled typeshed problem"))

pyrefly/lib/state/memory.rs

@@ -7,17 +7,17 @@
 
 use std::collections::HashMap;
 use std::fmt::Debug;
-use std::path::Path;
-use std::path::PathBuf;
 use std::sync::Arc;
 
+use pyrefly_python::module_path::ModulePath;
+
 use crate::state::load::FileContents;
 
 #[derive(Debug, Clone, Default)]
-pub struct MemoryFiles(HashMap<PathBuf, Arc<FileContents>>);
+pub struct MemoryFiles(HashMap<ModulePath, Arc<FileContents>>);
 
 #[derive(Debug, Default, Clone)]
-pub struct MemoryFilesOverlay(HashMap<PathBuf, Option<Arc<FileContents>>>);
+pub struct MemoryFilesOverlay(HashMap<ModulePath, Option<Arc<FileContents>>>);
 
 impl MemoryFiles {
     pub fn apply_overlay(&mut self, overlay: MemoryFilesOverlay) {
@@ -31,7 +31,7 @@ impl MemoryFiles {
 }
 
 impl MemoryFilesOverlay {
-    pub fn set(&mut self, path: PathBuf, contents: Option<Arc<FileContents>>) {
+    pub fn set(&mut self, path: ModulePath, contents: Option<Arc<FileContents>>) {
         self.0.insert(path, contents);
     }
 }
@@ -46,7 +46,7 @@ impl<'a> MemoryFilesLookup<'a> {
         Self { base, overlay }
     }
 
-    pub fn get(&self, path: &Path) -> Option<&'a Arc<FileContents>> {
+    pub fn get(&self, path: &ModulePath) -> Option<&'a Arc<FileContents>> {
         match self.overlay.0.get(path) {
             Some(contents) => contents.as_ref(), // Might be a None if deleted
             None => self.base.0.get(path),

pyrefly/lib/state/state.rs

@@ -596,7 +596,7 @@ struct StateData {
     stdlib: SmallMap<SysInfo, Arc<Stdlib>>,
     modules: HashMap<Handle, ModuleData>,
     loaders: SmallMap<ArcId<ConfigFile>, Arc<LoaderFindCache>>,
-    /// The contents for ModulePath::memory values
+    /// In-memory contents that override each module path's usual source.
     memory: MemoryFiles,
     /// The current epoch, gets incremented every time we recompute
     now: Epoch,
@@ -1432,7 +1432,11 @@ impl<'a> Transaction<'a> {
                 if matches!(
                     module_data.handle.path().details(),
                     ModulePathDetails::FileSystem(_)
-                ) {
+                ) && self
+                    .memory_lookup()
+                    .get(module_data.handle.path())
+                    .is_none()
+                {
                     self.data.inputs_changed.store(true, Ordering::Relaxed);
                 }
                 guard.store_load(Some(Arc::new(Load::load_from_data(
@@ -2601,14 +2605,15 @@ impl<'a> Transaction<'a> {
         }
     }
 
-    /// Called if the `load_from_memory` portion of loading might have changed.
-    /// Specify which in-memory files might have changed, use None to say they don't exist anymore.
-    pub fn set_memory(&mut self, files: Vec<(PathBuf, Option<Arc<FileContents>>)>) {
+    /// Override sources by their exact module identity without changing import discovery.
+    /// None removes an override so that the module returns to its normal source.
+    /// Updating an override does not write or delete any filesystem source.
+    pub fn set_memory(&mut self, files: Vec<(ModulePath, Option<Arc<FileContents>>)>) {
         let mut changed = SmallSet::new();
         for (path, contents) in files {
             if self.memory_lookup().get(&path) != contents.as_ref() {
                 self.data.memory_overlay.set(path.clone(), contents);
-                changed.insert(ModulePath::memory(path));
+                changed.insert(path);
             }
         }
         self.stats.lock().set_memory_dirty = changed.len();

pyrefly/lib/test/incremental.rs

@@ -7,6 +7,7 @@
 
 //! Tests of the `State` object.
 
+use std::fs;
 use std::mem;
 use std::path::PathBuf;
 use std::sync::Arc;
@@ -22,6 +23,7 @@ use pyrefly_util::lock::Mutex;
 use pyrefly_util::prelude::SliceExt;
 use pyrefly_util::thread_pool::TEST_THREAD_COUNT;
 use starlark_map::small_map::SmallMap;
+use tempfile::TempDir;
 
 use crate::config::config::ConfigFile;
 use crate::config::error::ErrorDisplayConfig;
@@ -142,7 +144,7 @@ impl Incremental {
                 .lock()
                 .insert(ModuleName::from_str(&file), contents.dupe());
             transaction.as_mut().set_memory(vec![(
-                PathBuf::from(file),
+                ModulePath::memory(PathBuf::from(file)),
                 Some(Arc::new(FileContents::Source(contents))),
             )]);
         }
@@ -3191,3 +3193,114 @@ fn test_constant_defined_behind_skipped_star_import_rechecks_importer() {
         "{unsettled:?}"
     );
 }
+
+#[test]
+fn test_filesystem_overrides_preserve_incrementality_and_generation() {
+    init_test();
+    let directory = TempDir::new().unwrap();
+    let root = directory.path();
+    let foo = root.join("foo.py");
+    fs::write(&foo, "value: int = 1\n").unwrap();
+    fs::write(
+        root.join("main.py"),
+        "from foo import value\ncheck: int = value\n",
+    )
+    .unwrap();
+    let mut config = ConfigFile {
+        search_path_from_args: vec![root.to_owned()],
+        disable_search_path_heuristics: true,
+        ..Default::default()
+    };
+    config.python_environment.set_empty_to_default();
+    config.interpreters.skip_interpreter_query = true;
+    assert!(config.configure().is_empty());
+    let handles = ["main", "foo"].map(|name| {
+        Handle::new(
+            ModuleName::from_str(name),
+            ModulePath::filesystem(root.join(format!("{name}.py"))),
+            config.get_sys_info(),
+        )
+    });
+    let state = State::new(
+        ConfigFinder::new_constant(ArcId::new(config)),
+        TEST_THREAD_COUNT,
+    );
+    let check = |step: &str,
+                 memory: Option<Option<&str>>,
+                 disk: Option<&str>,
+                 rechecked: &[&str],
+                 errors: usize,
+                 generation_delta: u64| {
+        let generation = state.generation();
+        let subscriber = TestSubscriber::new();
+        let mut transaction =
+            state.new_committable_transaction(Require::Exports, Some(Box::new(subscriber.dupe())));
+        if let Some(contents) = memory {
+            transaction.as_mut().set_memory(vec![(
+                handles[1].path().dupe(),
+                contents.map(|contents| Arc::new(FileContents::from_source(contents.to_owned()))),
+            )]);
+        }
+        if let Some(contents) = disk {
+            fs::write(&foo, contents).unwrap();
+            transaction
+                .as_mut()
+                .invalidate_disk(std::slice::from_ref(&foo));
+        }
+        state.run_with_committing_transaction(
+            transaction,
+            std::slice::from_ref(&handles[0]),
+            Require::Everything,
+            None,
+            None,
+        );
+        let actual_errors = state
+            .transaction()
+            .get_errors([&handles[0]])
+            .collect_errors();
+        assert_eq!(
+            actual_errors.ordinary.len(),
+            errors,
+            "{step}: {actual_errors:?}"
+        );
+        let mut actual_rechecked = Vec::new();
+        for (handle, (count, _)) in subscriber.finish() {
+            if handles.contains(&handle) {
+                for _ in 0..count {
+                    actual_rechecked.push(handle.module().as_str().to_owned());
+                }
+            }
+        }
+        actual_rechecked.sort();
+        assert_eq!(actual_rechecked, rechecked, "{step}");
+        assert_eq!(state.generation(), generation + generation_delta, "{step}");
+    };
+
+    check("initial", None, None, &["foo", "main"], 0, 0);
+    check(
+        "interface change",
+        Some(Some("value: str = 'editor'\n")),
+        None,
+        &["foo", "main"],
+        1,
+        0,
+    );
+    check(
+        "comment-only change",
+        Some(Some("value: str = 'editor'\n# same interface\n")),
+        None,
+        &["foo"],
+        1,
+        0,
+    );
+    check(
+        "disk change under override",
+        None,
+        Some("value: int = 2\n"),
+        &[],
+        1,
+        0,
+    );
+    check("remove override", Some(None), None, &["foo", "main"], 0, 1);
+    check("remove absent override", Some(None), None, &[], 0, 0);
+}