Merge pull request #999 from ahoppen/ahoppen/folding-range-improvements

Improve code folding

Author: ahoppen

Sources/SourceKitLSP/CMakeLists.txt

@@ -24,6 +24,7 @@ target_sources(SourceKitLSP PRIVATE
   Swift/Diagnostic.swift
   Swift/DocumentSymbols.swift
   Swift/EditorPlaceholder.swift
+  Swift/FoldingRange.swift
   Swift/OpenInterface.swift
   Swift/RelatedIdentifiers.swift
   Swift/SemanticRefactorCommand.swift

Sources/SourceKitLSP/Swift/FoldingRange.swift

@@ -0,0 +1,261 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import LSPLogging
+import LanguageServerProtocol
+import SwiftSyntax
+
+fileprivate final class FoldingRangeFinder: SyntaxAnyVisitor {
+  private let snapshot: DocumentSnapshot
+  /// Some ranges might occur multiple times.
+  /// E.g. for `print("hi")`, `"hi"` is both the range of all call arguments and the range the first argument in the call.
+  /// It doesn't make sense to report them multiple times, so use a `Set` here.
+  private var ranges: Set<FoldingRange>
+  /// The client-imposed limit on the number of folding ranges it would
+  /// prefer to receive from the LSP server. If the value is `nil`, there
+  /// is no preset limit.
+  private var rangeLimit: Int?
+  /// If `true`, the client is only capable of folding entire lines. If
+  /// `false` the client can handle folding ranges.
+  private var lineFoldingOnly: Bool
+
+  init(snapshot: DocumentSnapshot, rangeLimit: Int?, lineFoldingOnly: Bool) {
+    self.snapshot = snapshot
+    self.ranges = []
+    self.rangeLimit = rangeLimit
+    self.lineFoldingOnly = lineFoldingOnly
+    super.init(viewMode: .sourceAccurate)
+  }
+
+  override func visit(_ node: TokenSyntax) -> SyntaxVisitorContinueKind {
+    // Index comments, so we need to see at least '/*', or '//'.
+    if node.leadingTriviaLength.utf8Length > 2 {
+      self.addTrivia(from: node, node.leadingTrivia)
+    }
+
+    if node.trailingTriviaLength.utf8Length > 2 {
+      self.addTrivia(from: node, node.trailingTrivia)
+    }
+
+    return .visitChildren
+  }
+
+  private func addTrivia(from node: TokenSyntax, _ trivia: Trivia) {
+    let pieces = trivia.pieces
+    var start = node.position
+    /// The index of the trivia piece we are currently inspecting.
+    var index = 0
+
+    while index < pieces.count {
+      let piece = pieces[index]
+      defer {
+        start = start.advanced(by: pieces[index].sourceLength.utf8Length)
+        index += 1
+      }
+      switch piece {
+      case .blockComment:
+        _ = self.addFoldingRange(
+          start: start,
+          end: start.advanced(by: piece.sourceLength.utf8Length),
+          kind: .comment
+        )
+      case .docBlockComment:
+        _ = self.addFoldingRange(
+          start: start,
+          end: start.advanced(by: piece.sourceLength.utf8Length),
+          kind: .comment
+        )
+      case .lineComment, .docLineComment:
+        let lineCommentBlockStart = start
+
+        // Keep scanning the upcoming trivia pieces to find the end of the
+        // block of line comments.
+        // As we find a new end of the block comment, we set `index` and
+        // `start` to `lookaheadIndex` and `lookaheadStart` resp. to
+        // commit the newly found end.
+        var lookaheadIndex = index
+        var lookaheadStart = start
+        var hasSeenNewline = false
+        LOOP: while lookaheadIndex < pieces.count {
+          let piece = pieces[lookaheadIndex]
+          defer {
+            lookaheadIndex += 1
+            lookaheadStart = lookaheadStart.advanced(by: piece.sourceLength.utf8Length)
+          }
+          switch piece {
+          case .newlines(let count), .carriageReturns(let count), .carriageReturnLineFeeds(let count):
+            if count > 1 || hasSeenNewline {
+              // More than one newline is separating the two line comment blocks.
+              // We have reached the end of this block of line comments.
+              break LOOP
+            }
+            hasSeenNewline = true
+          case .spaces, .tabs:
+            // We allow spaces and tabs because the comments might be indented
+            continue
+          case .lineComment, .docLineComment:
+            // We have found a new line comment in this block. Commit it.
+            index = lookaheadIndex
+            start = lookaheadStart
+            hasSeenNewline = false
+          default:
+            // We assume that any other trivia piece terminates the block
+            // of line comments.
+            break LOOP
+          }
+        }
+        _ = self.addFoldingRange(
+          start: lineCommentBlockStart,
+          end: start.advanced(by: pieces[index].sourceLength.utf8Length),
+          kind: .comment
+        )
+      default:
+        break
+      }
+    }
+  }
+
+  override func visitAny(_ node: Syntax) -> SyntaxVisitorContinueKind {
+    if let braced = node.asProtocol(BracedSyntax.self) {
+      return self.addFoldingRange(
+        start: braced.leftBrace.endPositionBeforeTrailingTrivia,
+        end: braced.rightBrace.positionAfterSkippingLeadingTrivia
+      )
+    }
+    if let parenthesized = node.asProtocol(ParenthesizedSyntax.self) {
+      return self.addFoldingRange(
+        start: parenthesized.leftParen.endPositionBeforeTrailingTrivia,
+        end: parenthesized.rightParen.positionAfterSkippingLeadingTrivia
+      )
+    }
+    return .visitChildren
+  }
+
+  override func visit(_ node: ArrayExprSyntax) -> SyntaxVisitorContinueKind {
+    return self.addFoldingRange(
+      start: node.leftSquare.endPositionBeforeTrailingTrivia,
+      end: node.rightSquare.positionAfterSkippingLeadingTrivia
+    )
+  }
+
+  override func visit(_ node: DictionaryExprSyntax) -> SyntaxVisitorContinueKind {
+    return self.addFoldingRange(
+      start: node.leftSquare.endPositionBeforeTrailingTrivia,
+      end: node.rightSquare.positionAfterSkippingLeadingTrivia
+    )
+  }
+
+  override func visit(_ node: FunctionCallExprSyntax) -> SyntaxVisitorContinueKind {
+    if let leftParen = node.leftParen, let rightParen = node.rightParen {
+      return self.addFoldingRange(
+        start: leftParen.endPositionBeforeTrailingTrivia,
+        end: rightParen.positionAfterSkippingLeadingTrivia
+      )
+    }
+    return .visitChildren
+  }
+
+  override func visit(_ node: SubscriptCallExprSyntax) -> SyntaxVisitorContinueKind {
+    return self.addFoldingRange(
+      start: node.leftSquare.endPositionBeforeTrailingTrivia,
+      end: node.rightSquare.positionAfterSkippingLeadingTrivia
+    )
+  }
+
+  override func visit(_ node: SwitchCaseSyntax) -> SyntaxVisitorContinueKind {
+    return self.addFoldingRange(
+      start: node.label.endPositionBeforeTrailingTrivia,
+      end: node.statements.endPosition
+    )
+  }
+
+  __consuming func finalize() -> Set<FoldingRange> {
+    return self.ranges
+  }
+
+  private func addFoldingRange(
+    start: AbsolutePosition,
+    end: AbsolutePosition,
+    kind: FoldingRangeKind? = nil
+  ) -> SyntaxVisitorContinueKind {
+    if let limit = self.rangeLimit, self.ranges.count >= limit {
+      return .skipChildren
+    }
+    if start == end {
+      // Don't report empty ranges
+      return .visitChildren
+    }
+
+    guard let start: Position = snapshot.positionOf(utf8Offset: start.utf8Offset),
+      let end: Position = snapshot.positionOf(utf8Offset: end.utf8Offset)
+    else {
+      logger.error(
+        "folding range failed to retrieve position of \(self.snapshot.uri.forLogging): \(start.utf8Offset)-\(end.utf8Offset)"
+      )
+      return .visitChildren
+    }
+    let range: FoldingRange
+    if lineFoldingOnly {
+      // Since the client cannot fold less than a single line, if the
+      // fold would span 1 line there's no point in reporting it.
+      guard end.line > start.line else {
+        return .visitChildren
+      }
+
+      // If the client only supports folding full lines, don't report
+      // the end of the range since there's nothing they could do with it.
+      range = FoldingRange(
+        startLine: start.line,
+        startUTF16Index: nil,
+        endLine: end.line,
+        endUTF16Index: nil,
+        kind: kind
+      )
+    } else {
+      range = FoldingRange(
+        startLine: start.line,
+        startUTF16Index: start.utf16index,
+        endLine: end.line,
+        endUTF16Index: end.utf16index,
+        kind: kind
+      )
+    }
+    ranges.insert(range)
+    return .visitChildren
+  }
+}
+
+extension SwiftLanguageServer {
+  public func foldingRange(_ req: FoldingRangeRequest) async throws -> [FoldingRange]? {
+    let foldingRangeCapabilities = capabilityRegistry.clientCapabilities.textDocument?.foldingRange
+    let snapshot = try self.documentManager.latestSnapshot(req.textDocument.uri)
+
+    let sourceFile = await syntaxTreeManager.syntaxTree(for: snapshot)
+
+    try Task.checkCancellation()
+
+    // If the limit is less than one, do nothing.
+    if let limit = foldingRangeCapabilities?.rangeLimit, limit <= 0 {
+      return []
+    }
+
+    let rangeFinder = FoldingRangeFinder(
+      snapshot: snapshot,
+      rangeLimit: foldingRangeCapabilities?.rangeLimit,
+      lineFoldingOnly: foldingRangeCapabilities?.lineFoldingOnly ?? false
+    )
+    rangeFinder.walk(sourceFile)
+    let ranges = rangeFinder.finalize()
+
+    return ranges.sorted()
+  }
+}