Remove the break stack, and update the exisiting visitors.

Author: dabelknap

Sources/PrettyPrint/PrettyPrint.swift

@@ -47,16 +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]
 
@@ -112,9 +111,7 @@ public class PrettyPrinter {
       }
 
       let indentValue = indentStack.last ?? 0
-      let breakValue = breakStack.last ?? 0
       indentStack.append(indentValue + offset + lastBreakOffset)
-      breakStack.append(breakValue)
       lastBreakOffset = 0
 
     case .close:
@@ -123,7 +120,6 @@ public class PrettyPrinter {
       }
       forceBreakStack.removeLast()
       indentStack.removeLast()
-      breakStack.removeLast()
 
     case .break(let size, let offset):
       if isDebugMode {
@@ -138,38 +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, let offset):
       // Newlines are treated as forced `breaks` with a size of zero.
       let indentValue = indentStack.last ?? 0
-      breakStack.removeLast()
-      breakStack.append(indentValue + offset)
 
       spaceRemaining = maxLineLength - indentValue - offset
       write(String(repeating: "\n", count: N))
+
       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

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: .newline(offset: 2), .open(.consistent, 0))
-      before(body.rightBrace, tokens: .close)
+      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)
   }

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
       }