Merge pull request swiftlang#119 from dabelknap/arrays

Add supports for Array declarations

Author: allevato

Sources/PrettyPrint/PrettyPrint.swift

@@ -167,13 +167,18 @@ public class PrettyPrinter {
         if !lastBreakConsecutive {
           writeSpaces(size)
           spaceRemaining -= size
+          lastBreakValue = 0
         }
 
         lastBreak = false
         lastBreakOffset = 0
-        lastBreakValue = 0
       }
 
+    // Print out the number of spaces according to the size, and adjust spaceRemaining.
+    case .space(let size):
+      spaceRemaining -= size
+      writeSpaces(size)
+
     // Apply N line breaks, calculate the indentation required, and adjust spaceRemaining.
     case .newlines(let N, let offset):
       let indentValue = indentStack.last ?? 0
@@ -259,6 +264,11 @@ public class PrettyPrinter {
         delimIndexStack.append(i)
         total += size
 
+      // Space tokens have a length equal to its size.
+      case .space(let size):
+        lengths.append(size)
+        total += size
+
       // The length of newlines are equal to the maximum allowed line length. Calculate the length
       // of any prior break tokens.
       case .newlines:

Sources/PrettyPrint/Token.swift

@@ -31,6 +31,7 @@ enum Token {
   case open(BreakStyle, Int)
   case close
   case `break`(size: Int, offset: Int)
+  case space(size: Int)
   case newlines(Int, offset: Int)
   case comment(Comment, hasTrailingSpace: Bool)
 
@@ -42,6 +43,8 @@ enum Token {
     return Token.newlines(1, offset: offset)
   }
 
+  static let space = Token.space(size: 1)
+
   static let `break` = Token.break(size: 1, offset: 0)
   static func `break`(offset: Int) -> Token {
     return Token.break(size: 1, offset: offset)

Sources/PrettyPrint/TokenStreamCreator.swift

@@ -86,6 +86,8 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: BinaryOperatorExprSyntax) {
+    before(node.operatorToken, tokens: .break)
+    after(node.operatorToken, tokens: .break)
     super.visit(node)
   }
 
@@ -94,6 +96,8 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: ArrayExprSyntax) {
+    after(node.leftSquare, tokens: .break(size: 0, offset: 2), .open(.consistent, 0), .break(size: 0), .open(.consistent, 0))
+    before(node.rightSquare, tokens: .close, .break(size: 0, offset: -2), .close)
     super.visit(node)
   }
 
@@ -288,7 +292,7 @@ private final class TokenStreamCreator: SyntaxVisitor {
 
   override func visit(_ node: SwitchStmtSyntax) {
     before(node.switchKeyword, tokens: .open(.inconsistent, 7))
-    after(node.switchKeyword, tokens: .break)
+    after(node.switchKeyword, tokens: .space)
     after(node.expression.lastToken, tokens: .close, .break)
     after(node.leftBrace, tokens: .newline, .open(.consistent, 0))
     before(node.rightBrace, tokens: .break, .close)
@@ -307,7 +311,7 @@ private final class TokenStreamCreator: SyntaxVisitor {
 
   override func visit(_ node: SwitchCaseLabelSyntax) {
     before(node.caseKeyword, tokens: .open(.inconsistent, 5))
-    after(node.caseKeyword, tokens: .break)
+    after(node.caseKeyword, tokens: .space)
     after(node.colon, tokens: .close)
     super.visit(node)
   }
@@ -366,7 +370,7 @@ private final class TokenStreamCreator: SyntaxVisitor {
 
   override func visit(_ node: IfStmtSyntax) {
     before(node.ifKeyword, tokens: .open(.inconsistent, 3))
-    after(node.ifKeyword, tokens: .break)
+    after(node.ifKeyword, tokens: .space)
     before(node.body.leftBrace, tokens: .break(offset: -3), .close)
 
     after(node.body.leftBrace, tokens: .newline(offset: 2), .open(.consistent, 0))
@@ -441,9 +445,9 @@ private final class TokenStreamCreator: SyntaxVisitor {
 
   override func visit(_ node: ForInStmtSyntax) {
     before(node.forKeyword, tokens: .open(.inconsistent, 4))
-    after(node.forKeyword, tokens: .break)
+    after(node.forKeyword, tokens: .space)
     before(node.inKeyword, tokens: .break)
-    after(node.inKeyword, tokens: .break)
+    after(node.inKeyword, tokens: .space)
 
     if let whereClause = node.whereClause {
       before(
@@ -487,8 +491,8 @@ private final class TokenStreamCreator: SyntaxVisitor {
 
   override func visit(_ node: WhileStmtSyntax) {
     before(node.firstToken, tokens: .open(.inconsistent, 6))
-    after(node.labelColon, tokens: .break)
-    after(node.whileKeyword, tokens: .break)
+    after(node.labelColon, tokens: .space)
+    after(node.whileKeyword, tokens: .space)
     before(node.body.leftBrace, tokens: .break(offset: -6), .close)
     after(node.body.leftBrace, tokens: .break(offset: 2), .open(.consistent, 0))
     before(node.body.rightBrace, tokens: .break(offset: -2), .close)
@@ -543,7 +547,7 @@ private final class TokenStreamCreator: SyntaxVisitor {
 
   override func visit(_ node: WhereClauseSyntax) {
     before(node.whereKeyword, tokens: .open(.inconsistent, 2))
-    after(node.whereKeyword, tokens: .break)
+    after(node.whereKeyword, tokens: .space)
     after(node.lastToken, tokens: .close)
     super.visit(node)
   }
@@ -553,6 +557,12 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   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)
   }
 
@@ -607,9 +617,8 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: SequenceExprSyntax) {
-    for index in 0..<(node.elements.count - 1) {
-      after(node.elements[index].lastToken, tokens: .break)
-    }
+    before(node.firstToken, tokens: .open)
+    after(node.lastToken, tokens: .close)
     super.visit(node)
   }
 
@@ -623,7 +632,7 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: VariableDeclSyntax) {
-    before(node.firstToken, tokens: .open(.inconsistent, 2))
+    before(node.firstToken, tokens: .open(.inconsistent, 0))
     after(node.lastToken, tokens: .close)
     after(node.letOrVarKeyword, tokens: .break)
     super.visit(node)
@@ -718,7 +727,7 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: TypeAnnotationSyntax) {
-    after(node.colon, tokens: .break)
+    after(node.colon, tokens: .break(offset: 2))
     super.visit(node)
   }
 
@@ -772,8 +781,8 @@ private final class TokenStreamCreator: SyntaxVisitor {
     after(node.repeatKeyword, tokens: .break)
     after(node.body.leftBrace, tokens: .break(offset: 2), .open(.consistent, 0))
     before(node.body.rightBrace, tokens: .break(offset: -2), .close, .open(.inconsistent, 8))
-    before(node.whileKeyword, tokens: .break)
-    after(node.whileKeyword, tokens: .break)
+    before(node.whileKeyword, tokens: .space)
+    after(node.whileKeyword, tokens: .space)
     after(node.condition.lastToken, tokens: .close)
     super.visit(node)
   }
@@ -839,7 +848,8 @@ private final class TokenStreamCreator: SyntaxVisitor {
 
   override func visit(_ node: InitializerClauseSyntax) {
     before(node.equal, tokens: .break)
-    after(node.equal, tokens: .break)
+    after(node.equal, tokens: .break(offset: 2))
+
     super.visit(node)
   }
 

Tests/PrettyPrinterTests/ArrayDeclTests.swift

@@ -0,0 +1,32 @@
+public class ArrayDeclTests: PrettyPrintTestCase {
+  public func testBasicArrays() {
+    let input =
+      """
+      let a = [1, 2, 3]
+      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 expected =
+      """
+      let a = [1, 2, 3]
+      let a: [Bool] = [false, true, true, false]
+      let a: [String] = [
+        "One", "Two", "Three", "Four"
+      ]
+      let a: [String] = [
+        "One",
+        "Two",
+        "Three",
+        "Four",
+        "Five",
+        "Six",
+        "Seven"
+      ]
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 45)
+  }
+}

Tests/PrettyPrinterTests/ForInStmtTests.swift

@@ -20,8 +20,8 @@ public class ForInStmtTests: PrettyPrintTestCase {
         let b = i
       }
 
-      for item in
-          mylargecontainer {
+      for item
+          in mylargecontainer {
         let a = 123
         let b = item
       }
@@ -75,18 +75,18 @@ public class ForInStmtTests: PrettyPrintTestCase {
   public func testForLoopFullWrap() {
     let input =
       """
-      for item in aVeryLargeContainterObject where aVeryLargeContainerObject.hasProperty() {
+      for item in aVeryLargeContainterObject where largeObject.hasProperty() && condition {
         let a = 123
         let b = 456
       }
       """
 
     let expected =
       """
-      for item in
-          aVeryLargeContainterObject
-      where
-        aVeryLargeContainerObject.hasProperty()
+      for item
+          in aVeryLargeContainterObject
+      where largeObject.hasProperty()
+        && condition
       {
         let a = 123
         let b = 456

Tests/PrettyPrinterTests/VariableDeclTests.swift

@@ -10,12 +10,13 @@ public class VariableDeclarationTests: PrettyPrintTestCase {
 
     let expected =
       """
-      let x = firstVariable +
+      let x =
+        firstVariable +
         secondVariable /
         thirdVariable +
         fourthVariable
-      let y: Int = anotherVar +
-        moreVar
+      let y: Int =
+        anotherVar + moreVar
       let (w, z, s):
         (Int, Double, Bool) =
         firstTuple + secondTuple