fix(lsp): more accurate command name highlight/rename (#15540)

# Description

The `command` version of #15523 

# User-Facing Changes

Before:

<img width="394" alt="image"
src="https://github.com/user-attachments/assets/cdd1954d-c120-4aa4-8625-8a0f817ddebf"
/>

After:

<img width="431" alt="image"
src="https://github.com/user-attachments/assets/66fa17cd-2e6f-4305-a08a-df1c1617cfe8"
/>

And the renaming of that command finally works as expected.

Of course the identification of module prefixes in command calls is
still missing. I kinda feel there's no power-efficient way to do it.
I'll put low priority to that feature.

# Tests + Formatting

+1

# After Submitting

crates/nu-lsp/src/ast.rs

@@ -2,7 +2,7 @@ use crate::Id;
 use nu_protocol::{
     ast::{Argument, Block, Call, Expr, Expression, FindMapResult, ListItem, PathMember, Traverse},
     engine::StateWorkingSet,
-    ModuleId, Span,
+    DeclId, ModuleId, Span,
 };
 use std::sync::Arc;
 
@@ -24,6 +24,55 @@ fn strip_quotes(span: Span, working_set: &StateWorkingSet) -> (Box<[u8]>, Span)
     }
 }
 
+/// Trim leading `$` sign For variable references `$foo`
+fn strip_dollar_sign(span: Span, working_set: &StateWorkingSet<'_>) -> (Box<[u8]>, Span) {
+    let content = working_set.get_span_contents(span);
+    if content.starts_with(b"$") {
+        (
+            content[1..].into(),
+            Span::new(span.start.saturating_add(1), span.end),
+        )
+    } else {
+        (content.into(), span)
+    }
+}
+
+/// For a command call with head span content of `module name command    name`,
+/// return the span of `command    name`,
+/// while the actual command name is simply `command name`
+fn command_name_span_from_call_head(
+    working_set: &StateWorkingSet,
+    decl_id: DeclId,
+    head_span: Span,
+) -> Span {
+    let name = working_set.get_decl(decl_id).name();
+    let head_content = working_set.get_span_contents(head_span);
+    let mut head_words = head_content.split(|c| *c == b' ').collect::<Vec<_>>();
+    let mut name_words = name.split(' ').collect::<Vec<_>>();
+    let mut matched_len = name_words.len() - 1;
+    while let Some(name_word) = name_words.pop() {
+        while let Some(head_word) = head_words.pop() {
+            // for extra spaces, like those in the `command    name` example
+            if head_word.is_empty() && !name_word.is_empty() {
+                matched_len += 1;
+                continue;
+            }
+            if name_word.as_bytes() == head_word {
+                matched_len += head_word.len();
+                break;
+            } else {
+                // no such command name substring in head span
+                // probably an alias command, returning the whole head span
+                return head_span;
+            }
+        }
+        if name_words.len() > head_words.len() {
+            return head_span;
+        }
+    }
+    Span::new(head_span.end.saturating_sub(matched_len), head_span.end)
+}
+
 fn try_find_id_in_misc(
     call: &Call,
     working_set: &StateWorkingSet,
@@ -86,8 +135,8 @@ fn try_find_id_in_def(
         // for defs inside def
         // TODO: get scope by position
         // https://github.com/nushell/nushell/issues/15291
-        (0..working_set.num_decls()).find_map(|id| {
-            let decl_id = nu_protocol::DeclId::new(id);
+        (0..working_set.num_decls()).rev().find_map(|id| {
+            let decl_id = DeclId::new(id);
             let decl = working_set.get_decl(decl_id);
             let span = working_set.get_block(decl.block_id()?).span?;
             call.span().contains_span(span).then_some(decl_id)
@@ -139,7 +188,7 @@ fn try_find_id_in_mod(
                     }
                     let block_span = call.arguments.last()?.span();
                     (0..working_set.num_modules())
-                        .find(|id| {
+                        .rfind(|id| {
                             (any_id || id_num_ref == *id)
                                 && working_set.get_module(ModuleId::new(*id)).span.is_some_and(
                                     |mod_span| {
@@ -360,19 +409,6 @@ fn try_find_id_in_overlay(
     None
 }
 
-/// Trim leading `$` sign For variable references `$foo`
-fn strip_dollar_sign(span: Span, working_set: &StateWorkingSet<'_>) -> (Box<[u8]>, Span) {
-    let content = working_set.get_span_contents(span);
-    if content.starts_with(b"$") {
-        (
-            content[1..].into(),
-            Span::new(span.start.saturating_add(1), span.end),
-        )
-    } else {
-        (content.into(), span)
-    }
-}
-
 fn find_id_in_expr(
     expr: &Expression,
     working_set: &StateWorkingSet,
@@ -390,7 +426,8 @@ fn find_id_in_expr(
         }
         Expr::Call(call) => {
             if call.head.contains(*location) {
-                FindMapResult::Found((Id::Declaration(call.decl_id), call.head))
+                let span = command_name_span_from_call_head(working_set, call.decl_id, call.head);
+                FindMapResult::Found((Id::Declaration(call.decl_id), span))
             } else {
                 try_find_id_in_misc(call, working_set, Some(location), None)
                     .map(FindMapResult::Found)
@@ -482,7 +519,11 @@ fn find_reference_by_id_in_expr(
                 .flat_map(|e| e.flat_map(working_set, &closure))
                 .collect();
             if matches!(id, Id::Declaration(decl_id) if call.decl_id == *decl_id) {
-                occurs.push(call.head);
+                occurs.push(command_name_span_from_call_head(
+                    working_set,
+                    call.decl_id,
+                    call.head,
+                ));
                 return Some(occurs);
             }
             if let Some((_, span_found)) = try_find_id_in_misc(call, working_set, None, Some(id)) {

crates/nu-lsp/src/workspace.rs

@@ -70,6 +70,8 @@ impl IDTracker {
     fn new(id: Id, span: Span, file_span: Span, working_set: &StateWorkingSet) -> Self {
         let name = match &id {
             Id::Variable(_, name) | Id::Module(_, name) => name.clone(),
+            // NOTE: search by the canonical command name, some weird aliasing will be missing
+            Id::Declaration(decl_id) => working_set.get_decl(*decl_id).name().as_bytes().into(),
             _ => working_set.get_span_contents(span).into(),
         };
         Self {
@@ -899,6 +901,95 @@ mod tests {
         }
     }
 
+    #[test]
+    fn rename_module_command() {
+        let mut script = fixtures();
+        script.push("lsp");
+        script.push("workspace");
+        let (client_connection, _recv) = initialize_language_server(
+            None,
+            serde_json::to_value(InitializeParams {
+                workspace_folders: Some(vec![WorkspaceFolder {
+                    uri: path_to_uri(&script),
+                    name: "random name".to_string(),
+                }]),
+                ..Default::default()
+            })
+            .ok(),
+        );
+        script.push("baz.nu");
+        let script = path_to_uri(&script);
+
+        open_unchecked(&client_connection, script.clone());
+
+        let message_num = 6;
+        let messages = send_rename_prepare_request(
+            &client_connection,
+            script.clone(),
+            1,
+            47,
+            message_num,
+            false,
+        );
+        assert_eq!(messages.len(), message_num);
+        let mut has_response = false;
+        for message in messages {
+            match message {
+                Message::Notification(n) => assert_eq!(n.method, "$/progress"),
+                Message::Response(r) => {
+                    has_response = true;
+                    assert_json_eq!(
+                        r.result,
+                        serde_json::json!({
+                            "start": { "line": 1, "character": 41 },
+                            "end": { "line": 1, "character": 56 }
+                        }),
+                    )
+                }
+                _ => panic!("unexpected message type"),
+            }
+        }
+        assert!(has_response);
+
+        if let Message::Response(r) = send_rename_request(&client_connection, script.clone(), 6, 11)
+        {
+            let changes = r.result.unwrap()["changes"].clone();
+            assert_json_eq!(
+                changes[script.to_string().replace("baz", "foo")],
+                serde_json::json!([
+                    {
+                        "range": { "start": { "line": 10, "character": 16 }, "end": { "line": 10, "character": 29 } },
+                        "newText": "new"
+                    }
+                ])
+            );
+            let changs_baz = changes[script.to_string()].as_array().unwrap();
+            assert!(
+                changs_baz.contains(
+                &serde_json::json!({
+                    "range": { "start": { "line": 1, "character": 41 }, "end": { "line": 1, "character": 56 } },
+                    "newText": "new"
+                })
+            ));
+            assert!(
+                changs_baz.contains(
+                &serde_json::json!({
+                    "range": { "start": { "line": 2, "character": 0 }, "end": { "line": 2, "character": 5 } },
+                    "newText": "new"
+                })
+            ));
+            assert!(
+                changs_baz.contains(
+                &serde_json::json!({
+                    "range": { "start": { "line": 9, "character": 20 }, "end": { "line": 9, "character": 33 } },
+                    "newText": "new"
+                })
+            ));
+        } else {
+            panic!()
+        }
+    }
+
     #[test]
     fn rename_command_argument() {
         let mut script = fixtures();

tests/fixtures/lsp/workspace/baz.nu

@@ -1,5 +1,5 @@
 overlay use ./foo.nu as prefix --prefix
-alias aname = prefix mod name sub module cmd name
+alias aname = prefix mod name sub module cmd  name  long
 aname
 prefix foo str
 overlay hide prefix
@@ -7,6 +7,6 @@ overlay hide prefix
 use ./foo.nu [ "mod name" cst_mod ]
 
 $cst_mod."sub module"."sub sub module".var_name
-mod name sub module cmd name
+mod name sub module cmd name long
 let $cst_mod = 1
 $cst_mod

tests/fixtures/lsp/workspace/foo.nu

@@ -8,7 +8,7 @@ export  def "foo str" [] { "foo" }
 
 export module "mod name" {
   export module "sub module" {
-    export def "cmd name" [] { }
+    export def "cmd name long" [] { }
   }
 }