🚥 Implement Document Highlight for all return / throws (function/closure/accessor/init/deinit) and break / continue (for/while/repeat...while)

Package.swift

@@ -508,6 +508,7 @@ var targets: [Target] = [
         "SwiftBasicFormat",
         "SwiftDiagnostics",
         "SwiftIDEUtils",
+        "SwiftLexicalLookup",
         "SwiftParser",
         "SwiftParserDiagnostics",
         "SwiftRefactor",
@@ -785,7 +786,8 @@ var dependencies: [Package.Dependency] {
       .package(url: "https://github.com/swiftlang/swift-tools-protocols.git", exact: "0.0.9"),
       .package(url: "https://github.com/swiftlang/swift-tools-support-core.git", branch: relatedDependenciesBranch),
       .package(url: "https://github.com/apple/swift-argument-parser.git", from: "1.5.1"),
-      .package(url: "https://github.com/swiftlang/swift-syntax.git", branch: relatedDependenciesBranch),
+      // .package(url: "https://github.com/swiftlang/swift-syntax.git", branch: relatedDependenciesBranch),
+      .package(path: "../swift-syntax"),
       .package(url: "https://github.com/apple/swift-crypto.git", from: "3.0.0"),
       // Not a build dependency. Used so the "Format Source Code" command plugin can be used to format sourcekit-lsp
       .package(url: "https://github.com/swiftlang/swift-format.git", branch: relatedDependenciesBranch),

Sources/SwiftLanguageService/SwiftLanguageService.swift

@@ -31,6 +31,8 @@ package import SwiftSyntax
 package import ToolchainRegistry
 @_spi(SourceKitLSP) import ToolsProtocolsSwiftExtensions
 
+@_spi(Experimental) import SwiftLexicalLookup
+
 #if os(Windows)
 import WinSDK
 #endif
@@ -875,6 +877,13 @@ extension SwiftLanguageService {
 
   package func documentSymbolHighlight(_ req: DocumentHighlightRequest) async throws -> [DocumentHighlight]? {
     let snapshot = try await self.latestSnapshot(for: req.textDocument.uri)
+
+    if let highlights = try await documentSymbolHighlightHelper(
+      at: req.position,
+      in: snapshot
+    ) {
+      return highlights
+    }
 
     let relatedIdentifiers = try await self.relatedIdentifiers(
       at: req.position,
@@ -889,6 +898,67 @@ extension SwiftLanguageService {
     }
   }
 
+  private func documentSymbolHighlightHelper(
+    at position: Position,
+    in snapshot: DocumentSnapshot
+  ) async throws -> [DocumentHighlight]? {
+    let syntaxTree = await syntaxTreeManager.syntaxTree(for: snapshot)
+
+    guard let tokenSyntax = syntaxTree.token(at: snapshot.absolutePosition(of: position)) else {
+      return nil
+    }
+
+    guard let targetStructure = Syntax(tokenSyntax).parent!.lookupControlStructure() else {
+      return nil
+    }
+    class ControlFlowHighlighter: SyntaxVisitor {
+      var highlights: [DocumentHighlight] = []
+      let targetStructure: Syntax
+      let snapshot: DocumentSnapshot
+
+      init(targetStructure: Syntax, snapshot: DocumentSnapshot) {
+        self.targetStructure = targetStructure
+        self.snapshot = snapshot
+        super.init(viewMode: .sourceAccurate)
+      }
+
+      private func addHighlightIfMatches(_ node: some SyntaxProtocol) {
+        if node.lookupControlStructure() == targetStructure {
+          highlights.append(
+            DocumentHighlight(
+              range: snapshot.absolutePositionRange(of: node.positionAfterSkippingLeadingTrivia..<node.endPosition),
+              kind: .read
+            )
+          )
+        }
+      }
+
+      override func visit(_ node: BreakStmtSyntax) -> SyntaxVisitorContinueKind {
+        addHighlightIfMatches(node)
+        return .skipChildren
+      }
+
+      override func visit(_ node: ContinueStmtSyntax) -> SyntaxVisitorContinueKind {
+        addHighlightIfMatches(node)
+        return .skipChildren
+      }
+
+      override func visit(_ node: ReturnStmtSyntax) -> SyntaxVisitorContinueKind {
+        addHighlightIfMatches(node)
+        return .skipChildren
+      }
+
+      override func visit(_ node: ThrowStmtSyntax) -> SyntaxVisitorContinueKind {
+        addHighlightIfMatches(node)
+        return .skipChildren
+      }
+    }
+
+    let highlighter = ControlFlowHighlighter(targetStructure: targetStructure, snapshot: snapshot)
+    highlighter.walk(targetStructure)
+    return highlighter.highlights.isEmpty ? nil : highlighter.highlights
+  }
+
   package func codeAction(_ req: CodeActionRequest) async throws -> CodeActionRequestResponse? {
     if (try? ReferenceDocumentURL(from: req.textDocument.uri)) != nil {
       // Do not show code actions in reference documents