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,51 @@ fileprivate final class DocumentSymbolsFinder: SyntaxAnyVisitor {
     )
   }
 
+  private func visit(_ trivia: Trivia, position: AbsolutePosition) {
+    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.trimmingPrefix(while: { $0 == "/" || $0 == "*" || $0 == " " })
+        if trimmedComment.starts(with: "MARK: ") {
+          let markText = String(trimmedComment[trimmedComment.index(trimmedComment.startIndex, offsetBy: 6)...])
+          guard let rangeLowerBound = snapshot.position(of: position),
+            let rangeUpperBound = snapshot.position(of: position.advanced(by: piece.sourceLength.utf8Length))
+          else {
+            break
+          }
+          let documentSymbolName: String
+          switch piece {
+          case .lineComment: documentSymbolName = markText.trimmingCharacters(in: .whitespaces)
+          case .blockComment:
+            documentSymbolName = markText.trimmingCharacters(in: CharacterSet(["/", "*"]).union(.whitespaces))
+          default: preconditionFailure("Trivia piece kind not covered in the outer case")
+          }
+          result.append(
+            DocumentSymbol(
+              name: documentSymbolName,
+              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 {
@@ -253,3 +299,15 @@ fileprivate extension TokenKind {
     }
   }
 }
+
+fileprivate extension String {
+  /// Compatibility function of the `trimmingPrefix` function in the stdlib. We should be able to remove this once we
+  /// only support macOS 13+
+  func trimmingPrefix(while condition: (Character) -> Bool) -> Substring {
+    var result = self[...]
+    while let first = result.first, condition(first) {
+      result = result[result.index(after: result.startIndex)...]
+    }
+    return result
+  }
+}

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(