Merge pull request swiftlang#104 from dabelknap/cleanup-v2

Fix `newline` token length calculation

Author: dabelknap

Sources/PrettyPrint/PrettyPrint.swift

@@ -47,14 +47,15 @@ public class PrettyPrinter {
   // Did the previous token trigger a newline?
   private var lastBreak = false
 
+  // What is the offset value of the last triggering break?
+  private var lastBreakOffset = 0
+
+  // What is the total number of spaces we need to indent from the last break?
+  private var lastBreakValue = 0
+
   // Keep track of the indentation level of the current group as the number of blank spaces.
   private var indentStack = [0]
 
-  // The first token of a group must be indented according to it's outer group. Use this stack when
-  // printing the indent of the first token in a group. Subsequent tokens will use the indent of the
-  // current group.
-  private var breakStack = [0]
-
   // Do we force breaks (for consistent breaking) within the current group?
   private var forceBreakStack = [false]
 
@@ -110,17 +111,15 @@ public class PrettyPrinter {
       }
 
       let indentValue = indentStack.last ?? 0
-      let breakValue = breakStack.last ?? 0
-      indentStack.append(indentValue + offset)
-      breakStack.append(breakValue)
+      indentStack.append(indentValue + offset + lastBreakOffset)
+      lastBreakOffset = 0
 
     case .close:
       if isDebugMode {
         writeCloseGroupDebugMarker()
       }
       forceBreakStack.removeLast()
       indentStack.removeLast()
-      breakStack.removeLast()
 
     case .break(let size, let offset):
       if isDebugMode {
@@ -135,35 +134,41 @@ public class PrettyPrinter {
       if length > spaceRemaining || forcebreak {
         // Update the top of the breakStack to reflect the indentation level of the current group.
         let indentValue = indentStack.last ?? 0
-        breakStack.removeLast()
-        breakStack.append(indentValue + offset)
 
         spaceRemaining = maxLineLength - indentValue - offset
         write("\n")
+
         lastBreak = true
+        lastBreakOffset = offset
+        lastBreakValue = indentValue + offset
       } else {
         writeSpaces(size)
         spaceRemaining -= size
+
         lastBreak = false
+        lastBreakOffset = 0
+        lastBreakValue = 0
       }
 
-    case .newlines(let N):
+    case .newlines(let N, let offset):
       // Newlines are treated as forced `breaks` with a size of zero.
       let indentValue = indentStack.last ?? 0
-      breakStack.removeLast()
-      breakStack.append(indentValue)
 
-      spaceRemaining = maxLineLength - indentValue
+      spaceRemaining = maxLineLength - indentValue - offset
       write(String(repeating: "\n", count: N))
-      lastBreak = true
 
+      lastBreak = true
+      lastBreakOffset = offset
+      lastBreakValue = indentValue + offset
 
     case .syntax(let syntaxToken):
       if lastBreak {
-        // If the last token created  a newline, we need to apply indentation.
-        let indentValue = breakStack.last ?? 0
-        writeSpaces(indentValue)
+        // If the last token created a new line, we need to apply indentation.
+        writeSpaces(lastBreakValue)
+
         lastBreak = false
+        lastBreakOffset = 0
+        lastBreakValue = 0
       }
       write(syntaxToken.text)
       spaceRemaining -= syntaxToken.text.count
@@ -223,9 +228,7 @@ public class PrettyPrinter {
           delimIndexStack.removeLast()
         }
 
-        lengths.append(-total)
-        delimIndexStack.append(i)
-        total += 1
+        lengths.append(0)
 
       case .syntax(let syntaxToken):
         lengths.append(syntaxToken.text.count)

Sources/PrettyPrint/Token.swift

@@ -28,13 +28,17 @@ enum BreakStyle {
 
 enum Token {
   case comment(Comment, hasTrailingSpace: Bool)
-  case newlines(Int)
+  case newlines(Int, offset: Int)
   case `break`(size: Int, offset: Int)
   case open(BreakStyle, Int)
   case close
   case syntax(TokenSyntax)
 
-  static let newline = Token.newlines(1)
+  static let newline = Token.newlines(1, offset: 0)
+  static func newline(offset: Int) -> Token {
+    return Token.newlines(1, offset: offset)
+  }
+
   static let open = Token.open(.inconsistent, 0)
 
   static let `break` = Token.break(size: 1, offset: 0)

Sources/PrettyPrint/TokenStreamCreator.swift

@@ -169,12 +169,15 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: ParameterClauseSyntax) {
-    after(node.leftParen, tokens: .break(size: 0), .open(.consistent, 0))
-    before(node.rightParen, tokens: .close)
+    after(node.leftParen, tokens: .break(size: 0, offset: 2), .open(.consistent, 0))
+    before(node.rightParen, tokens: .break(size: 0, offset: -2), .close)
     super.visit(node)
   }
 
   override func visit(_ node: ReturnClauseSyntax) {
+    before(node.firstToken, tokens: .open)
+    before(node.returnType.firstToken, tokens: .break)
+    after(node.lastToken, tokens: .close)
     super.visit(node)
   }
 
@@ -260,6 +263,16 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: CodeBlockSyntax) {
+    for i in 0..<(node.statements.count - 1) {
+      after(node.statements[i].lastToken, tokens: .newline)
+    }
+    super.visit(node)
+  }
+
+  override func visit(_ node: CodeBlockItemSyntax) {
+    if !(node.parent?.parent is CodeBlockSyntax) {
+      after(node.lastToken, tokens: .newline)
+    }
     super.visit(node)
   }
 
@@ -315,17 +328,17 @@ private final class TokenStreamCreator: SyntaxVisitor {
   override func visit(_ node: IfStmtSyntax) {
     before(node.ifKeyword, tokens: .open(.inconsistent, 3))
     after(node.ifKeyword, tokens: .break)
-    before(node.body.leftBrace, tokens: .break, .close)
+    before(node.body.leftBrace, tokens: .break(offset: -3), .close)
 
-    after(node.body.leftBrace, tokens: .open(.consistent, 2), .newline)
-    before(node.body.rightBrace, tokens: .close)
+    after(node.body.leftBrace, tokens: .newline(offset: 2), .open(.consistent, 0))
+    before(node.body.rightBrace, tokens: .newline(offset: -2), .close)
 
     before(node.elseKeyword, tokens: .break)
     after(node.elseKeyword, tokens: .break)
 
     if let elseBody = node.elseBody as? CodeBlockSyntax {
-      after(elseBody.leftBrace, tokens: .open(.consistent, 2), .newline)
-      before(elseBody.rightBrace, tokens: .close)
+      after(elseBody.leftBrace, tokens: .newline(offset: 2), .open(.consistent, 0))
+      before(elseBody.rightBrace, tokens: .newline(offset: -2), .close)
     }
     super.visit(node)
   }
@@ -395,7 +408,9 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: ReturnStmtSyntax) {
+    before(node.firstToken, tokens: .open)
     after(node.returnKeyword, tokens: .break)
+    after(node.lastToken, tokens: .close)
     super.visit(node)
   }
 
@@ -448,30 +463,20 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: FunctionDeclSyntax) {
-    if let token = node.firstToken {
-      before(token, tokens: .open(.inconsistent, 2))
-    }
-    before(node.signature.input.rightParen, tokens: .break(size: 0), .close)
     after(node.modifiers?.lastToken, tokens: .break)
     after(node.funcKeyword, tokens: .break)
 
     if let body = node.body {
       before(body.leftBrace, tokens: .break)
-      after(body.leftBrace, tokens: .open(.consistent, 2), .newline)
-      before(body.rightBrace, tokens: .close)
+      after(body.leftBrace, tokens: .newline(offset: 2), .open(.consistent, 0))
+      before(body.rightBrace, tokens: .newline(offset: -2), .close)
     }
 
     super.visit(node)
   }
 
   override func visit(_ node: FunctionSignatureSyntax) {
-    if node.output != nil {
-      after(node.input.rightParen, tokens: .break)
-    }
-    before(node.output?.arrow, tokens: .open(.consistent, 2))
-    after(node.output?.arrow, tokens: .break)
-    after(node.output?.returnType.lastToken, tokens: .close)
-
+    before(node.output?.firstToken, tokens: .break)
     super.visit(node)
   }
 
@@ -513,11 +518,6 @@ private final class TokenStreamCreator: SyntaxVisitor {
     super.visit(node)
   }
 
-  override func visit(_ node: CodeBlockItemSyntax) {
-    after(node.lastToken, tokens: .newline)
-    super.visit(node)
-  }
-
   override func visit(_ node: ExtensionDeclSyntax) {
     super.visit(node)
   }
@@ -905,7 +905,7 @@ private final class TokenStreamCreator: SyntaxVisitor {
         }
       case .newlines(let n), .carriageReturns(let n), .carriageReturnLineFeeds(let n):
         if n > 1 {
-          appendToken(.newlines(min(n - 1, config.maximumBlankLines)))
+          appendToken(.newlines(min(n - 1, config.maximumBlankLines), offset: 0))
         }
       default:
         break

Tests/PrettyPrinterTests/IfStmtTests.swift

@@ -88,8 +88,8 @@ public class IfStmtTests: PrettyPrintTestCase {
 
       if var1 < var2 {
         let a = 23
-      } else if var3 <
-         var4 {
+      } else
+      if var3 < var4 {
         var b = 123
       }