Merge pull request #973 from ahoppen/ahoppen/no-local-vars-in-document-symbols

Don’t include local variables in document symbols

Author: ahoppen

Sources/SourceKitLSP/Swift/DocumentSymbols.swift

@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+import Foundation
 import LSPLogging
 import LanguageServerProtocol
 import SwiftSyntax
@@ -109,6 +110,45 @@ fileprivate final class DocumentSymbolsFinder: SyntaxAnyVisitor {
     )
   }
 
+  private func visit(_ trivia: Trivia, position: AbsolutePosition) {
+    let markPrefix = "MARK: "
+    var position = position
+    for piece in trivia.pieces {
+      defer {
+        position = position.advanced(by: piece.sourceLength.utf8Length)
+      }
+      switch piece {
+      case .lineComment(let commentText), .blockComment(let commentText):
+        let trimmedComment = commentText.trimmingCharacters(in: CharacterSet(["/", "*"]).union(.whitespaces))
+        if trimmedComment.starts(with: markPrefix) {
+          let markText = trimmedComment.dropFirst(markPrefix.count)
+          guard let rangeLowerBound = snapshot.position(of: position),
+            let rangeUpperBound = snapshot.position(of: position.advanced(by: piece.sourceLength.utf8Length))
+          else {
+            break
+          }
+          result.append(
+            DocumentSymbol(
+              name: String(markText),
+              kind: .namespace,
+              range: rangeLowerBound..<rangeUpperBound,
+              selectionRange: rangeLowerBound..<rangeUpperBound,
+              children: nil
+            )
+          )
+        }
+      default:
+        break
+      }
+    }
+  }
+
+  override func visit(_ node: TokenSyntax) -> SyntaxVisitorContinueKind {
+    self.visit(node.leadingTrivia, position: node.position)
+    self.visit(node.trailingTrivia, position: node.endPositionBeforeTrailingTrivia)
+    return .skipChildren
+  }
+
   override func visit(_ node: EnumCaseElementSyntax) -> SyntaxVisitorContinueKind {
     let rangeEnd =
       if let parameterClause = node.parameterClause {
@@ -180,7 +220,9 @@ fileprivate final class DocumentSymbolsFinder: SyntaxAnyVisitor {
     // If there is only one pattern binding within the variable decl, consider the entire variable decl as the
     // referenced range. If there are multiple, consider each pattern binding separately since the `var` keyword doesn't
     // belong to any pattern binding in particular.
-    guard let variableDecl = node.parent?.parent?.as(VariableDeclSyntax.self) else {
+    guard let variableDecl = node.parent?.parent?.as(VariableDeclSyntax.self),
+      variableDecl.isMemberOrTopLevelDeclaration
+    else {
       return .visitChildren
     }
     let rangeNode: Syntax = variableDecl.bindings.count == 1 ? Syntax(variableDecl) : Syntax(node)
@@ -228,6 +270,19 @@ fileprivate extension SyntaxProtocol {
   var rangeWithoutTrivia: Range<AbsolutePosition> {
     return positionAfterSkippingLeadingTrivia..<endPositionBeforeTrailingTrivia
   }
+
+  /// Whether this is a top-level constant or a member of a type, ie. if this is not a local variable.
+  var isMemberOrTopLevelDeclaration: Bool {
+    if self.parent?.is(MemberBlockItemSyntax.self) ?? false {
+      return true
+    }
+    if let codeBlockItem = self.parent?.as(CodeBlockItemSyntax.self),
+      codeBlockItem.parent?.parent?.is(SourceFileSyntax.self) ?? false
+    {
+      return true
+    }
+    return false
+  }
 }
 
 fileprivate extension TokenKind {

Tests/SourceKitLSPTests/DocumentSymbolTests.swift

@@ -571,8 +571,8 @@ final class DocumentSymbolTests: XCTestCase {
     try await assertDocumentSymbols(
       """
       1️⃣func 2️⃣f()3️⃣ {
-        4️⃣let 5️⃣localConstant6️⃣ = 07️⃣
-      }8️⃣
+        let localConstant = 0
+      }4️⃣
       """
     ) { positions in
       [
@@ -581,19 +581,9 @@ final class DocumentSymbolTests: XCTestCase {
           detail: nil,
           kind: .function,
           deprecated: nil,
-          range: positions["1️⃣"]..<positions["8️⃣"],
+          range: positions["1️⃣"]..<positions["4️⃣"],
           selectionRange: positions["2️⃣"]..<positions["3️⃣"],
-          children: [
-            DocumentSymbol(
-              name: "localConstant",
-              detail: nil,
-              kind: .variable,
-              deprecated: nil,
-              range: positions["4️⃣"]..<positions["7️⃣"],
-              selectionRange: positions["5️⃣"]..<positions["6️⃣"],
-              children: []
-            )
-          ]
+          children: []
         )
       ]
     }
@@ -630,6 +620,80 @@ final class DocumentSymbolTests: XCTestCase {
       ]
     }
   }
+
+  func testIncludeMarkComments() async throws {
+    try await assertDocumentSymbols(
+      """
+      1️⃣// MARK: Marker2️⃣
+      """
+    ) { positions in
+      [
+        DocumentSymbol(
+          name: "Marker",
+          kind: .namespace,
+          range: positions["1️⃣"]..<positions["2️⃣"],
+          selectionRange: positions["1️⃣"]..<positions["2️⃣"]
+        )
+      ]
+    }
+
+    try await assertDocumentSymbols(
+      """
+      1️⃣// MARK: - Marker2️⃣
+      """
+    ) { positions in
+      [
+        DocumentSymbol(
+          name: "- Marker",
+          kind: .namespace,
+          range: positions["1️⃣"]..<positions["2️⃣"],
+          selectionRange: positions["1️⃣"]..<positions["2️⃣"]
+        )
+      ]
+    }
+
+    try await assertDocumentSymbols(
+      """
+      1️⃣/* MARK: Marker */2️⃣
+      """
+    ) { positions in
+      [
+        DocumentSymbol(
+          name: "Marker",
+          kind: .namespace,
+          range: positions["1️⃣"]..<positions["2️⃣"],
+          selectionRange: positions["1️⃣"]..<positions["2️⃣"]
+        )
+      ]
+    }
+  }
+
+  func testIncludeNestedMarkComments() async throws {
+    try await assertDocumentSymbols(
+      """
+      1️⃣struct 2️⃣Foo3️⃣ {
+        4️⃣// MARK: Marker5️⃣
+      }6️⃣
+      """
+    ) { positions in
+      [
+        DocumentSymbol(
+          name: "Foo",
+          kind: .struct,
+          range: positions["1️⃣"]..<positions["6️⃣"],
+          selectionRange: positions["2️⃣"]..<positions["3️⃣"],
+          children: [
+            DocumentSymbol(
+              name: "Marker",
+              kind: .namespace,
+              range: positions["4️⃣"]..<positions["5️⃣"],
+              selectionRange: positions["4️⃣"]..<positions["5️⃣"]
+            )
+          ]
+        )
+      ]
+    }
+  }
 }
 
 fileprivate func assertDocumentSymbols(