Respect indentation in the line for sub-tree in BasicFormat

Sources/SwiftBasicFormat/BasicFormat.swift

@@ -150,6 +150,18 @@ open class BasicFormat: SyntaxRewriter {
 
   // MARK: - Customization points
 
+  /// If we are formatting a subtree and the line that the initial token occurs on is indented,
+  /// use that line indentation for the first token in the subtree to format.
+  ///
+  /// For example, when formatting only the code block in the following,
+  /// then the opening `{` should be indented by four spaces.
+  /// ```
+  ///     func test() {
+  ///         print(1)
+  ///     }
+  /// ```
+  open var inferInitialTokenIndentaiton: Bool { true }
+
   /// Whether a leading newline on `token` should be added.
   open func requiresIndent(_ node: some SyntaxProtocol) -> Bool {
     return node.requiresIndent
@@ -413,6 +425,16 @@ open class BasicFormat: SyntaxRewriter {
       anchorPoints[token] = currentIndentationLevel
     }
 
+    if inferInitialTokenIndentaiton
+      && isInitialToken
+      && token.presence == .present
+    {
+      let indentationOfLine = token.indentationOfLine
+      if token.leadingTrivia.pieces.suffix(indentationOfLine.pieces.count) != indentationOfLine.pieces {
+        leadingTrivia += indentationOfLine
+      }
+    }
+
     // Add a trailing space to the token unless
     //  - it already ends with a whitespace or
     //  - the next token will start starts with a newline after the rewrite

Sources/SwiftBasicFormat/CMakeLists.txt

@@ -9,6 +9,7 @@
 add_swift_host_library(SwiftBasicFormat
   BasicFormat.swift
   generated/BasicFormat+Extensions.swift
+  Syntax+Extensions.swift
   SyntaxProtocol+Formatted.swift
   Trivia+FormatExtensions.swift
 )

Sources/SwiftBasicFormat/Syntax+Extensions.swift

@@ -0,0 +1,40 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 SwiftSyntax
+
+extension TokenSyntax {
+  /// The indentation of this token
+  ///
+  /// In contrast to `indentationOfLine`, this does not walk to previous tokens to
+  /// find the indentation of the line this token occurs on.
+  private var indentation: Trivia? {
+    let previous = self.previousToken(viewMode: .sourceAccurate)
+    return ((previous?.trailingTrivia ?? []) + leadingTrivia).indentation(isOnNewline: false)
+  }
+
+  /// Returns the indentation of the line this token occurs on
+  public var indentationOfLine: Trivia {
+    var token: TokenSyntax = self
+    if let indentation = token.indentation {
+      return indentation
+    }
+    while let previous = token.previousToken(viewMode: .sourceAccurate) {
+      token = previous
+      if let indentation = token.indentation {
+        return indentation
+      }
+    }
+
+    return []
+  }
+}

Sources/SwiftParserDiagnostics/MissingNodesError.swift

@@ -22,6 +22,8 @@ fileprivate func findCommonAncestor(_ nodes: [Syntax]) -> Syntax? {
 }
 
 class NoNewlinesFormat: BasicFormat {
+  override var inferInitialTokenIndentaiton: Bool { false }
+
   override func requiresNewline(between first: TokenSyntax?, and second: TokenSyntax?) -> Bool {
     return false
   }

Sources/SwiftParserDiagnostics/ParseDiagnosticsGenerator.swift

@@ -51,31 +51,6 @@ fileprivate func getTokens(between first: TokenSyntax, and second: TokenSyntax)
 }
 
 fileprivate extension TokenSyntax {
-  /// The indentation of this token
-  ///
-  /// In contrast to `indentation`, this does not walk to previous tokens to
-  /// find the indentation of the line this token occurs on.
-  private var indentation: Trivia? {
-    let previous = self.previousToken(viewMode: .sourceAccurate)
-    return ((previous?.trailingTrivia ?? []) + leadingTrivia).indentation(isOnNewline: false)
-  }
-
-  /// Returns the indentation of the line this token occurs on
-  var indentationOfLine: Trivia {
-    var token: TokenSyntax = self
-    if let indentation = token.indentation {
-      return indentation
-    }
-    while let previous = token.previousToken(viewMode: .sourceAccurate) {
-      token = previous
-      if let indentation = token.indentation {
-        return indentation
-      }
-    }
-
-    return []
-  }
-
   /// Assuming this token is a `poundAvailableKeyword` or `poundUnavailableKeyword`
   /// returns the opposite keyword.
   var negatedAvailabilityKeyword: TokenSyntax {

Tests/SwiftBasicFormatTest/BasicFormatTests.swift

@@ -351,7 +351,7 @@ final class BasicFormatTest: XCTestCase {
     let body = decl.cast(FunctionDeclSyntax.self).body!
 
     assertFormatted(
-      source: body.formatted().description,
+      tree: body,
       expected: """
         {
             print(1)
@@ -372,9 +372,9 @@ final class BasicFormatTest: XCTestCase {
     let body = decl.cast(StructDeclSyntax.self).memberBlock.members.first!.decl.cast(FunctionDeclSyntax.self).body!
 
     assertFormatted(
-      source: body.formatted().description,
+      tree: body,
       expected: """
-        {
+            {
                 print(1)
             }
         """