Report MARK comments in the document symbols request

Fixes #963
rdar://117811210

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 {
@@ -235,13 +275,13 @@ fileprivate extension SyntaxProtocol {
   var isMemberOrTopLevelDeclaration: Bool {
     if self.parent?.is(MemberBlockItemSyntax.self) ?? false {
       return true
-    } else if let codeBlockItem = self.parent?.as(CodeBlockItemSyntax.self),
+    }
+    if let codeBlockItem = self.parent?.as(CodeBlockItemSyntax.self),
       codeBlockItem.parent?.parent?.is(SourceFileSyntax.self) ?? false
     {
       return true
-    } else {
-      return false
     }
+    return false
   }
 }
 

Tests/SourceKitLSPTests/DocumentSymbolTests.swift

@@ -620,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(