Add an indented(by:includeFirstLine:) method to SyntaxProtocol

Author: j-f1

Sources/SwiftBasicFormat/Syntax+Extensions.swift

@@ -42,3 +42,87 @@ extension TokenSyntax {
     return []
   }
 }
+
+extension SyntaxProtocol {
+  /// Indent this node’s lines by the provided amount.
+  ///
+  /// - Parameter includeFirstLine: Whether the first token of this node should be indented.
+  ///   Pass `true` if you know that this node will be placed at the beginning of a line, even if its
+  ///   current leading trivia does not start with a newline (such as at the very start of a file).
+  public func indented(by indentation: Trivia, indentFirstLine: Bool = false) -> Syntax {
+    Indenter(indentation: indentation, indentFirstLine: indentFirstLine)
+      .rewrite(self)
+  }
+}
+
+private class Indenter: SyntaxRewriter {
+  private let indentation: Trivia
+  private var shouldIndent: Bool
+
+  init(indentation: Trivia, indentFirstLine: Bool) {
+    self.indentation = indentation
+    self.shouldIndent = indentFirstLine
+  }
+
+  private func indentIfNeeded() -> [TriviaPiece] {
+    if shouldIndent {
+      shouldIndent = false
+      return indentation.pieces
+    } else {
+      return []
+    }
+  }
+
+  private func indentAfterNewlines(_ content: String) -> String {
+    content.split(separator: "\n").joined(separator: "\n" + indentation.description)
+  }
+
+  private func indent(_ trivia: Trivia, skipEmpty: Bool) -> Trivia {
+    if skipEmpty, trivia.isEmpty { return trivia }
+
+    var result: [TriviaPiece] = []
+    // most times, we won’t have anything to insert so this will
+    // reserve enough space
+    result.reserveCapacity(trivia.count)
+
+    for piece in trivia.pieces {
+      result.append(contentsOf: indentIfNeeded())
+      switch piece {
+      case .newlines, .carriageReturns, .carriageReturnLineFeeds:
+        shouldIndent = true
+        // style decision: don’t indent totally blank lines
+        result.append(piece)
+      case .blockComment(let content):
+        result.append(.blockComment(indentAfterNewlines(content)))
+      case .docBlockComment(let content):
+        result.append(.docBlockComment(indentAfterNewlines(content)))
+      case .unexpectedText(let content):
+        result.append(.unexpectedText(indentAfterNewlines(content)))
+      default:
+        result.append(piece)
+      }
+    }
+    result.append(contentsOf: indentIfNeeded())
+    return Trivia(pieces: result)
+  }
+
+  override func visit(_ token: TokenSyntax) -> TokenSyntax {
+    let indentedLeadingTrivia = indent(token.leadingTrivia, skipEmpty: false)
+
+    // compute this before indenting the trailing trivia since the
+    // newline here is before the start of the trailing trivia (since
+    // it is part of the string’s value)
+    if case .stringSegment(let content) = token.tokenKind,
+       let last = content.last,
+       last.isNewline {
+      shouldIndent = true
+    }
+
+    return token
+      .with(\.leadingTrivia, indentedLeadingTrivia)
+      // source files as parsed can’t have anything requiring indentation
+      // here, but it’s easy to do `.with(\.trailingTrivia, .newline)` so
+      // we should still check if there’s something to indent.
+      .with(\.trailingTrivia, indent(token.trailingTrivia, skipEmpty: true))
+  }
+}

Sources/SwiftSyntax/generated/TriviaPieces.swift

@@ -31,11 +31,11 @@ public enum TriviaPiece: Sendable {
   case carriageReturnLineFeeds(Int)
   /// A documentation block comment, starting with '/**' and ending with '*/'.
   case docBlockComment(String)
-  /// A documentation line comment, starting with '///'.
+  /// A documentation line comment, starting with '///' and excluding the trailing newline.
   case docLineComment(String)
   /// A form-feed 'f' character.
   case formfeeds(Int)
-  /// A developer line comment, starting with '//'
+  /// A developer line comment, starting with '//' and excluding the trailing newline.
   case lineComment(String)
   /// A newline '\n' character.
   case newlines(Int)

Tests/SwiftBasicFormatTest/IndentTests.swift

@@ -0,0 +1,187 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 SwiftBasicFormat
+import SwiftParser
+import SwiftSyntax
+@_spi(Testing) import SwiftSyntaxBuilder
+import XCTest
+import _SwiftSyntaxTestSupport
+
+fileprivate func assertIndented(
+  by indentation: Trivia = .tab,
+  indentFirstLine: Bool = true,
+  source: String,
+  expected: String,
+  file: StaticString = #filePath,
+  line: UInt = #line
+) {
+  assertStringsEqualWithDiff(Parser.parse(source: source).indented(by: indentation, indentFirstLine: indentFirstLine).description, expected, file: file, line: line)
+}
+
+final class IndentTests: XCTestCase {
+  func testNotIndented() {
+    assertIndented(
+      source: """
+        func foo() {
+          let bar = 2
+        }
+        """,
+      expected: """
+        \tfunc foo() {
+        \t  let bar = 2
+        \t}
+        """
+    )
+  }
+
+  func testSingleLineComments() {
+    assertIndented(
+      source: """
+        func foo() {
+          // This is a comment
+          // that extends onto
+          // multiple lines \\
+          let bar = 2
+          // and another one
+        }
+        """,
+      expected: """
+        \tfunc foo() {
+        \t  // This is a comment
+        \t  // that extends onto
+        \t  // multiple lines \\
+        \t  let bar = 2
+        \t  // and another one
+        \t}
+        """
+    )
+  }
+
+  func testMultiLineComments() {
+    assertIndented(
+      source: """
+        func foo() {
+          /* This is a multiline comment
+                that extends onto
+        multiple lines*/
+          let bar = 2
+          /* on a single line */
+          let another = "Hello, world!" /* on a single line */
+        }
+        """,
+      expected: """
+        \tfunc foo() {
+        \t  /* This is a multiline comment
+        \t        that extends onto
+        \tmultiple lines*/
+        \t  let bar = 2
+        \t  /* on a single line */
+        \t  let another = "Hello, world!" /* on a single line */
+        \t}
+        """
+    )
+  }
+
+  func testMultiLineString() {
+    assertIndented(
+      source: #"""
+        func foo() {
+          let page = """
+            <h1>Hello, world!</h1>
+            <p>This is my web site</p>
+            """
+        }
+        """#,
+      expected: #"""
+        \#tfunc foo() {
+        \#t  let page = """
+        \#t    <h1>Hello, world!</h1>
+        \#t    <p>This is my web site</p>
+        \#t    """
+        \#t}
+        """#
+    )
+  }
+
+  func testIndented() {
+    assertIndented(
+      source: """
+          func foo() {
+            let bar = 2
+          }
+        """,
+      expected: """
+        \t  func foo() {
+        \t    let bar = 2
+        \t  }
+        """
+    )
+    assertIndented(
+      source: """
+        \tfunc foo() {
+        \t  let bar = 2
+        \t}
+        """,
+      expected: """
+        \t\tfunc foo() {
+        \t\t  let bar = 2
+        \t\t}
+        """
+    )
+  }
+
+  func testIndentBySpaces() {
+    assertIndented(
+      by: .spaces(4),
+      source: """
+        func foo() {
+          let bar = 2
+        }
+        """,
+      expected: """
+            func foo() {
+              let bar = 2
+            }
+        """
+    )
+  }
+
+  func testSkipFirstLine() {
+    assertIndented(
+      indentFirstLine: false,
+      source: """
+        \nfunc foo() {
+          let bar = 2
+        }
+        """,
+      expected: """
+        \n\tfunc foo() {
+        \t  let bar = 2
+        \t}
+        """
+    )
+    assertIndented(
+      indentFirstLine: false,
+      source: """
+        func foo() {
+          let bar = 2
+        }
+        """,
+      expected: """
+        func foo() {
+        \t  let bar = 2
+        \t}
+        """
+    )
+  }
+}