Use ModuleId instead of SourceId as the metrics map key (#5871)

## Description
Fixes a critical bug where the server would crash if compilation was
triggered from anything other than the root `main.sw` file. This was due
to the metrics_map using the `SourceId` rather then the `ModuleId` as
the key.

This was introduced last week in #5813 as we now need to metrics to
decide if the engines should be mem swapped or discarded.

closes #5870 

## Checklist

- [x] I have linked to any relevant issues.
- [x] I have commented my code, particularly in hard-to-understand
areas.
- [x] I have updated the documentation where relevant (API docs, the
reference, and the Sway book).
- [x] If my change requires substantial documentation changes, I have
[requested support from the DevRel
team](https://github.com/FuelLabs/devrel-requests/issues/new/choose)
- [x] I have added tests that prove my fix is effective or that my
feature works.
- [x] I have added (or requested a maintainer to add) the necessary
`Breaking*` or `New Feature` labels where relevant.
- [x] I have done my best to ensure that my PR adheres to [the Fuel Labs
Code Review
Standards](https://github.com/FuelLabs/rfcs/blob/master/text/code-standards/external-contributors.md).
- [x] I have requested a review from the relevant team or maintainers.

sway-lsp/src/core/session.rs

@@ -10,7 +10,7 @@ use crate::{
         token::{self, TypedAstToken},
         token_map::{TokenMap, TokenMapExt},
     },
-    error::{DocumentError, LanguageServerError},
+    error::{DirectoryError, DocumentError, LanguageServerError},
     traverse::{
         dependency, lexed_tree, parsed_tree::ParsedTree, typed_tree::TypedTree, ParseContext,
     },
@@ -40,7 +40,7 @@ use sway_core::{
     BuildTarget, Engines, LspConfig, Namespace, Programs,
 };
 use sway_error::{error::CompileError, handler::Handler, warning::CompileWarning};
-use sway_types::{SourceEngine, SourceId, Spanned};
+use sway_types::{ModuleId, SourceEngine, Spanned};
 use sway_utils::{helpers::get_sway_files, PerformanceData};
 
 pub type RunnableMap = DashMap<PathBuf, Vec<Box<dyn Runnable>>>;
@@ -66,7 +66,7 @@ pub struct Session {
     pub sync: SyncWorkspace,
     // Cached diagnostic results that require a lock to access. Readers will wait for writers to complete.
     pub diagnostics: Arc<RwLock<DiagnosticMap>>,
-    pub metrics: DashMap<SourceId, PerformanceData>,
+    pub metrics: DashMap<ModuleId, PerformanceData>,
 }
 
 impl Default for Session {
@@ -294,23 +294,27 @@ pub fn traverse(
             metrics,
         } = value.unwrap();
 
-        let source_id = lexed.root.tree.span().source_id().cloned();
-        if let Some(source_id) = source_id {
-            session.metrics.insert(source_id, metrics.clone());
-        }
-
-        let engines_ref = session.engines.read();
         // Check if the cached AST was returned by the compiler for the users workspace.
         // If it was, then we need to use the original engines for traversal.
         //
         // This is due to the garbage collector removing types from the engines_clone
         // and they have not been re-added due to compilation being skipped.
+        let engines_ref = session.engines.read();
         let engines = if i == results_len - 1 && metrics.reused_modules > 0 {
             &*engines_ref
         } else {
             engines_clone
         };
 
+        // Convert the source_id to a path so we can use the manifest path to get the module_id.
+        // This is used to store the metrics for the module.
+        let source_id = lexed.root.tree.span().source_id().cloned();
+        if let Some(source_id) = source_id {
+            let path = engines.se().get_path(&source_id);
+            let module_id = module_id_from_path(&path, engines)?;
+            session.metrics.insert(module_id, metrics);
+        }
+
         // Get a reference to the typed program AST.
         let typed_program = typed
             .as_ref()
@@ -357,7 +361,6 @@ pub fn traverse(
             });
         }
     }
-
     Ok(Some(diagnostics))
 }
 
@@ -482,6 +485,19 @@ fn create_runnables(
     }
 }
 
+/// Resolves a `ModuleId` from a given `path` using the manifest directory.
+pub(crate) fn module_id_from_path(
+    path: &PathBuf,
+    engines: &Engines,
+) -> Result<ModuleId, DirectoryError> {
+    let module_id = sway_utils::find_parent_manifest_dir(path)
+        .and_then(|manifest_path| engines.se().get_module_id(&manifest_path))
+        .ok_or_else(|| DirectoryError::ModuleIdNotFound {
+            path: path.to_string_lossy().to_string(),
+        })?;
+    Ok(module_id)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

sway-lsp/src/error.rs

@@ -72,6 +72,8 @@ pub enum DirectoryError {
     PathFromUrlFailed { url: String },
     #[error("Unable to create span from path {:?}", path)]
     SpanFromPathFailed { path: String },
+    #[error("No module ID found for path {:?}", path)]
+    ModuleIdNotFound { path: String },
 }
 
 #[derive(Debug, Error, PartialEq, Eq)]

sway-lsp/src/handlers/request.rs

@@ -502,7 +502,8 @@ pub(crate) async fn metrics(
                     .engines
                     .read()
                     .se()
-                    .get_path(kv.key())
+                    .get_path_from_module_id(kv.key())
+                    .unwrap()
                     .to_string_lossy()
                     .to_string();
                 metrics.push((path, kv.value().clone()));

sway-lsp/src/server_state.rs

@@ -158,27 +158,33 @@ impl ServerState {
                             experimental,
                         ) {
                             Ok(_) => {
-                                if let Ok(path) = uri.to_file_path() {
-                                    let path = Arc::new(path);
-                                    let source_id =
-                                        session.engines.read().se().get_source_id(&path);
-                                    let metrics = session
-                                        .metrics
-                                        .get(&source_id)
-                                        .expect("metrics not found for source_id");
-                                    // It's very important to check if the workspace AST was reused to determine if we need to overwrite the engines.
-                                    // Because the engines_clone has garbage collection applied. If the workspace AST was reused, we need to keep the old engines
-                                    // as the engines_clone might have cleared some types that are still in use.
-                                    if metrics.reused_modules == 0 {
-                                        // The compiler did not reuse the workspace AST.
-                                        // We need to overwrite the old engines with the engines clone.
-                                        mem::swap(
-                                            &mut *session.engines.write(),
-                                            &mut engines_clone,
-                                        );
+                                let path = uri.to_file_path().unwrap();
+                                // Find the module id from the path
+                                match session::module_id_from_path(&path, &engines_clone) {
+                                    Ok(module_id) => {
+                                        // Use the module id to get the metrics for the module
+                                        if let Some(metrics) = session.metrics.get(&module_id) {
+                                            // It's very important to check if the workspace AST was reused to determine if we need to overwrite the engines.
+                                            // Because the engines_clone has garbage collection applied. If the workspace AST was reused, we need to keep the old engines
+                                            // as the engines_clone might have cleared some types that are still in use.
+                                            if metrics.reused_modules == 0 {
+                                                // The compiler did not reuse the workspace AST.
+                                                // We need to overwrite the old engines with the engines clone.
+                                                mem::swap(
+                                                    &mut *session.engines.write(),
+                                                    &mut engines_clone,
+                                                );
+                                            }
+                                        }
+                                        *last_compilation_state.write() =
+                                            LastCompilationState::Success;
+                                    }
+                                    Err(err) => {
+                                        tracing::error!("{}", err.to_string());
+                                        *last_compilation_state.write() =
+                                            LastCompilationState::Failed;
                                     }
                                 }
-                                *last_compilation_state.write() = LastCompilationState::Success;
                             }
                             Err(_err) => {
                                 *last_compilation_state.write() = LastCompilationState::Failed;

sway-lsp/tests/lib.rs

@@ -345,6 +345,19 @@ fn lsp_syncs_with_workspace_edits() {
     });
 }
 
+#[test]
+fn compilation_succeeds_when_triggered_from_module() {
+    run_async!({
+        let server = ServerState::default();
+        let _ = open(
+            &server,
+            test_fixtures_dir().join("tokens/modules/src/test_mod.sw"),
+        )
+        .await;
+        let _ = server.shutdown_server().await;
+    });
+}
+
 #[test]
 fn show_ast() {
     run_async!({

sway-types/src/source_engine.rs

@@ -108,6 +108,15 @@ impl SourceEngine {
         self.path_to_module_map.read().unwrap().get(path).cloned()
     }
 
+    /// Returns the [PathBuf] associated with the provided [ModuleId], if it exists in the path_to_module_map.
+    pub fn get_path_from_module_id(&self, module_id: &ModuleId) -> Option<PathBuf> {
+        let path_to_module_map = self.path_to_module_map.read().unwrap();
+        path_to_module_map
+            .iter()
+            .find(|(_, &id)| id == *module_id)
+            .map(|(path, _)| path.clone())
+    }
+
     /// This function provides the file name (with extension) corresponding to a specified source ID.
     pub fn get_file_name(&self, source_id: &SourceId) -> Option<String> {
         self.get_path(source_id)