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()),
+            ModulePath::memory(module_path.as_path().to_path_buf()),
             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

@@ -3291,7 +3291,12 @@ impl Server {
             self.open_files
                 .read()
                 .iter()
-                .map(|x| (x.0.clone(), Some(Arc::new(x.1.to_file_contents()))))
+                .map(|x| {
+                    (
+                        ModulePath::memory(x.0.clone()),
+                        Some(Arc::new(x.1.to_file_contents())),
+                    )
+                })
                 .collect::<Vec<_>>(),
         );
         transaction.run(&handles, Require::Everything, custom_thread_pool);
@@ -4802,7 +4807,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,10 +46,127 @@ 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),
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::fs;
+    use std::path::PathBuf;
+
+    use tempfile::TempDir;
+
+    use super::*;
+    use crate::state::load::Load;
+
+    #[test]
+    fn lookup_keeps_module_identities_and_pending_removals_separate() {
+        let path = PathBuf::from("same.py");
+        let memory = ModulePath::memory(path.clone());
+        let filesystem = ModulePath::filesystem(path);
+        let memory_contents = Arc::new(FileContents::from_source("memory".to_owned()));
+        let filesystem_contents = Arc::new(FileContents::from_source("filesystem".to_owned()));
+        let mut base = MemoryFiles::default();
+        let mut initial = MemoryFilesOverlay::default();
+        initial.set(memory.clone(), Some(Arc::clone(&memory_contents)));
+        initial.set(filesystem.clone(), Some(Arc::clone(&filesystem_contents)));
+        base.apply_overlay(initial);
+
+        let mut overlay = MemoryFilesOverlay::default();
+        let lookup = MemoryFilesLookup::new(&base, &overlay);
+        assert_eq!(lookup.get(&memory), Some(&memory_contents));
+        assert_eq!(lookup.get(&filesystem), Some(&filesystem_contents));
+        overlay.set(filesystem.clone(), None);
+        let lookup = MemoryFilesLookup::new(&base, &overlay);
+        assert!(lookup.get(&filesystem).is_none());
+        assert_eq!(lookup.get(&memory), Some(&memory_contents));
+
+        base.apply_overlay(overlay);
+        let empty = MemoryFilesOverlay::default();
+        let lookup = MemoryFilesLookup::new(&base, &empty);
+        assert!(lookup.get(&filesystem).is_none());
+        assert_eq!(lookup.get(&memory), Some(&memory_contents));
+    }
+
+    #[test]
+    fn exact_overrides_and_removal_preserve_each_variants_normal_loader() {
+        let directory = TempDir::new().unwrap();
+        let file = directory.path().join("source.py");
+        fs::write(&file, "saved\n").unwrap();
+        let paths = [
+            ModulePath::filesystem(file.clone()),
+            ModulePath::memory(file.clone()),
+            ModulePath::namespace(file),
+            ModulePath::bundled_typeshed(PathBuf::from("builtins.pyi")),
+            ModulePath::bundled_typeshed_third_party(PathBuf::from("__overlay_probe__.pyi")),
+            ModulePath::bundled_third_party(PathBuf::from("__overlay_probe__.pyi")),
+        ];
+        let mut base = MemoryFiles::default();
+        let empty = MemoryFilesOverlay::default();
+        let normal: Vec<_> = paths
+            .iter()
+            .map(|path| {
+                let (contents, error) =
+                    Load::load_from_path(path, &MemoryFilesLookup::new(&base, &empty), None);
+                (contents, error.map(|error| error.to_string()))
+            })
+            .collect();
+        let mut initial = MemoryFilesOverlay::default();
+        for (index, path) in paths.iter().enumerate() {
+            initial.set(
+                path.clone(),
+                Some(Arc::new(FileContents::from_source(format!(
+                    "override {index}\n"
+                )))),
+            );
+        }
+        base.apply_overlay(initial);
+        for (index, path) in paths.iter().enumerate() {
+            let (contents, error) =
+                Load::load_from_path(path, &MemoryFilesLookup::new(&base, &empty), None);
+            assert!(error.is_none(), "{path:?}: {error:?}");
+            assert_eq!(
+                contents,
+                FileContents::from_source(format!("override {index}\n"))
+            );
+            let mut removed = MemoryFilesOverlay::default();
+            removed.set(path.clone(), None);
+            let (contents, error) =
+                Load::load_from_path(path, &MemoryFilesLookup::new(&base, &removed), None);
+            assert_eq!(
+                (contents, error.map(|error| error.to_string())),
+                normal[index]
+            );
+        }
+    }
+
+    #[test]
+    fn removing_an_override_exposes_missing_filesystem_source() {
+        let directory = TempDir::new().unwrap();
+        let path = ModulePath::filesystem(directory.path().join("missing.py"));
+        let mut base = MemoryFiles::default();
+        let mut initial = MemoryFilesOverlay::default();
+        initial.set(
+            path.clone(),
+            Some(Arc::new(FileContents::from_source("buffer\n".to_owned()))),
+        );
+        base.apply_overlay(initial);
+        let mut pending = MemoryFilesOverlay::default();
+        let (contents, error) =
+            Load::load_from_path(&path, &MemoryFilesLookup::new(&base, &pending), None);
+        assert!(error.is_none());
+        assert_eq!(contents, FileContents::from_source("buffer\n".to_owned()));
+        pending.set(path.clone(), None);
+        let (_, error) =
+            Load::load_from_path(&path, &MemoryFilesLookup::new(&base, &pending), None);
+        assert!(
+            error.is_some(),
+            "Removing an override must expose the missing saved source"
+        );
+    }
+}

pyrefly/lib/state/state.rs

@@ -2587,14 +2587,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

@@ -142,7 +142,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))),
             )]);
         }

pyrefly/lib/test/lsp/definition.rs

@@ -533,7 +533,7 @@ fn state_with_conftest(files: &[(&str, &str)], code: &str) -> (TempDir, State, H
     );
     let mut transaction = state.new_committable_transaction(Require::Exports, None);
     transaction.as_mut().set_memory(vec![(
-        path,
+        ModulePath::memory(path),
         Some(Arc::new(FileContents::from_source(code.to_owned()))),
     )]);
     transaction
@@ -1388,7 +1388,7 @@ fn parent_declarations_keep_the_declaring_snapshot() {
     );
     let mut transaction = state.new_committable_transaction(Require::Everything, None);
     transaction.as_mut().set_memory(vec![(
-        path,
+        ModulePath::memory(path),
         Some(Arc::new(FileContents::from_source(memory))),
     )]);
     transaction

pyrefly/lib/test/map_int_tuples_binding.rs

@@ -135,7 +135,7 @@ fn map_int_tuples_mapper_parameter_is_order_independent_and_incremental() {
 
     let mut transaction = state.new_committable_transaction(Require::Everything, None);
     transaction.as_mut().set_memory(vec![(
-        PathBuf::from("mapmod.py"),
+        ModulePath::memory(PathBuf::from("mapmod.py")),
         Some(Arc::new(FileContents::from_source(map_module("Item")))),
     )]);
     transaction