Minor fixes to BasicFormat that I found while using it to run CodeGeneration

Author: ahoppen

CodeGeneration/Sources/Utils/CodeGenerationFormat.swift

@@ -106,7 +106,7 @@ public class CodeGenerationFormat: BasicFormat {
   }
 
   private func formatChildrenSeparatedByNewline<SyntaxType: SyntaxProtocol>(children: SyntaxChildren, elementType: SyntaxType.Type) -> [SyntaxType] {
-    pushIndentationLevel(increasingIndentationBy: indentationWidth)
+    increaseIndentationLevel()
     var formattedChildren = children.map {
       self.visit($0).as(SyntaxType.self)!
     }
@@ -117,7 +117,7 @@ public class CodeGenerationFormat: BasicFormat {
         return $0.with(\.leadingTrivia, indentedNewline + $0.leadingTrivia)
       }
     }
-    popIndentationLevel()
+    decreaseIndentationLevel()
     if !formattedChildren.isEmpty {
       formattedChildren[formattedChildren.count - 1] = formattedChildren[formattedChildren.count - 1].with(\.trailingTrivia, indentedNewline)
     }

Sources/SwiftBasicFormat/BasicFormat.swift

@@ -18,12 +18,15 @@ open class BasicFormat: SyntaxRewriter {
 
   /// As we reach a new indendation level, its indentation will be added to the
   /// stack. As we exit that indentation level, the indendation will be popped.
-  public var indentationStack: [Trivia]
+  /// `isUserDefined` is `true` if the indentation was inferred from something
+  /// the user provided manually instead of being inferred from the nesting
+  /// level.
+  public var indentationStack: [(indentation: Trivia, isUserDefined: Bool)]
 
   /// The trivia by which tokens should currently be indented.
   public var currentIndentationLevel: Trivia {
     // `popIndentationLevel` guarantees that there is always one item on the stack.
-    return indentationStack.last!
+    return indentationStack.last!.indentation
   }
 
   /// For every token that is being put on a new line but did not have
@@ -39,16 +42,16 @@ open class BasicFormat: SyntaxRewriter {
 
   public init(indentationWidth: Trivia = .spaces(4), initialIndentation: Trivia = []) {
     self.indentationWidth = indentationWidth
-    self.indentationStack = [initialIndentation]
+    self.indentationStack = [(indentation: initialIndentation, isUserDefined: false)]
   }
 
   // MARK: - Updating indentation level
 
-  public func pushIndentationLevel(increasingIndentationBy: Trivia) {
-    indentationStack.append(currentIndentationLevel + increasingIndentationBy)
+  public func increaseIndentationLevel() {
+    indentationStack.append((currentIndentationLevel + indentationWidth, false))
   }
 
-  public func popIndentationLevel() {
+  public func decreaseIndentationLevel() {
     precondition(indentationStack.count > 1, "Popping more indentation levels than have been pushed")
     indentationStack.removeLast()
   }
@@ -60,16 +63,17 @@ open class BasicFormat: SyntaxRewriter {
         !tokenIndentation.isEmpty
       {
         // If the first token in this block is indented, infer the indentation level from it.
-        pushIndentationLevel(increasingIndentationBy: tokenIndentation)
+        let lastNonUserDefinedIndentation = indentationStack.last(where: { !$0.isUserDefined })!.indentation
+        indentationStack.append((indentation: lastNonUserDefinedIndentation + tokenIndentation, isUserDefined: true))
       } else {
-        pushIndentationLevel(increasingIndentationBy: indentationWidth)
+        increaseIndentationLevel()
       }
     }
   }
 
   open override func visitPost(_ node: Syntax) {
     if requiresIndent(node) {
-      popIndentationLevel()
+      decreaseIndentationLevel()
     }
   }
 
@@ -197,7 +201,25 @@ open class BasicFormat: SyntaxRewriter {
         // don't add a leading newline to the file.
         return true
       }
-      return previousToken.trailingTrivia.endsWithNewline
+      if previousToken.trailingTrivia.pieces.last?.isNewline ?? false {
+        return true
+      }
+      if case .stringSegment(let segment) = previousToken.tokenKind, segment.last?.isNewline ?? false {
+        return true
+      }
+      return false
+    }()
+
+    lazy var previousTokenIsStringLiteralEndingInNewline: Bool = {
+      guard let previousToken = previousToken else {
+        // Assume that the start of the tree is equivalent to a newline so we
+        // don't add a leading newline to the file.
+        return true
+      }
+      if case .stringSegment(let segment) = previousToken.tokenKind, segment.last?.isNewline ?? false {
+        return true
+      }
+      return false
     }()
 
     lazy var nextTokenWillStartWithNewline: Bool = {
@@ -270,7 +292,7 @@ open class BasicFormat: SyntaxRewriter {
       trailingTriviaIndentation = anchorPointIndentation
     }
 
-    leadingTrivia = leadingTrivia.indented(indentation: leadingTriviaIndentation, isOnNewline: false)
+    leadingTrivia = leadingTrivia.indented(indentation: leadingTriviaIndentation, isOnNewline: previousTokenIsStringLiteralEndingInNewline)
     trailingTrivia = trailingTrivia.indented(indentation: trailingTriviaIndentation, isOnNewline: false)
 
     leadingTrivia = leadingTrivia.trimmingTrailingWhitespaceBeforeNewline(isBeforeNewline: false)