Merge pull request swiftlang#15 from allevato/sr-11109

Preserve order of case elements in OneCasePerLine.

Author: allevato

Sources/SwiftFormatPrettyPrint/TokenStreamCreator.swift

@@ -834,6 +834,9 @@ private final class TokenStreamCreator: SyntaxVisitor {
 
   func visit(_ node: EnumCaseDeclSyntax) -> SyntaxVisitorContinueKind {
     before(node.firstToken, tokens: .open)
+
+    arrangeAttributeList(node.attributes)
+
     after(node.caseKeyword, tokens: .break)
     after(node.lastToken, tokens: .close)
     return .visitChildren

Sources/SwiftFormatRules/OneCasePerLine.swift

@@ -14,98 +14,125 @@ import Foundation
 import SwiftFormatCore
 import SwiftSyntax
 
-/// Each enum case with associated values should appear on its own line.
+/// Each enum case with associated values or a raw value should appear in its own case declaration.
 ///
 /// Lint: If a single `case` declaration declares multiple cases, and any of them have associated
-///       values, a lint error is raised.
+///       values or raw values, a lint error is raised.
 ///
-/// Format: All case declarations with associated values will be moved to a new line.
+/// Format: All case declarations with associated values or raw values will be moved to their own
+///         case declarations.
 ///
 /// - SeeAlso: https://google.github.io/swift#enum-cases
 public final class OneCasePerLine: SyntaxFormatRule {
 
+  /// A state machine that collects case elements encountered during visitation and allows new case
+  /// declarations to be created with those elements.
+  private struct CaseElementCollector {
+
+    /// The case declaration used as the source from which additional new declarations will be
+    /// created; thus, all new cases will share the same attributes and modifiers as the basis.
+    public private(set) var basis: EnumCaseDeclSyntax
+
+    /// Case elements collected so far.
+    private var elements = [EnumCaseElementSyntax]()
+
+    /// Indicates whether the full leading trivia of basis case declaration should be preserved by
+    /// the next case declaration that will be created by copying the basis declaration.
+    ///
+    /// This is true for the first case (to preserve any leading comments on the original case
+    /// declaration) and false for all subsequent cases (so that we don't repeat those comments).
+    private var shouldKeepLeadingTrivia = true
+
+    /// Creates a new case element collector based on the given case declaration.
+    init(basedOn basis: EnumCaseDeclSyntax) {
+      self.basis = basis
+    }
+
+    /// Adds a new case element to the collector.
+    mutating func addElement(_ element: EnumCaseElementSyntax) {
+      elements.append(element)
+    }
+
+    /// Creates a new case declaration with the elements collected so far, then resets the internal
+    /// state to start a new empty declaration again.
+    ///
+    /// This will return nil if there are no elements collected since the last time this was called
+    /// (or the collector was created).
+    mutating func makeCaseDeclAndReset() -> EnumCaseDeclSyntax? {
+      guard let last = elements.last else { return nil }
+
+      // Remove the trailing comma on the final element, if there was one.
+      if last.trailingComma != nil {
+        elements[elements.count - 1] = last.withTrailingComma(nil)
+      }
+
+      defer { elements.removeAll() }
+      return makeCaseDeclFromBasis(elements: elements)
+    }
+
+    /// Creates and returns a new `EnumCaseDeclSyntax` with the given elements, based on the current
+    /// basis declaration, and updates the comment preserving state if needed.
+    mutating func makeCaseDeclFromBasis(elements: [EnumCaseElementSyntax]) -> EnumCaseDeclSyntax {
+      let caseDecl = basis.withElements(SyntaxFactory.makeEnumCaseElementList(elements))
+
+      if shouldKeepLeadingTrivia {
+        shouldKeepLeadingTrivia = false
+
+        // We don't bother preserving any indentation because the pretty printer will fix that up.
+        // All we need to do here is ensure that there is a newline.
+        basis = basis.withLeadingTrivia(Trivia.newlines(1))
+      }
+
+      return caseDecl
+    }
+  }
+
   public override func visit(_ node: EnumDeclSyntax) -> DeclSyntax {
-    let enumMembers = node.members.members
     var newMembers: [MemberDeclListItemSyntax] = []
-    var newIndx = 0
-
-    for member in enumMembers {
-      var numNewMembers = 0
-      if let caseMember = member.decl as? EnumCaseDeclSyntax {
-        var otherDecl: EnumCaseDeclSyntax? = caseMember
-        // Add and skip single element case declarations
-        guard caseMember.elements.count > 1 else {
-          newMembers.append(member.withDecl(caseMember))
-          newIndx += 1
-          continue
-        }
-        // Move all cases with associated/raw values to new declarations
-        for element in caseMember.elements {
-          if element.associatedValue != nil || element.rawValue != nil {
-            diagnose(.moveAssociatedOrRawValueCase(name: element.identifier.text), on: element)
-            let newRemovedDecl = createAssociateOrRawCaseDecl(
-              fullDecl: caseMember,
-              removedElement: element)
-            otherDecl = removeAssociateOrRawCaseDecl(fullDecl: otherDecl)
-            newMembers.append(member.withDecl(newRemovedDecl))
-            numNewMembers += 1
-          }
-        }
-        // Add case declaration of remaining elements without associated/raw values, if any
-        if let otherDecl = otherDecl {
-          newMembers.insert(member.withDecl(otherDecl), at: newIndx)
-          newIndx += 1
-        }
-        // Add any member that isn't an enum case declaration
-      } else {
+
+    for member in node.members.members {
+      // If it's not a case declaration, or it's a case declaration with only one element, leave it
+      // alone.
+      guard let caseDecl = member.decl as? EnumCaseDeclSyntax, caseDecl.elements.count > 1 else {
         newMembers.append(member)
-        newIndx += 1
+        continue
       }
-      newIndx += numNewMembers
-    }
 
-    let newMemberBlock = SyntaxFactory.makeMemberDeclBlock(
-      leftBrace: node.members.leftBrace,
-      members: SyntaxFactory.makeMemberDeclList(newMembers),
-      rightBrace: node.members.rightBrace)
-    return node.withMembers(newMemberBlock)
-  }
+      var collector = CaseElementCollector(basedOn: caseDecl)
 
-  func createAssociateOrRawCaseDecl(
-    fullDecl: EnumCaseDeclSyntax,
-    removedElement: EnumCaseElementSyntax
-  ) -> EnumCaseDeclSyntax {
-    let formattedElement = removedElement.withTrailingComma(nil)
-    let newElementList = SyntaxFactory.makeEnumCaseElementList([formattedElement])
-    let newDecl = SyntaxFactory.makeEnumCaseDecl(
-      attributes: fullDecl.attributes,
-      modifiers: fullDecl.modifiers,
-      caseKeyword: fullDecl.caseKeyword,
-      elements: newElementList)
-    return newDecl
-  }
+      // Collect the elements of the case declaration until we see one that has either an associated
+      // value or a raw value.
+      for element in caseDecl.elements {
+        if element.associatedValue != nil || element.rawValue != nil {
+          // Once we reach one of these, we need to write out the ones we've collected so far, then
+          // emit a separate case declaration with the associated/raw value element.
+          diagnose(.moveAssociatedOrRawValueCase(name: element.identifier.text), on: element)
 
-  // Returns formatted declaration of cases without associated/raw values, or nil if all cases had
-  // a raw or associate value
-  func removeAssociateOrRawCaseDecl(fullDecl: EnumCaseDeclSyntax?) -> EnumCaseDeclSyntax? {
-    guard let fullDecl = fullDecl else { return nil }
-    var newList: [EnumCaseElementSyntax] = []
+          if let caseDeclForCollectedElements = collector.makeCaseDeclAndReset() {
+            newMembers.append(member.withDecl(caseDeclForCollectedElements))
+          }
 
-    for element in fullDecl.elements {
-      if element.associatedValue == nil && element.rawValue == nil { newList.append(element) }
-    }
+          let separatedCaseDecl =
+            collector.makeCaseDeclFromBasis(elements: [element.withTrailingComma(nil)])
+          newMembers.append(member.withDecl(separatedCaseDecl))
+        } else {
+          collector.addElement(element)
+        }
+      }
 
-    guard newList.count > 0 else { return nil }
-    let (last, indx) = (newList[newList.count - 1], newList.count - 1)
-    if last.trailingComma != nil {
-      newList[indx] = last.withTrailingComma(nil)
+      // Make sure to emit any trailing collected elements.
+      if let caseDeclForCollectedElements = collector.makeCaseDeclAndReset() {
+        newMembers.append(member.withDecl(caseDeclForCollectedElements))
+      }
     }
-    return fullDecl.withElements(SyntaxFactory.makeEnumCaseElementList(newList))
+
+    let newMemberBlock = node.members.withMembers(SyntaxFactory.makeMemberDeclList(newMembers))
+    return node.withMembers(newMemberBlock)
   }
 }
 
 extension Diagnostic.Message {
   static func moveAssociatedOrRawValueCase(name: String) -> Diagnostic.Message {
-    return .init(.warning, "move \(name) case to a new line")
+    return .init(.warning, "move '\(name)' to its own case declaration")
   }
 }

Tests/SwiftFormatRulesTests/OneCasePerLineTests.swift

@@ -5,35 +5,112 @@ import XCTest
 @testable import SwiftFormatRules
 
 public class OneCasePerLineTests: DiagnosingTestCase {
+
+  // The inconsistent leading whitespace in the expected text is intentional. This transform does
+  // not attempt to preserve leading indentation since the pretty printer will correct it when
+  // running the full formatter.
+
   func testInvalidCasesOnLine() {
-    XCTAssertFormatting(OneCasePerLine.self,
-                        input: """
-                               public enum Token {
-                                 case arrow
-                                 case comma, identifier(String), semicolon, stringSegment(String)
-                                 case period
-                                 case ifKeyword(String), forKeyword(String)
-                                 indirect case guardKeyword, elseKeyword, contextualKeyword(String)
-                                 var x: Bool
-                                 case leftParen, rightParen = ")", leftBrace, rightBrace = "}"
-                               }
-                               """,
-                        expected: """
-                                  public enum Token {
-                                    case arrow
-                                    case comma, semicolon
-                                    case identifier(String)
-                                    case stringSegment(String)
-                                    case period
-                                    case ifKeyword(String)
-                                    case forKeyword(String)
-                                    indirect case guardKeyword, elseKeyword
-                                    indirect case contextualKeyword(String)
-                                    var x: Bool
-                                    case leftParen, leftBrace
-                                    case rightParen = ")"
-                                    case rightBrace = "}"
-                                  }
-                                  """)
+    XCTAssertFormatting(
+      OneCasePerLine.self,
+      input:
+        """
+        public enum Token {
+          case arrow
+          case comma, identifier(String), semicolon, stringSegment(String)
+          case period
+          case ifKeyword(String), forKeyword(String)
+          indirect case guardKeyword, elseKeyword, contextualKeyword(String)
+          var x: Bool
+          case leftParen, rightParen = ")", leftBrace, rightBrace = "}"
+        }
+        """,
+      expected:
+        """
+        public enum Token {
+          case arrow
+          case comma
+        case identifier(String)
+        case semicolon
+        case stringSegment(String)
+          case period
+          case ifKeyword(String)
+        case forKeyword(String)
+          indirect case guardKeyword, elseKeyword
+        indirect case contextualKeyword(String)
+          var x: Bool
+          case leftParen
+        case rightParen = ")"
+        case leftBrace
+        case rightBrace = "}"
+        }
+        """)
+  }
+
+  func testElementOrderIsPreserved() {
+    XCTAssertFormatting(
+      OneCasePerLine.self,
+      input:
+        """
+        enum Foo: Int {
+          case a = 0, b, c, d
+        }
+        """,
+      expected:
+        """
+        enum Foo: Int {
+          case a = 0
+        case b, c, d
+        }
+        """)
+  }
+
+  func testCommentsAreNotRepeated() {
+    XCTAssertFormatting(
+      OneCasePerLine.self,
+      input:
+        """
+        enum Foo: Int {
+          /// This should only be above `a`.
+          case a = 0, b, c, d
+          // This should only be above `e`.
+          case e, f = 100
+        }
+        """,
+      expected:
+        """
+        enum Foo: Int {
+          /// This should only be above `a`.
+          case a = 0
+        case b, c, d
+          // This should only be above `e`.
+          case e
+        case f = 100
+        }
+        """)
+  }
+
+  func testAttributesArePropagated() {
+    XCTAssertFormatting(
+      OneCasePerLine.self,
+      input:
+        """
+        enum Foo {
+          @someAttr case a(String), b, c, d
+          case e, f(Int)
+          @anotherAttr case g, h(Float)
+        }
+        """,
+      expected:
+        """
+        enum Foo {
+          @someAttr case a(String)
+        @someAttr case b, c, d
+          case e
+        case f(Int)
+          @anotherAttr case g
+        @anotherAttr case h(Float)
+        }
+        """)
   }
 }

Tests/SwiftFormatRulesTests/XCTestManifests.swift

@@ -240,6 +240,9 @@ extension OneCasePerLineTests {
     //   `swift test --generate-linuxmain`
     // to regenerate.
     static let __allTests__OneCasePerLineTests = [
+        ("testAttributesArePropagated", testAttributesArePropagated),
+        ("testCommentsAreNotRepeated", testCommentsAreNotRepeated),
+        ("testElementOrderIsPreserved", testElementOrderIsPreserved),
         ("testInvalidCasesOnLine", testInvalidCasesOnLine),
     ]
 }