Merge pull request swiftlang#143 from dabelknap/fix-containers

Fix wrapping for containers with a trailing comma

Author: allevato

Sources/SwiftFormatPrettyPrint/TokenStreamCreator.swift

@@ -77,6 +77,17 @@ private final class TokenStreamCreator: SyntaxVisitor {
     afterMap[tok, default: []].append(tokens)
   }
 
+  private func insertToken<Node: Collection>(
+    _ token: Token,
+    betweenChildrenOf collectionNode: Node
+  ) where Node.Element: Syntax, Node.Index == Int {
+    if collectionNode.count > 0 {
+      for i in 0..<(collectionNode.count - 1) {
+        after(collectionNode[i].lastToken, tokens: token)
+      }
+    }
+  }
+
   override func visitPre(_ node: Syntax) {}
 
   override func visit(_ node: DeclNameArgumentsSyntax) {
@@ -99,13 +110,14 @@ private final class TokenStreamCreator: SyntaxVisitor {
     super.visit(node)
   }
 
+  override func visit(_ node: TupleElementListSyntax) {
+    insertToken(.break, betweenChildrenOf: node)
+    super.visit(node)
+  }
+
   override func visit(_ node: TupleElementSyntax) {
     before(node.firstToken, tokens: .open)
-    if let trailingComma = node.trailingComma {
-      after(trailingComma, tokens: .close, .break)
-    } else {
-      after(node.lastToken, tokens: .close)
-    }
+    after(node.lastToken, tokens: .close)
     super.visit(node)
   }
 
@@ -119,6 +131,17 @@ private final class TokenStreamCreator: SyntaxVisitor {
     super.visit(node)
   }
 
+  override func visit(_ node: ArrayElementListSyntax) {
+    insertToken(.break, betweenChildrenOf: node)
+    super.visit(node)
+  }
+
+  override func visit(_ node: ArrayElementSyntax) {
+    before(node.firstToken, tokens: .open)
+    after(node.lastToken, tokens: .close)
+    super.visit(node)
+  }
+
   override func visit(_ node: DictionaryExprSyntax) {
     after(
       node.leftSquare,
@@ -129,6 +152,11 @@ private final class TokenStreamCreator: SyntaxVisitor {
     super.visit(node)
   }
 
+  override func visit(_ node: DictionaryElementListSyntax) {
+    insertToken(.break, betweenChildrenOf: node)
+    super.visit(node)
+  }
+
   override func visit(_ node: DictionaryTypeSyntax) {
     after(node.colon, tokens: .space)
     super.visit(node)
@@ -137,11 +165,7 @@ private final class TokenStreamCreator: SyntaxVisitor {
   override func visit(_ node: DictionaryElementSyntax) {
     before(node.firstToken, tokens: .open)
     after(node.colon, tokens: .break(offset: 2))
-    if let trailingComma = node.trailingComma {
-      after(trailingComma, tokens: .close, .break)
-    } else {
-      after(node.lastToken, tokens: .close)
-    }
+    after(node.lastToken, tokens: .close)
     super.visit(node)
   }
 
@@ -398,11 +422,7 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: CodeBlockSyntax) {
-    if node.statements.count > 0 {
-      for i in 0..<(node.statements.count - 1) {
-        after(node.statements[i].lastToken, tokens: .newline)
-      }
-    }
+    insertToken(.newline, betweenChildrenOf: node.statements)
     super.visit(node)
   }
 
@@ -441,9 +461,7 @@ private final class TokenStreamCreator: SyntaxVisitor {
   override func visit(_ node: SwitchCaseSyntax) {
     before(node.firstToken, tokens: .open)
     after(node.label.lastToken, tokens: .newline(offset: 2), .open(.consistent, 0))
-    for i in 0..<(node.statements.count - 1) {
-      after(node.statements[i].lastToken, tokens: .newline)
-    }
+    insertToken(.newline, betweenChildrenOf: node.statements)
     after(node.lastToken, tokens: .break(offset: -2), .close, .close)
     super.visit(node)
   }
@@ -757,16 +775,6 @@ private final class TokenStreamCreator: SyntaxVisitor {
     super.visit(node)
   }
 
-  override func visit(_ node: ArrayElementSyntax) {
-    before(node.firstToken, tokens: .open)
-    if let trailingComma = node.trailingComma {
-      after(trailingComma, tokens: .close, .break)
-    } else {
-      after(node.lastToken, tokens: .close)
-    }
-    super.visit(node)
-  }
-
   override func visit(_ node: ContinueStmtSyntax) {
     super.visit(node)
   }

Tests/SwiftFormatPrettyPrintTests/ArrayDeclTests.swift

@@ -6,6 +6,7 @@ public class ArrayDeclTests: PrettyPrintTestCase {
       let a: [Bool] = [false, true, true, false]
       let a: [String] = ["One", "Two", "Three", "Four"]
       let a: [String] = ["One", "Two", "Three", "Four", "Five", "Six", "Seven"]
+      let a: [String] = ["One", "Two", "Three", "Four", "Five", "Six", "Seven",]
       """
 
     let expected =
@@ -24,6 +25,15 @@ public class ArrayDeclTests: PrettyPrintTestCase {
         "Six",
         "Seven"
       ]
+      let a: [String] = [
+        "One",
+        "Two",
+        "Three",
+        "Four",
+        "Five",
+        "Six",
+        "Seven",
+      ]
 
       """
 

Tests/SwiftFormatPrettyPrintTests/DictionaryDeclTests.swift

@@ -6,6 +6,7 @@ public class DictionaryDeclTests: PrettyPrintTestCase {
       let a: [Int: String] = [1: "a", 2: "b", 3: "c"]
       let a: [Int: String] = [1: "a", 2: "b", 3: "c", 4: "d"]
       let a: [Int: String] = [1: "a", 2: "b", 3: "c", 4: "d", 5: "e", 6: "f", 7: "g"]
+      let a: [Int: String] = [1: "a", 2: "b", 3: "c", 4: "d", 5: "e", 6: "f", 7: "g",]
       """
 
     let expected =
@@ -24,6 +25,15 @@ public class DictionaryDeclTests: PrettyPrintTestCase {
         6: "f",
         7: "g"
       ]
+      let a: [Int: String] = [
+        1: "a",
+        2: "b",
+        3: "c",
+        4: "d",
+        5: "e",
+        6: "f",
+        7: "g",
+      ]
 
       """
 

Tests/SwiftFormatPrettyPrintTests/TupleDeclTests.swift

@@ -6,6 +6,7 @@ public class TupleDeclTests: PrettyPrintTestCase {
       let a: (Int, Int, Int) = (1, 2, 3)
       let a = (1, 2, 3, 4, 5, 6, 7, 8, 9, 10)
       let a = (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12)
+      let a = (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12,)
       """
 
     let expected =
@@ -29,6 +30,20 @@ public class TupleDeclTests: PrettyPrintTestCase {
         11,
         12
       )
+      let a = (
+        1,
+        2,
+        3,
+        4,
+        5,
+        6,
+        7,
+        8,
+        9,
+        10,
+        11,
+        12,
+      )
 
       """