state: key in-memory sources by ModulePath [NFC]

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

@@ -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

@@ -99,7 +99,7 @@ impl Load {
         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
+            ModulePathDetails::Memory(_) => memory_lookup
                 .get(path)
                 .duped()
                 .as_deref()

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

@@ -2603,12 +2603,12 @@ 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>>)>) {
+    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