[analysis_server] [lsp] Support return/yield/break/continue in Document Highlights

This adds support for keywords like `return`, `yield`, `break`, `continue` to Document Highlights. For break/continue, the matching loop keyword (`do`/`while`/`for`) is also highlighted (and this works in both directions). This behaviour matches what I see for TypeScript in VS Code.

This is only supported for LSP because the legacy protocol groups these by "Elements" which we don't have for loops (though since IntelliJ uses its own data for occurrences, it's not clear to me if anyone is using Occurrences over the legacy protocol).

Fixes #61170

Change-Id: I5563ec0a91a6fe33d10a6317d33e9256b3c50209
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/442061
Reviewed-by: Brian Wilkerson <[email protected]>
Reviewed-by: Samuel Rawlins <[email protected]>
Commit-Queue: Samuel Rawlins <[email protected]>

Author: DanTup

pkg/analysis_server/lib/src/domains/analysis/occurrences_dart.dart

@@ -9,11 +9,12 @@ import 'package:analyzer/dart/ast/token.dart';
 import 'package:analyzer/dart/ast/visitor.dart';
 import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/src/dart/ast/extensions.dart';
+import 'package:analyzer/src/utilities/extensions/collection.dart';
 
 void addDartOccurrences(OccurrencesCollector collector, CompilationUnit unit) {
   var visitor = DartUnitOccurrencesComputerVisitor();
   unit.accept(visitor);
-  visitor.elementsOffsetLengths.forEach((engineElement, offsetLengths) {
+  visitor.elementOccurrences.forEach((engineElement, nodes) {
     // For legacy protocol, we only support occurrences with the same
     // length, so we must filter the offset to only those that match the length
     // from the element.
@@ -24,18 +25,51 @@ void addDartOccurrences(OccurrencesCollector collector, CompilationUnit unit) {
     var length =
         serverElement.location?.length ?? engineElement.name?.length ?? 0;
     var offsets =
-        offsetLengths
-            .where((offsetLength) => offsetLength.$2 == length)
-            .map((offsetLength) => offsetLength.$1)
+        nodes
+            .where((node) => node.length == length)
+            .map((node) => node.offset)
             .toList();
 
     var occurrences = protocol.Occurrences(serverElement, offsets, length);
     collector.addOccurrences(occurrences);
   });
 }
 
-class DartUnitOccurrencesComputerVisitor extends RecursiveAstVisitor<void> {
-  final Map<Element, List<(int, int)>> elementsOffsetLengths = {};
+/// Returns both element-based and node-based occurrences for a file.
+///
+/// The legacy protocol does not support occurrences that are not based on
+/// elements and therefore use [addDartOccurrences]. This method is for LSP
+/// which does not use elements and instead only cares about the resulting
+/// ranges of code, allowing ranges for keywords like `return` and `break` to be
+/// included.
+List<Occurrences> getAllOccurrences(CompilationUnit unit) {
+  // TODO(dantup): Since LSP always starts with a target element or node,
+  //  we should consider passing it in here to avoid building the occurrences
+  // for the whole file and then extracting only the matches we want.
+  var visitor = DartUnitOccurrencesComputerVisitor();
+  unit.accept(visitor);
+
+  return [
+    // Element-based occurrences
+    for (var MapEntry(key: element, value: tokens)
+        in visitor.elementOccurrences.entries)
+      ElementOccurrences(element, tokens),
+    // Node-based occurrences
+    for (var MapEntry(key: node, value: tokens)
+        in visitor.nodeOccurrences.entries)
+      NodeOccurrences(node, tokens),
+  ];
+}
+
+class DartUnitOccurrencesComputerVisitor extends GeneralizingAstVisitor<void> {
+  /// Occurrences tracked by their elements.
+  final Map<Element, List<Token>> elementOccurrences = {};
+
+  /// Occurrences tracked by nodes (such as loops and their exit keywords).
+  final Map<AstNode, List<Token>> nodeOccurrences = {};
+
+  // Stack to track the current function for return/yield keywords
+  final List<AstNode> _functionStack = [];
 
   @override
   void visitAssignedVariablePattern(AssignedVariablePattern node) {
@@ -47,6 +81,13 @@ class DartUnitOccurrencesComputerVisitor extends RecursiveAstVisitor<void> {
     super.visitAssignedVariablePattern(node);
   }
 
+  @override
+  void visitBreakStatement(BreakStatement node) {
+    _addNodeOccurrence(node.target, node.breakKeyword);
+
+    super.visitBreakStatement(node);
+  }
+
   @override
   void visitClassDeclaration(ClassDeclaration node) {
     _addOccurrence(node.declaredFragment!.element, node.name);
@@ -66,10 +107,9 @@ class DartUnitOccurrencesComputerVisitor extends RecursiveAstVisitor<void> {
     if (node.name case var name?) {
       _addOccurrence(node.declaredFragment!.element, name);
     } else {
-      _addOccurrenceAt(
+      _addOccurrence(
         node.declaredFragment!.element,
-        node.returnType.offset,
-        node.returnType.length,
+        node.returnType.beginToken,
       );
     }
 
@@ -91,6 +131,13 @@ class DartUnitOccurrencesComputerVisitor extends RecursiveAstVisitor<void> {
     super.visitConstructorName(node);
   }
 
+  @override
+  void visitContinueStatement(ContinueStatement node) {
+    _addNodeOccurrence(node.target, node.continueKeyword);
+
+    super.visitContinueStatement(node);
+  }
+
   @override
   void visitDeclaredIdentifier(DeclaredIdentifier node) {
     _addOccurrence(node.declaredFragment!.element, node.name);
@@ -109,6 +156,13 @@ class DartUnitOccurrencesComputerVisitor extends RecursiveAstVisitor<void> {
     super.visitDeclaredVariablePattern(node);
   }
 
+  @override
+  void visitDoStatement(DoStatement node) {
+    _addNodeOccurrence(node, node.doKeyword);
+
+    super.visitDoStatement(node);
+  }
+
   @override
   void visitEnumConstantDeclaration(EnumConstantDeclaration node) {
     _addOccurrence(node.declaredFragment!.element, node.name);
@@ -159,6 +213,20 @@ class DartUnitOccurrencesComputerVisitor extends RecursiveAstVisitor<void> {
     super.visitFieldFormalParameter(node);
   }
 
+  @override
+  void visitForStatement(ForStatement node) {
+    _addNodeOccurrence(node, node.forKeyword);
+
+    super.visitForStatement(node);
+  }
+
+  @override
+  void visitFunctionBody(FunctionBody node) {
+    _functionStack.add(node);
+    super.visitFunctionBody(node);
+    _functionStack.removeLastOrNull();
+  }
+
   @override
   void visitFunctionDeclaration(FunctionDeclaration node) {
     _addOccurrence(node.declaredFragment!.element, node.name);
@@ -235,6 +303,13 @@ class DartUnitOccurrencesComputerVisitor extends RecursiveAstVisitor<void> {
     super.visitRepresentationDeclaration(node);
   }
 
+  @override
+  void visitReturnStatement(ReturnStatement node) {
+    _addNodeOccurrence(_functionStack.lastOrNull, node.returnKeyword);
+
+    super.visitReturnStatement(node);
+  }
+
   @override
   void visitSimpleFormalParameter(SimpleFormalParameter node) {
     var nameToken = node.name;
@@ -270,6 +345,13 @@ class DartUnitOccurrencesComputerVisitor extends RecursiveAstVisitor<void> {
     super.visitSuperFormalParameter(node);
   }
 
+  @override
+  void visitSwitchStatement(SwitchStatement node) {
+    _addNodeOccurrence(node, node.switchKeyword);
+
+    super.visitSwitchStatement(node);
+  }
+
   @override
   void visitTypeParameter(TypeParameter node) {
     if (node case TypeParameter(:var declaredFragment?)) {
@@ -285,21 +367,32 @@ class DartUnitOccurrencesComputerVisitor extends RecursiveAstVisitor<void> {
     super.visitVariableDeclaration(node);
   }
 
-  void _addOccurrence(Element element, Token token) {
-    _addOccurrenceAt(element, token.offset, token.length);
+  @override
+  void visitWhileStatement(WhileStatement node) {
+    _addNodeOccurrence(node, node.whileKeyword);
+
+    super.visitWhileStatement(node);
+  }
+
+  @override
+  void visitYieldStatement(YieldStatement node) {
+    _addNodeOccurrence(_functionStack.lastOrNull, node.yieldKeyword);
+
+    super.visitYieldStatement(node);
+  }
+
+  void _addNodeOccurrence(AstNode? node, Token token) {
+    if (node == null) return;
+
+    (nodeOccurrences[node] ??= []).add(token);
   }
 
-  void _addOccurrenceAt(Element element, int offset, int length) {
+  void _addOccurrence(Element element, Token token) {
     var canonicalElement = _canonicalizeElement(element);
     if (canonicalElement == null) {
       return;
     }
-    var offsetLengths = elementsOffsetLengths[canonicalElement];
-    if (offsetLengths == null) {
-      offsetLengths = <(int, int)>[];
-      elementsOffsetLengths[canonicalElement] = offsetLengths;
-    }
-    offsetLengths.add((offset, length));
+    (elementOccurrences[canonicalElement] ??= []).add(token);
   }
 
   Element? _canonicalizeElement(Element element) {
@@ -314,3 +407,28 @@ class DartUnitOccurrencesComputerVisitor extends RecursiveAstVisitor<void> {
     return canonicalElement?.baseElement;
   }
 }
+
+/// Occurrences grouped by an Element.
+class ElementOccurrences extends Occurrences {
+  final Element element;
+
+  @override
+  final List<Token> tokens;
+
+  ElementOccurrences(this.element, this.tokens);
+}
+
+/// Occurrences grouped by a node (for example exit keywords grouped by a loop).
+class NodeOccurrences extends Occurrences {
+  final AstNode node;
+
+  @override
+  final List<Token> tokens;
+
+  NodeOccurrences(this.node, this.tokens);
+}
+
+/// Base class for protocol-agnostic occurrences.
+sealed class Occurrences {
+  List<Token> get tokens;
+}

pkg/analysis_server/lib/src/legacy_analysis_server.dart

@@ -19,7 +19,8 @@ import 'package:analysis_server/src/analytics/analytics_manager.dart';
 import 'package:analysis_server/src/channel/channel.dart';
 import 'package:analysis_server/src/computer/computer_highlights.dart';
 import 'package:analysis_server/src/domains/analysis/occurrences.dart';
-import 'package:analysis_server/src/domains/analysis/occurrences_dart.dart';
+import 'package:analysis_server/src/domains/analysis/occurrences_dart.dart'
+    hide Occurrences;
 import 'package:analysis_server/src/flutter/flutter_notifications.dart';
 import 'package:analysis_server/src/handler/legacy/analysis_get_errors.dart';
 import 'package:analysis_server/src/handler/legacy/analysis_get_hover.dart';

pkg/analysis_server/lib/src/lsp/handlers/handler_document_highlights.dart

@@ -8,7 +8,7 @@ import 'package:analysis_server/src/lsp/error_or.dart';
 import 'package:analysis_server/src/lsp/handlers/handlers.dart';
 import 'package:analysis_server/src/lsp/mapping.dart';
 import 'package:analysis_server/src/lsp/registration/feature_registration.dart';
-import 'package:collection/collection.dart';
+import 'package:analyzer/dart/ast/token.dart';
 
 typedef StaticOptions = Either2<bool, DocumentHighlightOptions>;
 
@@ -45,42 +45,37 @@ class DocumentHighlightsHandler
     var offset = unit.mapResultSync((unit) => toOffset(unit.lineInfo, pos));
 
     return (unit, offset).mapResults((unit, requestedOffset) async {
-      var visitor = DartUnitOccurrencesComputerVisitor();
-      unit.unit.accept(visitor);
+      var occurrences = getAllOccurrences(unit.unit);
 
-      /// Checks whether an Occurrence offset/length spans the requested
+      /// Checks whether an Occurrence token spans the requested
       /// offset.
       ///
       /// It's possible multiple occurrences might match because some nodes
       /// such as object destructuring might match multiple elements (for
       /// example the object getter and a declared variable).
-      bool spansRequestedPosition((int, int) offsetLength) {
-        var (offset, length) = offsetLength;
-        return offset <= requestedOffset && offset + length >= requestedOffset;
+      bool spansRequestedPosition(Token token) {
+        return token.offset <= requestedOffset && token.end >= requestedOffset;
       }
 
-      // Find the groups (elements) of offset/lengths that contains an
+      // Find the groups of tokens that contains an
       // offset/length that spans the requested range. There may be multiple
       // matches here if the source element is in multiple groups.
-      var matchingSet =
-          visitor.elementsOffsetLengths.values
-              .where(
-                (offsetLengths) => offsetLengths.any(spansRequestedPosition),
-              )
-              .flattened
-              .toSet();
+      var matchingSet = <Token>{};
+
+      for (var occurrence in occurrences) {
+        var tokens = occurrence.tokens;
+        if (tokens.any(spansRequestedPosition)) {
+          matchingSet.addAll(tokens);
+        }
+      }
 
       // No matches will return an empty list (not null) because that prevents
       // the editor falling back to a text search.
       var highlights =
           matchingSet
               .map(
-                (offsetLength) => DocumentHighlight(
-                  range: toRange(
-                    unit.lineInfo,
-                    offsetLength.$1,
-                    offsetLength.$2,
-                  ),
+                (token) => DocumentHighlight(
+                  range: toRange(unit.lineInfo, token.offset, token.length),
                 ),
               )
               .toList();

pkg/analysis_server/lib/src/services/refactoring/legacy/rename_parameter.dart

@@ -16,7 +16,7 @@ import 'package:analyzer_plugin/protocol/protocol_common.dart';
 /// A [Refactoring] for renaming [analyzer.FormalParameterElement]s.
 class RenameParameterRefactoringImpl extends RenameRefactoringImpl {
   List<analyzer.FormalParameterElement> elements = [];
-  bool _renameAllPositionalOccurences = false;
+  bool _renameAllPositionalOccurrences = false;
 
   RenameParameterRefactoringImpl(
     super.workspace,
@@ -51,8 +51,8 @@ class RenameParameterRefactoringImpl extends RenameRefactoringImpl {
       if (resolvedUnit != null) {
         // If any of the resolved units have the lint enabled, we should avoid
         // renaming method parameters separately from the other implementations.
-        if (element.isPositional && !_renameAllPositionalOccurences) {
-          _renameAllPositionalOccurences |=
+        if (element.isPositional && !_renameAllPositionalOccurrences) {
+          _renameAllPositionalOccurrences |=
               getCodeStyleOptions(
                 resolvedUnit.file,
               ).avoidRenamingMethodParameters;
@@ -69,7 +69,7 @@ class RenameParameterRefactoringImpl extends RenameRefactoringImpl {
         );
       }
     }
-    if (_renameAllPositionalOccurences && elements.length > 1) {
+    if (_renameAllPositionalOccurrences && elements.length > 1) {
       result.addStatus(
         _RefactoringStatusExt.from(
           'This will also rename all related positional parameters '
@@ -97,7 +97,7 @@ class RenameParameterRefactoringImpl extends RenameRefactoringImpl {
     for (var element in elements) {
       if (element != this.element &&
           element.isPositional &&
-          !_renameAllPositionalOccurences) {
+          !_renameAllPositionalOccurrences) {
         continue;
       }
       var fieldRenamed = false;