feat(lsp): document highlight (#14898)

<!--
if this PR closes one or more issues, you can automatically link the PR
with
them by using one of the [*linking
keywords*](https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword),
e.g.
- this PR should close #xxxx
- fixes #xxxx

you can also mention related issues, PRs or discussions!
-->

# Description
<!--
Thank you for improving Nushell. Please, check our [contributing
guide](../CONTRIBUTING.md) and talk to the core team before making major
changes.

Description of your pull request goes here. **Provide examples and/or
screenshots** if your changes affect the user experience.
-->

This is a minor feature that highlights all occurrences of current
variable/command in current file:

<img width="346" alt="image"
src="https://github.com/user-attachments/assets/f1078e79-d02e-480e-b84a-84efb222c9a4"
/>

Since this kind of request happens a lot with fixed document content, to
avoid unnecessary parsing, this PR caches the `StateDelta` to the
server.

# User-Facing Changes
<!-- List of all changes that impact the user experience here. This
helps us keep track of breaking changes. -->

Can be disabled on the client side.

# Tests + Formatting
<!--
Don't forget to add tests that cover your changes.

Make sure you've run and fixed any issues with these commands:

- `cargo fmt --all -- --check` to check standard code formatting (`cargo
fmt --all` applies these changes)
- `cargo clippy --workspace -- -D warnings -D clippy::unwrap_used` to
check that you're using the standard code style
- `cargo test --workspace` to check that all tests pass (on Windows make
sure to [enable developer
mode](https://learn.microsoft.com/en-us/windows/apps/get-started/developer-mode-features-and-debugging))
- `cargo run -- -c "use toolkit.nu; toolkit test stdlib"` to run the
tests for the standard library

> **Note**
> from `nushell` you can also use the `toolkit` as follows
> ```bash
> use toolkit.nu # or use an `env_change` hook to activate it
automatically
> toolkit check pr
> ```
-->

Implementation is directly borrowed from `references`, only one simple
test case added.

# After Submitting
<!-- If your PR had any user-facing changes, update [the
documentation](https://github.com/nushell/nushell.github.io) after the
PR is merged, if necessary. This will help us keep the docs up to date.
-->

crates/nu-lsp/src/ast.rs

@@ -286,43 +286,22 @@ fn try_find_id_in_use(
     if call_name != "use".as_bytes() {
         return None;
     }
-    let find_by_name = |name: &str| {
-        match id {
-            Some(Id::Variable(var_id_ref)) => {
-                if let Some(var_id) = working_set.find_variable(name.as_bytes()) {
-                    if var_id == *var_id_ref {
-                        return Some(Id::Variable(var_id));
-                    }
-                }
-            }
-            Some(Id::Declaration(decl_id_ref)) => {
-                if let Some(decl_id) = working_set.find_decl(name.as_bytes()) {
-                    if decl_id == *decl_id_ref {
-                        return Some(Id::Declaration(decl_id));
-                    }
-                }
-            }
-            Some(Id::Module(module_id_ref)) => {
-                if let Some(module_id) = working_set.find_module(name.as_bytes()) {
-                    if module_id == *module_id_ref {
-                        return Some(Id::Module(module_id));
-                    }
-                }
-            }
-            None => {
-                if let Some(var_id) = working_set.find_variable(name.as_bytes()) {
-                    return Some(Id::Variable(var_id));
-                }
-                if let Some(decl_id) = working_set.find_decl(name.as_bytes()) {
-                    return Some(Id::Declaration(decl_id));
-                }
-                if let Some(module_id) = working_set.find_module(name.as_bytes()) {
-                    return Some(Id::Module(module_id));
-                }
-            }
-            _ => (),
-        }
-        None
+    let find_by_name = |name: &[u8]| match id {
+        Some(Id::Variable(var_id_ref)) => working_set
+            .find_variable(name)
+            .and_then(|var_id| (var_id == *var_id_ref).then_some(Id::Variable(var_id))),
+        Some(Id::Declaration(decl_id_ref)) => working_set
+            .find_decl(name)
+            .and_then(|decl_id| (decl_id == *decl_id_ref).then_some(Id::Declaration(decl_id))),
+        Some(Id::Module(module_id_ref)) => working_set
+            .find_module(name)
+            .and_then(|module_id| (module_id == *module_id_ref).then_some(Id::Module(module_id))),
+        None => working_set
+            .find_variable(name)
+            .map(Id::Variable)
+            .or(working_set.find_decl(name).map(Id::Declaration))
+            .or(working_set.find_module(name).map(Id::Module)),
+        _ => None,
     };
     let check_location = |span: &Span| location.map_or(true, |pos| span.contains(*pos));
     let get_module_id = |span: Span| {
@@ -330,8 +309,7 @@ fn try_find_id_in_use(
         let name = String::from_utf8_lossy(working_set.get_span_contents(span));
         let path = PathBuf::from(name.as_ref());
         let stem = path.file_stem().and_then(|fs| fs.to_str()).unwrap_or(&name);
-        let module_id = working_set.find_module(stem.as_bytes())?;
-        let found_id = Id::Module(module_id);
+        let found_id = Id::Module(working_set.find_module(stem.as_bytes())?);
         id.map_or(true, |id_r| found_id == *id_r)
             .then_some((found_id, span))
     };
@@ -359,7 +337,7 @@ fn try_find_id_in_use(
                 .and_then(|e| {
                     let name = e.as_string()?;
                     Some((
-                        find_by_name(&name)?,
+                        find_by_name(name.as_bytes())?,
                         strip_quotes(item_expr.span, working_set),
                     ))
                 })
@@ -367,31 +345,25 @@ fn try_find_id_in_use(
     };
 
     // the imported name is always at the second argument
-    if let Argument::Positional(expr) = call.arguments.get(1)? {
-        if check_location(&expr.span) {
-            match &expr.expr {
-                Expr::String(name) => {
-                    if let Some(id) = find_by_name(name) {
-                        return Some((id, strip_quotes(expr.span, working_set)));
-                    }
-                }
-                Expr::List(items) => {
-                    if let Some(res) = search_in_list_items(items) {
-                        return Some(res);
-                    }
-                }
-                Expr::FullCellPath(fcp) => {
-                    if let Expr::List(items) = &fcp.head.expr {
-                        if let Some(res) = search_in_list_items(items) {
-                            return Some(res);
-                        }
-                    }
-                }
-                _ => (),
-            }
-        }
+    let Argument::Positional(expr) = call.arguments.get(1)? else {
+        return None;
+    };
+    if !check_location(&expr.span) {
+        return None;
+    }
+    match &expr.expr {
+        Expr::String(name) => {
+            find_by_name(name.as_bytes()).map(|id| (id, strip_quotes(expr.span, working_set)))
+        }
+        Expr::List(items) => search_in_list_items(items),
+        Expr::FullCellPath(fcp) => {
+            let Expr::List(items) = &fcp.head.expr else {
+                return None;
+            };
+            search_in_list_items(items)
+        }
+        _ => None,
     }
-    None
 }
 
 fn find_id_in_expr(
@@ -483,6 +455,7 @@ fn find_reference_by_id_in_expr(
             if let Id::Declaration(decl_id) = id {
                 if *decl_id == call.decl_id {
                     occurs.push(call.head);
+                    return Some(occurs);
                 }
             }
             if let Some((_, span_found)) = try_find_id_in_def(call, working_set, None, Some(id))

crates/nu-lsp/src/diagnostics.rs

@@ -10,7 +10,7 @@ impl LanguageServer {
         let mut engine_state = self.new_engine_state();
         engine_state.generate_nu_constant();
 
-        let Some((_, offset, working_set)) = self.parse_file(&mut engine_state, &uri, true) else {
+        let Some((_, span, working_set)) = self.parse_file(&mut engine_state, &uri, true) else {
             return Ok(());
         };
 
@@ -31,7 +31,7 @@ impl LanguageServer {
             let message = err.to_string();
 
             diagnostics.diagnostics.push(Diagnostic {
-                range: span_to_range(&err.span(), file, offset),
+                range: span_to_range(&err.span(), file, span.start),
                 severity: Some(DiagnosticSeverity::ERROR),
                 message,
                 ..Default::default()

crates/nu-lsp/src/lib.rs

@@ -16,7 +16,7 @@ use nu_cli::{NuCompleter, SuggestionKind};
 use nu_parser::parse;
 use nu_protocol::{
     ast::Block,
-    engine::{CachedFile, EngineState, Stack, StateWorkingSet},
+    engine::{CachedFile, EngineState, Stack, StateDelta, StateWorkingSet},
     DeclId, ModuleId, Span, Type, Value, VarId,
 };
 use std::{collections::BTreeMap, sync::Mutex};
@@ -50,13 +50,17 @@ pub struct LanguageServer {
     connection: Connection,
     io_threads: Option<IoThreads>,
     docs: Arc<Mutex<TextDocuments>>,
-    engine_state: EngineState,
+    initial_engine_state: EngineState,
     symbol_cache: SymbolCache,
     inlay_hints: BTreeMap<Uri, Vec<InlayHint>>,
     workspace_folders: BTreeMap<String, WorkspaceFolder>,
-    // for workspace wide requests
+    /// for workspace wide requests
     occurrences: BTreeMap<Uri, Vec<Range>>,
     channels: Option<(Sender<bool>, Arc<Receiver<InternalMessage>>)>,
+    /// set to true when text changes
+    need_parse: bool,
+    /// cache `StateDelta` to avoid repeated parsing
+    cached_state_delta: Option<StateDelta>,
 }
 
 pub fn path_to_uri(path: impl AsRef<Path>) -> Uri {
@@ -96,12 +100,14 @@ impl LanguageServer {
             connection,
             io_threads,
             docs: Arc::new(Mutex::new(TextDocuments::new())),
-            engine_state,
+            initial_engine_state: engine_state,
             symbol_cache: SymbolCache::new(),
             inlay_hints: BTreeMap::new(),
             workspace_folders: BTreeMap::new(),
             occurrences: BTreeMap::new(),
             channels: None,
+            need_parse: true,
+            cached_state_delta: None,
         })
     }
 
@@ -110,22 +116,22 @@ impl LanguageServer {
             work_done_progress: Some(true),
         };
         let server_capabilities = serde_json::to_value(ServerCapabilities {
-            text_document_sync: Some(lsp_types::TextDocumentSyncCapability::Kind(
-                TextDocumentSyncKind::INCREMENTAL,
-            )),
-            definition_provider: Some(OneOf::Left(true)),
-            hover_provider: Some(lsp_types::HoverProviderCapability::Simple(true)),
             completion_provider: Some(lsp_types::CompletionOptions::default()),
+            definition_provider: Some(OneOf::Left(true)),
+            document_highlight_provider: Some(OneOf::Left(true)),
             document_symbol_provider: Some(OneOf::Left(true)),
-            workspace_symbol_provider: Some(OneOf::Left(true)),
+            hover_provider: Some(lsp_types::HoverProviderCapability::Simple(true)),
             inlay_hint_provider: Some(OneOf::Left(true)),
+            references_provider: Some(OneOf::Right(ReferencesOptions {
+                work_done_progress_options,
+            })),
             rename_provider: Some(OneOf::Right(RenameOptions {
                 prepare_provider: Some(true),
                 work_done_progress_options,
             })),
-            references_provider: Some(OneOf::Right(ReferencesOptions {
-                work_done_progress_options,
-            })),
+            text_document_sync: Some(lsp_types::TextDocumentSyncCapability::Kind(
+                TextDocumentSyncKind::INCREMENTAL,
+            )),
             workspace: Some(WorkspaceServerCapabilities {
                 workspace_folders: Some(WorkspaceFoldersServerCapabilities {
                     supported: Some(true),
@@ -133,18 +139,19 @@ impl LanguageServer {
                 }),
                 ..Default::default()
             }),
+            workspace_symbol_provider: Some(OneOf::Left(true)),
             ..Default::default()
         })
         .expect("Must be serializable");
         let init_params = self
             .connection
             .initialize_while(server_capabilities, || {
-                !self.engine_state.signals().interrupted()
+                !self.initial_engine_state.signals().interrupted()
             })
             .into_diagnostic()?;
         self.initialize_workspace_folders(init_params)?;
 
-        while !self.engine_state.signals().interrupted() {
+        while !self.initial_engine_state.signals().interrupted() {
             // first check new messages from child thread
             self.handle_internal_messages()?;
 
@@ -175,30 +182,22 @@ impl LanguageServer {
                     }
 
                     let resp = match request.method.as_str() {
+                        request::Completion::METHOD => {
+                            Self::handle_lsp_request(request, |params| self.complete(params))
+                        }
+                        request::DocumentHighlightRequest::METHOD => {
+                            Self::handle_lsp_request(request, |params| {
+                                self.document_highlight(params)
+                            })
+                        }
                         request::GotoDefinition::METHOD => {
                             Self::handle_lsp_request(request, |params| self.goto_definition(params))
                         }
                         request::HoverRequest::METHOD => {
                             Self::handle_lsp_request(request, |params| self.hover(params))
                         }
-                        request::Completion::METHOD => {
-                            Self::handle_lsp_request(request, |params| self.complete(params))
-                        }
-                        request::DocumentSymbolRequest::METHOD => {
-                            Self::handle_lsp_request(request, |params| self.document_symbol(params))
-                        }
-                        request::References::METHOD => {
-                            Self::handle_lsp_request(request, |params| {
-                                self.references(params, 5000)
-                            })
-                        }
-                        request::WorkspaceSymbolRequest::METHOD => {
-                            Self::handle_lsp_request(request, |params| {
-                                self.workspace_symbol(params)
-                            })
-                        }
-                        request::Rename::METHOD => {
-                            Self::handle_lsp_request(request, |params| self.rename(params))
+                        request::InlayHintRequest::METHOD => {
+                            Self::handle_lsp_request(request, |params| self.get_inlay_hints(params))
                         }
                         request::PrepareRenameRequest::METHOD => {
                             let id = request.id.clone();
@@ -207,8 +206,21 @@ impl LanguageServer {
                             }
                             continue;
                         }
-                        request::InlayHintRequest::METHOD => {
-                            Self::handle_lsp_request(request, |params| self.get_inlay_hints(params))
+                        request::References::METHOD => {
+                            Self::handle_lsp_request(request, |params| {
+                                self.references(params, 5000)
+                            })
+                        }
+                        request::Rename::METHOD => {
+                            Self::handle_lsp_request(request, |params| self.rename(params))
+                        }
+                        request::DocumentSymbolRequest::METHOD => {
+                            Self::handle_lsp_request(request, |params| self.document_symbol(params))
+                        }
+                        request::WorkspaceSymbolRequest::METHOD => {
+                            Self::handle_lsp_request(request, |params| {
+                                self.workspace_symbol(params)
+                            })
                         }
                         _ => {
                             continue;
@@ -223,6 +235,7 @@ impl LanguageServer {
                 Message::Response(_) => {}
                 Message::Notification(notification) => {
                     if let Some(updated_file) = self.handle_lsp_notification(notification) {
+                        self.need_parse = true;
                         self.symbol_cache.mark_dirty(updated_file.clone(), true);
                         self.publish_diagnostics_for_file(updated_file)?;
                     }
@@ -274,9 +287,19 @@ impl LanguageServer {
     }
 
     pub fn new_engine_state(&self) -> EngineState {
-        let mut engine_state = self.engine_state.clone();
+        let mut engine_state = self.initial_engine_state.clone();
         let cwd = std::env::current_dir().expect("Could not get current working directory.");
         engine_state.add_env_var("PWD".into(), Value::test_string(cwd.to_string_lossy()));
+        // merge the cached `StateDelta` if text not changed
+        if !self.need_parse {
+            engine_state
+                .merge_delta(
+                    self.cached_state_delta
+                        .to_owned()
+                        .expect("Tried to merge a non-existing state delta"),
+                )
+                .expect("Failed to merge state delta");
+        }
         engine_state
     }
 
@@ -286,7 +309,7 @@ impl LanguageServer {
         uri: &Uri,
         pos: Position,
     ) -> Result<(StateWorkingSet<'a>, Id, Span, usize)> {
-        let (block, file_offset, mut working_set) = self
+        let (block, file_span, working_set) = self
             .parse_file(engine_state, uri, false)
             .ok_or_else(|| miette!("\nFailed to parse current file"))?;
 
@@ -297,12 +320,10 @@ impl LanguageServer {
         let file = docs
             .get_document(uri)
             .ok_or_else(|| miette!("\nFailed to get document"))?;
-        let location = file.offset_at(pos) as usize + file_offset;
+        let location = file.offset_at(pos) as usize + file_span.start;
         let (id, span) = find_id(&block, &working_set, &location)
             .ok_or_else(|| miette!("\nFailed to find current name"))?;
-        // add block to working_set for later references
-        working_set.add_block(block);
-        Ok((working_set, id, span, file_offset))
+        Ok((working_set, id, span, file_span.start))
     }
 
     pub fn parse_file<'a>(
@@ -310,7 +331,7 @@ impl LanguageServer {
         engine_state: &'a mut EngineState,
         uri: &Uri,
         need_hints: bool,
-    ) -> Option<(Arc<Block>, usize, StateWorkingSet<'a>)> {
+    ) -> Option<(Arc<Block>, Span, StateWorkingSet<'a>)> {
         let mut working_set = StateWorkingSet::new(engine_state);
         let docs = self.docs.lock().ok()?;
         let file = docs.get_document(uri)?;
@@ -319,15 +340,19 @@ impl LanguageServer {
         let contents = file.get_content(None).as_bytes();
         let _ = working_set.files.push(file_path.clone(), Span::unknown());
         let block = parse(&mut working_set, Some(file_path_str), contents, false);
-        let offset = working_set.get_span_for_filename(file_path_str)?.start;
-        // TODO: merge delta back to engine_state?
-        // self.engine_state.merge_delta(working_set.render());
-
+        let span = working_set.get_span_for_filename(file_path_str)?;
         if need_hints {
-            let file_inlay_hints = self.extract_inlay_hints(&working_set, &block, offset, file);
+            let file_inlay_hints = self.extract_inlay_hints(&working_set, &block, span.start, file);
             self.inlay_hints.insert(uri.clone(), file_inlay_hints);
         }
-        Some((block, offset, working_set))
+        if self.need_parse {
+            // TODO: incremental parsing
+            // add block to working_set for later references
+            working_set.add_block(block.clone());
+            self.cached_state_delta = Some(working_set.delta.clone());
+            self.need_parse = false;
+        }
+        Some((block, span, working_set))
     }
 
     fn get_location_by_span<'a>(
@@ -352,7 +377,7 @@ impl LanguageServer {
                     let temp_doc = FullTextDocument::new(
                         "nu".to_string(),
                         0,
-                        String::from_utf8((*cached_file.content).to_vec()).expect("Invalid UTF-8"),
+                        String::from_utf8_lossy(cached_file.content.as_ref()).to_string(),
                     );
                     return Some(Location {
                         uri: target_uri,
@@ -587,8 +612,10 @@ impl LanguageServer {
         let docs = self.docs.lock().ok()?;
         let file = docs.get_document(&path_uri)?;
 
-        let mut completer =
-            NuCompleter::new(Arc::new(self.engine_state.clone()), Arc::new(Stack::new()));
+        let mut completer = NuCompleter::new(
+            Arc::new(self.initial_engine_state.clone()),
+            Arc::new(Stack::new()),
+        );
 
         let location = file.offset_at(params.text_document_position.position) as usize;
         let results = completer.fetch_completions_at(&file.get_content(None)[..location], location);

crates/nu-lsp/src/symbols.rs

@@ -203,16 +203,17 @@ impl SymbolCache {
                     continue;
                 }
                 let target_uri = path_to_uri(path);
-                let new_symbols = if let Some(doc) = docs.get_document(&target_uri) {
-                    Self::extract_all_symbols(&working_set, doc, cached_file)
-                } else {
-                    let temp_doc = FullTextDocument::new(
-                        "nu".to_string(),
-                        0,
-                        String::from_utf8((*cached_file.content).to_vec()).expect("Invalid UTF-8"),
-                    );
-                    Self::extract_all_symbols(&working_set, &temp_doc, cached_file)
-                };
+                let new_symbols = Self::extract_all_symbols(
+                    &working_set,
+                    docs.get_document(&target_uri)
+                        .unwrap_or(&FullTextDocument::new(
+                            "nu".to_string(),
+                            0,
+                            String::from_utf8((*cached_file.content).to_vec())
+                                .expect("Invalid UTF-8"),
+                        )),
+                    cached_file,
+                );
                 self.cache.insert(target_uri.clone(), new_symbols);
                 self.mark_dirty(target_uri, false);
             }

crates/nu-lsp/src/workspace.rs

@@ -8,12 +8,14 @@ use std::{
 };
 
 use crate::{
-    ast::find_reference_by_id, path_to_uri, span_to_range, uri_to_path, Id, LanguageServer,
+    ast::{find_id, find_reference_by_id},
+    path_to_uri, span_to_range, uri_to_path, Id, LanguageServer,
 };
 use crossbeam_channel::{Receiver, Sender};
 use lsp_server::{Message, Request, Response};
 use lsp_types::{
-    Location, PrepareRenameResponse, ProgressToken, Range, ReferenceParams, RenameParams,
+    DocumentHighlight, DocumentHighlightKind, DocumentHighlightParams, Location,
+    PrepareRenameResponse, ProgressToken, Range, ReferenceParams, RenameParams,
     TextDocumentPositionParams, TextEdit, Uri, WorkspaceEdit, WorkspaceFolder,
 };
 use miette::{miette, IntoDiagnostic, Result};
@@ -69,7 +71,7 @@ fn find_reference_in_file(
 }
 
 impl LanguageServer {
-    /// get initial workspace folders from initialization response
+    /// Get initial workspace folders from initialization response
     pub fn initialize_workspace_folders(&mut self, init_params: Value) -> Result<()> {
         if let Some(array) = init_params.get("workspaceFolders") {
             let folders: Vec<WorkspaceFolder> =
@@ -81,6 +83,43 @@ impl LanguageServer {
         Ok(())
     }
 
+    /// Highlight all occurrences of the text at cursor, in current file
+    pub fn document_highlight(
+        &mut self,
+        params: &DocumentHighlightParams,
+    ) -> Option<Vec<DocumentHighlight>> {
+        let mut engine_state = self.new_engine_state();
+        let path_uri = params
+            .text_document_position_params
+            .text_document
+            .uri
+            .to_owned();
+        let (block, file_span, working_set) =
+            self.parse_file(&mut engine_state, &path_uri, false)?;
+        let docs = &self.docs.lock().ok()?;
+        let file = docs.get_document(&path_uri)?;
+        let location = file.offset_at(params.text_document_position_params.position) as usize
+            + file_span.start;
+        let (id, cursor_span) = find_id(&block, &working_set, &location)?;
+        let mut refs = find_reference_by_id(&block, &working_set, &id);
+        let definition_span = Self::find_definition_span_by_id(&working_set, &id);
+        if let Some(extra_span) =
+            Self::reference_not_in_ast(&id, &working_set, definition_span, file_span, cursor_span)
+        {
+            if !refs.contains(&extra_span) {
+                refs.push(extra_span);
+            }
+        }
+        Some(
+            refs.iter()
+                .map(|span| DocumentHighlight {
+                    range: span_to_range(span, file, file_span.start),
+                    kind: Some(DocumentHighlightKind::TEXT),
+                })
+                .collect(),
+        )
+    }
+
     /// The rename request only happens after the client received a `PrepareRenameResponse`,
     /// and a new name typed in, could happen before ranges ready for all files in the workspace folder
     pub fn rename(&mut self, params: &RenameParams) -> Option<WorkspaceEdit> {
@@ -118,7 +157,7 @@ impl LanguageServer {
         self.occurrences = BTreeMap::new();
         let mut engine_state = self.new_engine_state();
         let path_uri = params.text_document_position.text_document.uri.to_owned();
-        let (working_set, id, span, _) = self
+        let (_, id, span, _) = self
             .parse_and_find(
                 &mut engine_state,
                 &path_uri,
@@ -126,8 +165,7 @@ impl LanguageServer {
             )
             .ok()?;
         // have to clone it again in order to move to another thread
-        let mut engine_state = self.new_engine_state();
-        engine_state.merge_delta(working_set.render()).ok()?;
+        let engine_state = self.new_engine_state();
         let current_workspace_folder = self.get_workspace_folder_by_uri(&path_uri)?;
         let token = params
             .work_done_progress_params
@@ -235,6 +273,33 @@ impl LanguageServer {
         Ok(())
     }
 
+    /// NOTE: for arguments whose declaration is in a signature
+    /// which is not covered in the AST
+    fn reference_not_in_ast(
+        id: &Id,
+        working_set: &StateWorkingSet,
+        definition_span: Option<Span>,
+        file_span: Span,
+        sample_span: Span,
+    ) -> Option<Span> {
+        if let (Id::Variable(_), Some(decl_span)) = (&id, definition_span) {
+            if file_span.contains_span(decl_span) && decl_span.end > decl_span.start {
+                let leading_dashes = working_set
+                    .get_span_contents(decl_span)
+                    .iter()
+                    // remove leading dashes for flags
+                    .take_while(|c| *c == &b'-')
+                    .count();
+                let start = decl_span.start + leading_dashes;
+                return Some(Span {
+                    start,
+                    end: start + sample_span.end - sample_span.start,
+                });
+            }
+        }
+        None
+    }
+
     fn find_reference_in_workspace(
         &self,
         engine_state: EngineState,
@@ -310,24 +375,15 @@ impl LanguageServer {
                     let file_span = working_set
                         .get_span_for_filename(fp.to_string_lossy().as_ref())
                         .unwrap_or(Span::unknown());
-                    // NOTE: for arguments whose declaration is in a signature
-                    // which is not covered in the AST
-                    if let (Id::Variable(_), Some(decl_span)) = (&id, definition_span) {
-                        if file_span.contains_span(decl_span)
-                            && decl_span.end > decl_span.start
-                            && !refs.contains(&decl_span)
-                        {
-                            let leading_dashes = working_set
-                                .get_span_contents(decl_span)
-                                .iter()
-                                // remove leading dashes for flags
-                                .take_while(|c| *c == &b'-')
-                                .count();
-                            let start = decl_span.start + leading_dashes;
-                            refs.push(Span {
-                                start,
-                                end: start + span.end - span.start,
-                            });
+                    if let Some(extra_span) = Self::reference_not_in_ast(
+                        &id,
+                        &working_set,
+                        definition_span,
+                        file_span,
+                        span,
+                    ) {
+                        if !refs.contains(&extra_span) {
+                            refs.push(extra_span)
                         }
                     }
                     let ranges = refs
@@ -367,12 +423,12 @@ impl LanguageServer {
 mod tests {
     use assert_json_diff::assert_json_eq;
     use lsp_server::{Connection, Message};
-    use lsp_types::RenameParams;
     use lsp_types::{
         request, request::Request, InitializeParams, PartialResultParams, Position,
         ReferenceContext, ReferenceParams, TextDocumentIdentifier, TextDocumentPositionParams, Uri,
         WorkDoneProgressParams, WorkspaceFolder,
     };
+    use lsp_types::{DocumentHighlightParams, RenameParams};
     use nu_parser::parse;
     use nu_protocol::engine::StateWorkingSet;
     use nu_test_support::fs::fixtures;
@@ -482,6 +538,35 @@ mod tests {
             .unwrap()
     }
 
+    fn send_document_highlight_request(
+        client_connection: &Connection,
+        uri: Uri,
+        line: u32,
+        character: u32,
+    ) -> Message {
+        client_connection
+            .sender
+            .send(Message::Request(lsp_server::Request {
+                id: 1.into(),
+                method: request::DocumentHighlightRequest::METHOD.to_string(),
+                params: serde_json::to_value(DocumentHighlightParams {
+                    text_document_position_params: TextDocumentPositionParams {
+                        text_document: TextDocumentIdentifier { uri },
+                        position: Position { line, character },
+                    },
+                    partial_result_params: PartialResultParams::default(),
+                    work_done_progress_params: WorkDoneProgressParams::default(),
+                })
+                .unwrap(),
+            }))
+            .unwrap();
+
+        client_connection
+            .receiver
+            .recv_timeout(std::time::Duration::from_secs(2))
+            .unwrap()
+    }
+
     #[test]
     fn command_reference_in_workspace() {
         let mut script = fixtures();
@@ -803,4 +888,28 @@ mod tests {
             .unwrap();
         assert!(working_set.find_block_by_span(span_foo).is_some())
     }
+
+    #[test]
+    fn document_highlight_variable() {
+        let mut script = fixtures();
+        script.push("lsp");
+        script.push("workspace");
+        script.push("foo.nu");
+        let script = path_to_uri(&script);
+
+        let (client_connection, _recv) = initialize_language_server(None);
+        open_unchecked(&client_connection, script.clone());
+
+        let message = send_document_highlight_request(&client_connection, script.clone(), 3, 5);
+        let Message::Response(r) = message else {
+            panic!("unexpected message type");
+        };
+        assert_json_eq!(
+            r.result,
+            serde_json::json!([
+                { "range": { "start": { "line": 3, "character": 3 }, "end": { "line": 3, "character": 8 } }, "kind": 1 },
+                { "range": { "start": { "line": 1, "character": 4 }, "end": { "line": 1, "character": 9 } }, "kind": 1 }
+            ]),
+        );
+    }
 }

crates/nu-protocol/src/engine/state_delta.rs

@@ -14,6 +14,7 @@ use crate::{PluginRegistryItem, RegisteredPlugin};
 /// A delta (or change set) between the current global state and a possible future global state. Deltas
 /// can be applied to the global state to update it to contain both previous state and the state held
 /// within the delta.
+#[derive(Clone)]
 pub struct StateDelta {
     pub(super) files: Vec<CachedFile>,
     pub(super) virtual_paths: Vec<(String, VirtualPath)>,

tests/fixtures/lsp/symbols/foo.nu

@@ -1,4 +1,4 @@
-use bar.nu [ var_bar def_bar ]
+use bar.nu [ def_bar ]
 
 let var_foo = 1
 def_bar