2 files changed, 74 insertions, 61 deletions

diff --git a/crates/base_db/src/lib.rs b/crates/base_db/src/lib.rs
index 030b96829..9733e1fd3 100644
--- a/crates/base_db/src/lib.rs
+++ b/crates/base_db/src/lib.rs
@@ -167,7 +167,7 @@ impl<T: SourceDatabaseExt> FileLoader for FileLoaderDelegate<&'_ T> {
     }
 
     fn possible_sudmobule_names(&self, module_file: FileId) -> Vec<String> {
-        self.source_root(module_file).file_set.possible_sudmobule_names(module_file)
+        possible_sudmobule_names(&self.source_root(module_file).file_set, module_file)
     }
 }
 
@@ -177,3 +177,71 @@ impl<T: SourceDatabaseExt> FileLoaderDelegate<&'_ T> {
         self.0.source_root(source_root)
     }
 }
+
+fn possible_sudmobule_names(module_files: &FileSet, module_file: FileId) -> Vec<String> {
+    let directory_to_look_for_submodules = match module_files
+        .path_for_file(&module_file)
+        .and_then(|module_file_path| get_directory_with_submodules(module_file_path))
+    {
+        Some(directory) => directory,
+        None => return Vec::new(),
+    };
+    module_files
+        .iter()
+        .filter(|submodule_file| submodule_file != &module_file)
+        .filter_map(|submodule_file| {
+            let submodule_path = module_files.path_for_file(&submodule_file)?;
+            if submodule_path.parent()? == directory_to_look_for_submodules {
+                submodule_path.file_name_and_extension()
+            } else {
+                None
+            }
+        })
+        .filter_map(|file_name_and_extension| {
+            match file_name_and_extension {
+                // TODO kb wrong resolution for nested non-file modules (mod tests {mod <|>)
+                // TODO kb in src/bin when a module is included into another,
+                // the included file gets "moved" into a directory below and now cannot add any other modules
+                ("mod", Some("rs")) | ("lib", Some("rs")) | ("main", Some("rs")) => None,
+                (file_name, Some("rs")) => Some(file_name.to_owned()),
+                (subdirectory_name, None) => {
+                    let mod_rs_path =
+                        directory_to_look_for_submodules.join(subdirectory_name)?.join("mod.rs")?;
+                    if module_files.file_for_path(&mod_rs_path).is_some() {
+                        Some(subdirectory_name.to_owned())
+                    } else {
+                        None
+                    }
+                }
+                _ => None,
+            }
+        })
+        .collect()
+}
+
+fn get_directory_with_submodules(module_file_path: &VfsPath) -> Option<VfsPath> {
+    let module_directory_path = module_file_path.parent()?;
+    match module_file_path.file_name_and_extension()? {
+        ("mod", Some("rs")) | ("lib", Some("rs")) | ("main", Some("rs")) => {
+            Some(module_directory_path)
+        }
+        (regular_rust_file_name, Some("rs")) => {
+            if matches!(
+                (
+                    module_directory_path
+                        .parent()
+                        .as_ref()
+                        .and_then(|path| path.file_name_and_extension()),
+                    module_directory_path.file_name_and_extension(),
+                ),
+                (Some(("src", None)), Some(("bin", None)))
+            ) {
+                // files in /src/bin/ can import each other directly
+                Some(module_directory_path)
+            } else {
+                module_directory_path.join(regular_rust_file_name)
+            }
+        }
+        _ => None,
+    }
+}
diff --git a/crates/vfs/src/file_set.rs b/crates/vfs/src/file_set.rs
index cb65c17e0..4aa2d6526 100644
--- a/crates/vfs/src/file_set.rs
+++ b/crates/vfs/src/file_set.rs
@@ -26,74 +26,19 @@ impl FileSet {
         self.files.get(&path).copied()
     }
 
-    pub fn possible_sudmobule_names(&self, module_file: FileId) -> Vec<String> {
-        let directory_to_look_for_submodules = match self.get_directory_with_submodules(module_file)
-        {
-            Some(directory) => directory,
-            None => return Vec::new(),
-        };
-        self.paths
-            .iter()
-            .filter_map(|(_, path)| {
-                if path.parent()? == directory_to_look_for_submodules {
-                    path.file_name_and_extension()
-                } else {
-                    None
-                }
-            })
-            .filter_map(|file_name_and_extension| match file_name_and_extension {
-                // TODO kb do not include the module file name itself, if present
-                // TODO kb wrong resolution for nesСпаted non-file modules (mod tests {mod <|>)
-                // TODO kb in src/bin when a module is included into another,
-                // the included file gets "moved" into a directory below and now cannot add any other modules
-                ("mod", Some("rs")) | ("lib", Some("rs")) | ("main", Some("rs")) => None,
-                (file_name, Some("rs")) => Some(file_name.to_owned()),
-                (subdirectory_name, None) => {
-                    let mod_rs_path =
-                        directory_to_look_for_submodules.join(subdirectory_name)?.join("mod.rs")?;
-                    if self.files.contains_key(&mod_rs_path) {
-                        Some(subdirectory_name.to_owned())
-                    } else {
-                        None
-                    }
-                }
-                _ => None,
-            })
-            .collect()
+    pub fn file_for_path(&self, path: &VfsPath) -> Option<&FileId> {
+        self.files.get(path)
     }
 
-    fn get_directory_with_submodules(&self, module_file: FileId) -> Option<VfsPath> {
-        let module_file_path = &self.paths[&module_file];
-        let module_directory_path = module_file_path.parent()?;
-        match module_file_path.file_name_and_extension() {
-            Some(("mod", Some("rs"))) | Some(("lib", Some("rs"))) | Some(("main", Some("rs"))) => {
-                Some(module_directory_path)
-            }
-            Some((regular_rust_file_name, Some("rs"))) => {
-                if matches!(
-                    (
-                        module_directory_path
-                            .parent()
-                            .as_ref()
-                            .and_then(|path| path.file_name_and_extension()),
-                        module_directory_path.file_name_and_extension(),
-                    ),
-                    (Some(("src", None)), Some(("bin", None)))
-                ) {
-                    // files in /src/bin/ can import each other directly
-                    Some(module_directory_path)
-                } else {
-                    module_directory_path.join(regular_rust_file_name)
-                }
-            }
-            _ => None,
-        }
+    pub fn path_for_file(&self, file: &FileId) -> Option<&VfsPath> {
+        self.paths.get(file)
     }
 
     pub fn insert(&mut self, file_id: FileId, path: VfsPath) {
         self.files.insert(path.clone(), file_id);
         self.paths.insert(file_id, path);
     }
+
     pub fn iter(&self) -> impl Iterator<Item = FileId> + '_ {
         self.paths.keys().copied()
     }