From ba41b21163722e3e3c76adf2b32751bb6d79b323 Mon Sep 17 00:00:00 2001
From: Alex Hoppen <ahoppen@apple.com>
Date: Mon, 19 Feb 2024 13:35:55 -0800
Subject: [PATCH] =?UTF-8?q?Don=E2=80=99t=20display=20MARK=20comments=20as?=
 =?UTF-8?q?=20children=20of=20the=20nodes=20they=20are=20attached=20to?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

We used to walk the leading and trailing trivia when walking the token that’s a child of the node. This caused us to eg. report the `MARK` comment in the following as a child of `myFunc` instead of a as a sibling of `myFunc`.

```swift
struct Foo {
  // MARK: Marker
  func myFunc() {}
}
```

rdar://121600714
---
 .../SourceKitLSP/Swift/DocumentSymbols.swift  | 42 +++++++++++++++----
 .../DocumentSymbolTests.swift                 | 37 +++++++++++++++-
 2 files changed, 71 insertions(+), 8 deletions(-)

diff --git a/Sources/SourceKitLSP/Swift/DocumentSymbols.swift b/Sources/SourceKitLSP/Swift/DocumentSymbols.swift
index e3092c1c..26e85112 100644
--- a/Sources/SourceKitLSP/Swift/DocumentSymbols.swift
+++ b/Sources/SourceKitLSP/Swift/DocumentSymbols.swift
@@ -22,7 +22,13 @@ extension SwiftLanguageServer {
     let syntaxTree = await syntaxTreeManager.syntaxTree(for: snapshot)
 
     try Task.checkCancellation()
-    return .documentSymbols(DocumentSymbolsFinder.find(in: [Syntax(syntaxTree)], snapshot: snapshot))
+    return .documentSymbols(
+      DocumentSymbolsFinder.find(
+        in: [Syntax(syntaxTree)],
+        snapshot: snapshot,
+        range: syntaxTree.position..<syntaxTree.endPosition
+      )
+    )
   }
 }
 
@@ -32,17 +38,25 @@ fileprivate final class DocumentSymbolsFinder: SyntaxAnyVisitor {
   /// The snapshot of the document for which we are getting document symbols.
   private let snapshot: DocumentSnapshot
 
+  /// Only document symbols that intersect with this range get reported.
+  private let range: Range<AbsolutePosition>
+
   /// Accumulating the result in here.
   private var result: [DocumentSymbol] = []
 
-  private init(snapshot: DocumentSnapshot) {
+  private init(snapshot: DocumentSnapshot, range: Range<AbsolutePosition>) {
     self.snapshot = snapshot
+    self.range = range
     super.init(viewMode: .sourceAccurate)
   }
 
   /// Designated entry point for `DocumentSymbolFinder`.
-  static func find(in nodes: some Sequence<Syntax>, snapshot: DocumentSnapshot) -> [DocumentSymbol] {
-    let visitor = Self(snapshot: snapshot)
+  static func find(
+    in nodes: some Sequence<Syntax>,
+    snapshot: DocumentSnapshot,
+    range: Range<AbsolutePosition>
+  ) -> [DocumentSymbol] {
+    let visitor = DocumentSymbolsFinder(snapshot: snapshot, range: range)
     for node in nodes {
       visitor.walk(node)
     }
@@ -57,6 +71,9 @@ fileprivate final class DocumentSymbolsFinder: SyntaxAnyVisitor {
     range: Range<AbsolutePosition>,
     selection: Range<AbsolutePosition>
   ) -> SyntaxVisitorContinueKind {
+    if !self.range.overlaps(range) {
+      return .skipChildren
+    }
     guard let rangeLowerBound = snapshot.position(of: range.lowerBound),
       let rangeUpperBound = snapshot.position(of: range.upperBound),
       let selectionLowerBound = snapshot.position(of: selection.lowerBound),
@@ -65,8 +82,14 @@ fileprivate final class DocumentSymbolsFinder: SyntaxAnyVisitor {
       return .skipChildren
     }
 
-    let children = DocumentSymbolsFinder.find(in: node.children(viewMode: .sourceAccurate), snapshot: snapshot)
+    // Record MARK comments on the node's leading and trailing trivia in `result` not as a child of `node`.
+    visit(node.leadingTrivia, position: node.position)
 
+    let children = DocumentSymbolsFinder.find(
+      in: node.children(viewMode: .sourceAccurate),
+      snapshot: snapshot,
+      range: node.positionAfterSkippingLeadingTrivia..<node.endPositionBeforeTrailingTrivia
+    )
     result.append(
       DocumentSymbol(
         name: name,
@@ -76,6 +99,7 @@ fileprivate final class DocumentSymbolsFinder: SyntaxAnyVisitor {
         children: children
       )
     )
+    visit(node.trailingTrivia, position: node.endPositionBeforeTrailingTrivia)
     return .skipChildren
   }
 
@@ -144,8 +168,12 @@ fileprivate final class DocumentSymbolsFinder: SyntaxAnyVisitor {
   }
 
   override func visit(_ node: TokenSyntax) -> SyntaxVisitorContinueKind {
-    self.visit(node.leadingTrivia, position: node.position)
-    self.visit(node.trailingTrivia, position: node.endPositionBeforeTrailingTrivia)
+    if self.range.overlaps(node.position..<node.positionAfterSkippingLeadingTrivia) {
+      self.visit(node.leadingTrivia, position: node.position)
+    }
+    if range.overlaps(node.endPositionBeforeTrailingTrivia..<node.endPosition) {
+      self.visit(node.trailingTrivia, position: node.endPositionBeforeTrailingTrivia)
+    }
     return .skipChildren
   }
 
diff --git a/Tests/SourceKitLSPTests/DocumentSymbolTests.swift b/Tests/SourceKitLSPTests/DocumentSymbolTests.swift
index 263d6396..f41020f5 100644
--- a/Tests/SourceKitLSPTests/DocumentSymbolTests.swift
+++ b/Tests/SourceKitLSPTests/DocumentSymbolTests.swift
@@ -668,7 +668,7 @@ final class DocumentSymbolTests: XCTestCase {
     }
   }
 
-  func testIncludeNestedMarkComments() async throws {
+  func testNestedMarkComment() async throws {
     try await assertDocumentSymbols(
       """
       1️⃣struct 2️⃣Foo3️⃣ {
@@ -694,6 +694,41 @@ final class DocumentSymbolTests: XCTestCase {
       ]
     }
   }
+
+  func testNestedMarkCommentFollowedAttachedToChild() async throws {
+    try await assertDocumentSymbols(
+      """
+      1️⃣struct 2️⃣Foo3️⃣ {
+        4️⃣// MARK: Marker5️⃣
+        6️⃣func 7️⃣myFunc()8️⃣  { }9️⃣
+      }🔟
+      """
+    ) { positions in
+      [
+        DocumentSymbol(
+          name: "Foo",
+          kind: .struct,
+          range: positions["1️⃣"]..<positions["🔟"],
+          selectionRange: positions["2️⃣"]..<positions["3️⃣"],
+          children: [
+            DocumentSymbol(
+              name: "Marker",
+              kind: .namespace,
+              range: positions["4️⃣"]..<positions["5️⃣"],
+              selectionRange: positions["4️⃣"]..<positions["5️⃣"]
+            ),
+            DocumentSymbol(
+              name: "myFunc()",
+              kind: .method,
+              range: positions["6️⃣"]..<positions["9️⃣"],
+              selectionRange: positions["7️⃣"]..<positions["8️⃣"],
+              children: []
+            ),
+          ]
+        )
+      ]
+    }
+  }
 }
 
 fileprivate func assertDocumentSymbols(