Add spaces after labels for overlooked stmt types.

Some statement types that allow a label were lacking the space token after the label. This caused the formatter to remove spaces when it encountered those stmt types.

Author: dylansturg

Sources/SwiftFormatPrettyPrint/TokenStreamCreator.swift

@@ -444,6 +444,7 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
   // MARK: - Control flow statement nodes
 
   override func visit(_ node: IfStmtSyntax) -> SyntaxVisitorContinueKind {
+    after(node.labelColon, tokens: .space)
     after(node.ifKeyword, tokens: .space)
 
     // Add break groups, using open continuation breaks, around any conditions after the first so
@@ -531,6 +532,7 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: RepeatWhileStmtSyntax) -> SyntaxVisitorContinueKind {
+    after(node.labelColon, tokens: .space)
     arrangeBracesAndContents(of: node.body, contentsKeyPath: \.statements)
 
     if config.lineBreakBeforeControlFlowKeywords {
@@ -550,6 +552,7 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: DoStmtSyntax) -> SyntaxVisitorContinueKind {
+    after(node.labelColon, tokens: .space)
     arrangeBracesAndContents(of: node.body, contentsKeyPath: \.statements)
     return .visitChildren
   }
@@ -591,6 +594,7 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: SwitchStmtSyntax) -> SyntaxVisitorContinueKind {
+    after(node.labelColon, tokens: .space)
     before(node.switchKeyword, tokens: .open)
     after(node.switchKeyword, tokens: .space)
     before(node.leftBrace, tokens: .break(.reset))

Tests/SwiftFormatPrettyPrintTests/DoStmtTests.swift

@@ -0,0 +1,61 @@
+import SwiftFormatConfiguration
+
+final class DoStmtTests: PrettyPrintTestCase {
+  func testBasicDoStmt() {
+    let input =
+      """
+      do {}
+      do { f() }
+      do { foo() }
+      do { let a = 123
+      var b = "abc"
+      }
+      do { veryLongFunctionCallThatShouldBeBrokenOntoANewLine() }
+      """
+
+    let expected =
+      """
+      do {}
+      do { f() }
+      do { foo() }
+      do {
+        let a = 123
+        var b = "abc"
+      }
+      do {
+        veryLongFunctionCallThatShouldBeBrokenOntoANewLine()
+      }
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 25)
+  }
+
+  func testLabeledDoStmt() {
+    let input = """
+      someLabel:do {
+        bar()
+        baz()
+      }
+      somePrettyLongLabelThatTakesUpManyColumns: do {
+        bar()
+        baz()
+      }
+      """
+
+    let expected = """
+      someLabel: do {
+        bar()
+        baz()
+      }
+      somePrettyLongLabelThatTakesUpManyColumns: do
+      {
+        bar()
+        baz()
+      }
+
+      """
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 45)
+  }
+}
+

Tests/SwiftFormatPrettyPrintTests/IfStmtTests.swift

@@ -488,4 +488,31 @@ final class IfStmtTests: PrettyPrintTestCase {
 
     assertPrettyPrintEqual(input: input, expected: expected, linelength: 50)
   }
+
+  func testLabeledIfStmt() {
+    let input =
+      """
+      someLabel:if foo && bar {
+        // do something
+      }
+      anotherVeryLongLabelThatTakesUpTooManyCharacters: if foo && bar {
+        // do something else
+      }
+      """
+
+    let expected =
+      """
+      someLabel: if foo && bar {
+        // do something
+      }
+      anotherVeryLongLabelThatTakesUpTooManyCharacters: if foo
+        && bar
+      {
+        // do something else
+      }
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 50)
+  }
 }

Tests/SwiftFormatPrettyPrintTests/RepeatStmtTests.swift

@@ -125,4 +125,31 @@ final class RepeatStmtTests: PrettyPrintTestCase {
       """
     assertPrettyPrintEqual(input: input, expected: input + "\n", linelength: 45)
   }
+
+  func testLabeledRepeat() {
+    let input = """
+      someLabel:repeat {
+        bar()
+        baz()
+      } while condition
+      somePrettyLongLabelThatTakesUpManyColumns: repeat {
+        bar()
+        baz()
+      } while condition
+      """
+
+    let expected = """
+      someLabel: repeat {
+        bar()
+        baz()
+      } while condition
+      somePrettyLongLabelThatTakesUpManyColumns: repeat
+      {
+        bar()
+        baz()
+      } while condition
+
+      """
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 45)
+  }
 }

Tests/SwiftFormatPrettyPrintTests/SwitchStmtTests.swift

@@ -288,4 +288,42 @@ final class SwitchStmtTests: PrettyPrintTestCase {
 
     assertPrettyPrintEqual(input: input, expected: expected, linelength: 40)
   }
+
+  func testLabeledSwitchStmt() {
+    let input =
+      """
+      label:switch foo {
+      case bar:
+        callForBar()
+      case baz:
+        callForBaz()
+      }
+      someVeryExtremelyLongLabel: switch foo {
+      case bar:
+        callForBar()
+      case baz:
+        callForBaz()
+      }
+      """
+
+    let expected =
+      """
+      label: switch foo {
+      case bar:
+        callForBar()
+      case baz:
+        callForBaz()
+      }
+      someVeryExtremelyLongLabel: switch foo
+      {
+      case bar:
+        callForBar()
+      case baz:
+        callForBaz()
+      }
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 20)
+  }
 }

Tests/SwiftFormatPrettyPrintTests/XCTestManifests.swift

@@ -215,6 +215,16 @@ extension DifferentiationAttributeTests {
     ]
 }
 
+extension DoStmtTests {
+    // DO NOT MODIFY: This is autogenerated, use:
+    //   `swift test --generate-linuxmain`
+    // to regenerate.
+    static let __allTests__DoStmtTests = [
+        ("testBasicDoStmt", testBasicDoStmt),
+        ("testLabeledDoStmt", testLabeledDoStmt),
+    ]
+}
+
 extension EnumDeclTests {
     // DO NOT MODIFY: This is autogenerated, use:
     //   `swift test --generate-linuxmain`
@@ -396,6 +406,7 @@ extension IfStmtTests {
         ("testIfElseStatement_noBreakBeforeElse", testIfElseStatement_noBreakBeforeElse),
         ("testIfLetStatements", testIfLetStatements),
         ("testIfStatement", testIfStatement),
+        ("testLabeledIfStmt", testLabeledIfStmt),
         ("testMatchingPatternConditions", testMatchingPatternConditions),
         ("testOptionalBindingConditions", testOptionalBindingConditions),
         ("testParenthesizedClauses", testParenthesizedClauses),
@@ -557,6 +568,7 @@ extension RepeatStmtTests {
     static let __allTests__RepeatStmtTests = [
         ("testBasicRepeatTests_breakBeforeWhile", testBasicRepeatTests_breakBeforeWhile),
         ("testBasicRepeatTests_noBreakBeforeWhile", testBasicRepeatTests_noBreakBeforeWhile),
+        ("testLabeledRepeat", testLabeledRepeat),
         ("testNestedRepeat", testNestedRepeat),
     ]
 }
@@ -693,6 +705,7 @@ extension SwitchStmtTests {
     // to regenerate.
     static let __allTests__SwitchStmtTests = [
         ("testBasicSwitch", testBasicSwitch),
+        ("testLabeledSwitchStmt", testLabeledSwitchStmt),
         ("testNestedSwitch", testNestedSwitch),
         ("testNewlinesDisambiguatingWhereClauses", testNewlinesDisambiguatingWhereClauses),
         ("testSwitchCases", testSwitchCases),
@@ -816,6 +829,7 @@ public func __allTests() -> [XCTestCaseEntry] {
         testCase(DeinitializerDeclTests.__allTests__DeinitializerDeclTests),
         testCase(DictionaryDeclTests.__allTests__DictionaryDeclTests),
         testCase(DifferentiationAttributeTests.__allTests__DifferentiationAttributeTests),
+        testCase(DoStmtTests.__allTests__DoStmtTests),
         testCase(EnumDeclTests.__allTests__EnumDeclTests),
         testCase(ExtensionDeclTests.__allTests__ExtensionDeclTests),
         testCase(ForInStmtTests.__allTests__ForInStmtTests),