Fix runnables code lens in multi-file sway projects

sway-lsp/benches/lsp_benchmarks/requests.rs

@@ -1,7 +1,7 @@
 use criterion::{black_box, criterion_group, Criterion};
 use lsp_types::{
-    CodeLens, CompletionResponse, DocumentSymbolResponse, Position, Range,
-    TextDocumentContentChangeEvent, TextDocumentIdentifier,
+    CompletionResponse, DocumentSymbolResponse, Position, Range, TextDocumentContentChangeEvent,
+    TextDocumentIdentifier,
 };
 use sway_lsp::{capabilities, lsp_ext::OnEnterParams, utils::keyword_docs::KeywordDocs};
 
@@ -79,17 +79,7 @@ fn benchmarks(c: &mut Criterion) {
     });
 
     c.bench_function("code_lens", |b| {
-        b.iter(|| {
-            let mut result = vec![];
-            session.runnables.iter().for_each(|item| {
-                let runnable = item.value();
-                result.push(CodeLens {
-                    range: runnable.range(),
-                    command: Some(runnable.command()),
-                    data: None,
-                });
-            });
-        })
+        b.iter(|| capabilities::code_lens::code_lens(&session, &uri.clone()))
     });
 
     c.bench_function("on_enter", |b| {

sway-lsp/src/capabilities/code_lens.rs

@@ -0,0 +1,27 @@
+use std::{path::PathBuf, sync::Arc};
+
+use lsp_types::{CodeLens, Url};
+
+use crate::core::session::Session;
+
+pub fn code_lens(session: &Arc<Session>, url: &Url) -> Vec<CodeLens> {
+    let url_path = PathBuf::from(url.path());
+
+    // Construct code lenses for runnable functions
+    let runnables_for_path = session.runnables.get(&url_path);
+    let mut result: Vec<CodeLens> = runnables_for_path
+        .map(|runnables| {
+            runnables
+                .iter()
+                .map(|runnable| CodeLens {
+                    range: *runnable.range(),
+                    command: Some(runnable.command()),
+                    data: None,
+                })
+                .collect()
+        })
+        .unwrap_or_default();
+    // Sort the results
+    result.sort_by(|a, b| a.range.start.line.cmp(&b.range.start.line));
+    result
+}

sway-lsp/src/capabilities/mod.rs

@@ -1,4 +1,5 @@
 pub mod code_actions;
+pub mod code_lens;
 pub mod completion;
 pub mod diagnostic;
 pub mod document_symbol;

sway-lsp/src/capabilities/runnable.rs

@@ -1,21 +1,19 @@
-use crate::core::token::get_range_from_span;
 use lsp_types::{Command, Range};
 use serde_json::{json, Value};
 use sway_core::language::parsed::TreeType;
-use sway_types::Span;
 
 #[derive(Debug, Eq, PartialEq, Clone)]
 pub struct RunnableMainFn {
     /// The location in the file where the runnable button should be displayed
-    pub span: Span,
+    pub range: Range,
     /// The program kind of the current file
     pub tree_type: TreeType,
 }
 
 #[derive(Debug, Eq, PartialEq, Clone)]
 pub struct RunnableTestFn {
-    /// The location in the file where the runnable button should be displayed.
-    pub span: Span,
+    /// The location in the file where the runnable button should be displayed
+    pub range: Range,
     /// The program kind of the current file.
     pub tree_type: TreeType,
     /// Additional arguments to use with the runnable command.
@@ -38,12 +36,8 @@ pub trait Runnable: core::fmt::Debug + Send + Sync + 'static {
     fn label_string(&self) -> String;
     /// The arguments to pass to the command.
     fn arguments(&self) -> Option<Vec<Value>>;
-    /// The span where the runnable button should be displayed.
-    fn span(&self) -> &Span;
     /// The range in the file where the runnable button should be displayed.
-    fn range(&self) -> Range {
-        get_range_from_span(self.span())
-    }
+    fn range(&self) -> &Range;
 }
 
 impl Runnable for RunnableMainFn {
@@ -56,8 +50,8 @@ impl Runnable for RunnableMainFn {
     fn arguments(&self) -> Option<Vec<Value>> {
         None
     }
-    fn span(&self) -> &Span {
-        &self.span
+    fn range(&self) -> &Range {
+        &self.range
     }
 }
 
@@ -73,7 +67,7 @@ impl Runnable for RunnableTestFn {
             .as_ref()
             .map(|test_name| vec![json!({ "name": test_name })])
     }
-    fn span(&self) -> &Span {
-        &self.span
+    fn range(&self) -> &Range {
+        &self.range
     }
 }

sway-lsp/src/core/session.rs

@@ -41,10 +41,12 @@ use sway_core::{
     },
     BuildTarget, Engines, Namespace, Programs,
 };
-use sway_types::{Span, Spanned};
+use sway_types::{SourceEngine, Spanned};
 use sway_utils::helpers::get_sway_files;
 use tokio::sync::Semaphore;
 
+use super::token::get_range_from_span;
+
 pub type Documents = DashMap<String, TextDocument>;
 pub type ProjectDirectory = PathBuf;
 
@@ -75,7 +77,7 @@ pub struct ParseResult {
 pub struct Session {
     token_map: TokenMap,
     pub documents: Documents,
-    pub runnables: DashMap<Span, Box<dyn Runnable>>,
+    pub runnables: DashMap<PathBuf, Vec<Box<dyn Runnable>>>,
     pub compiled_program: RwLock<CompiledProgram>,
     pub engines: RwLock<Engines>,
     pub sync: SyncWorkspace,
@@ -154,7 +156,11 @@ impl Session {
             self.token_map.insert(s.clone(), t.clone());
         });
 
-        self.create_runnables(&res.typed, self.engines.read().de());
+        self.create_runnables(
+            &res.typed,
+            self.engines.read().de(),
+            self.engines.read().se(),
+        );
         self.compiled_program.write().lexed = Some(res.lexed);
         self.compiled_program.write().parsed = Some(res.parsed);
         self.compiled_program.write().typed = Some(res.typed);
@@ -332,20 +338,31 @@ impl Session {
     }
 
     /// Create runnables if the `TyProgramKind` of the `TyProgram` is a script.
-    fn create_runnables(&self, typed_program: &ty::TyProgram, decl_engine: &DeclEngine) {
+    fn create_runnables(
+        &self,
+        typed_program: &ty::TyProgram,
+        decl_engine: &DeclEngine,
+        source_engine: &SourceEngine,
+    ) {
         // Insert runnable test functions.
         for (decl, _) in typed_program.test_fns(decl_engine) {
             // Get the span of the first attribute if it exists, otherwise use the span of the function name.
             let span = decl
                 .attributes
                 .first()
                 .map_or_else(|| decl.name.span(), |(_, attr)| attr.span.clone());
-            let runnable = Box::new(RunnableTestFn {
-                span,
-                tree_type: typed_program.kind.tree_type(),
-                test_name: Some(decl.name.to_string()),
-            });
-            self.runnables.insert(runnable.span().clone(), runnable);
+            if let Some(source_id) = span.source_id() {
+                let path = source_engine.get_path(source_id);
+                let runnable = Box::new(RunnableTestFn {
+                    range: get_range_from_span(&span.clone()),
+                    tree_type: typed_program.kind.tree_type(),
+                    test_name: Some(decl.name.to_string()),
+                });
+                self.runnables
+                    .entry(path)
+                    .or_insert(Vec::new())
+                    .push(runnable);
+            }
         }
 
         // Insert runnable main function if the program is a script.
@@ -354,11 +371,17 @@ impl Session {
         } = typed_program.kind
         {
             let span = main_function.name.span();
-            let runnable = Box::new(RunnableMainFn {
-                span,
-                tree_type: typed_program.kind.tree_type(),
-            });
-            self.runnables.insert(runnable.span().clone(), runnable);
+            if let Some(source_id) = span.source_id() {
+                let path = source_engine.get_path(source_id);
+                let runnable = Box::new(RunnableMainFn {
+                    range: get_range_from_span(&span.clone()),
+                    tree_type: typed_program.kind.tree_type(),
+                });
+                self.runnables
+                    .entry(path)
+                    .or_insert(Vec::new())
+                    .push(runnable);
+            }
         }
     }
 

sway-lsp/src/handlers/request.rs

@@ -240,25 +240,11 @@ pub fn handle_code_lens(
     state: &ServerState,
     params: lsp_types::CodeLensParams,
 ) -> Result<Option<Vec<CodeLens>>> {
-    let mut result = vec![];
     match state
         .sessions
         .uri_and_session_from_workspace(&params.text_document.uri)
     {
-        Ok((_, session)) => {
-            // Construct code lenses for runnable functions
-            session.runnables.iter().for_each(|item| {
-                let runnable = item.value();
-                result.push(CodeLens {
-                    range: runnable.range(),
-                    command: Some(runnable.command()),
-                    data: None,
-                });
-            });
-            // Sort the results
-            result.sort_by(|a, b| a.range.start.line.cmp(&b.range.start.line));
-            Ok(Some(result))
-        }
+        Ok((url, session)) => Ok(Some(capabilities::code_lens::code_lens(&session, &url))),
         Err(err) => {
             tracing::error!("{}", err.to_string());
             Ok(None)

sway-lsp/tests/fixtures/runnables/Forc.toml

@@ -0,0 +1,8 @@
+[project]
+authors = ["Fuel Labs <[email protected]>"]
+entry = "main.sw"
+license = "Apache-2.0"
+name = "script_multi_test"
+
+[dependencies]
+std = { path = "../../../../sway-lib-std" }

sway-lsp/tests/fixtures/runnables/src/main.sw

@@ -0,0 +1,16 @@
+script;
+
+fn main() {
+    revert(0);
+}
+
+#[test]
+fn test_foo() {
+    assert(true);
+}
+
+#[test]
+fn test_bar() {
+    log("test");
+    assert(4 / 2 == 2);
+}

sway-lsp/tests/fixtures/runnables/src/other.sw

@@ -0,0 +1 @@
+library;

sway-lsp/tests/integration/lsp.rs

@@ -224,6 +224,16 @@ pub(crate) fn highlight_request(server: &ServerState, uri: &Url) {
     assert_eq!(expected, response.unwrap());
 }
 
+pub(crate) fn code_lens_empty_request(server: &ServerState, uri: &Url) {
+    let params = CodeLensParams {
+        text_document: TextDocumentIdentifier { uri: uri.clone() },
+        work_done_progress_params: Default::default(),
+        partial_result_params: Default::default(),
+    };
+    let response = request::handle_code_lens(server, params.clone()).unwrap();
+    assert_eq!(response.unwrap().len(), 0);
+}
+
 pub(crate) fn code_lens_request(server: &ServerState, uri: &Url) {
     let params = CodeLensParams {
         text_document: TextDocumentIdentifier { uri: uri.clone() },

sway-lsp/tests/lib.rs

@@ -1723,6 +1723,11 @@ lsp_capability_test!(
     lsp::code_lens_request,
     runnables_test_dir().join("src/main.sw")
 );
+lsp_capability_test!(
+    code_lens_empty,
+    lsp::code_lens_empty_request,
+    runnables_test_dir().join("src/other.sw")
+);
 lsp_capability_test!(
     completion,
     lsp::completion_request,

sway-lsp/tests/utils/src/lib.rs

@@ -38,9 +38,7 @@ pub fn e2e_test_dir() -> PathBuf {
 }
 
 pub fn runnables_test_dir() -> PathBuf {
-    sway_workspace_dir()
-        .join(e2e_unit_dir())
-        .join("script_multi_test")
+    test_fixtures_dir().join("runnables")
 }
 
 pub fn test_fixtures_dir() -> PathBuf {