Add parameter inference

Sources/Inference.swift

@@ -315,7 +315,11 @@ private struct Inference {
     }
 
     let wrapArguments = OptionInferrer { formatter, options in
-        options.wrapArguments = formatter.wrapMode(for: "(", "<")
+        options.wrapArguments = formatter.wrapMode(forParameters: false)
+    }
+
+    let wrapParameters = OptionInferrer { formatter, options in
+        options.wrapParameters = formatter.wrapMode(forParameters: true)
     }
 
     let wrapCollections = OptionInferrer { formatter, options in
@@ -1226,6 +1230,30 @@ private struct Inference {
 }
 
 private extension Formatter {
+    func wrapMode(forParameters parameters: Bool) -> WrapMode {
+        var beforeFirst = 0, afterFirst = 0
+        forEachToken(where: { [.startOfScope("("), .startOfScope("<")].contains($0) }) { i, _ in
+            guard isParameterList(at: i) == parameters,
+                let closingBraceIndex = endOfScope(at: i),
+                index(of: .linebreak, in: i + 1 ..< closingBraceIndex) != nil else {
+                return
+            }
+            // Check if linebreak is after opening paren or first comma
+            if next(.nonSpaceOrComment, after: i)?.isLinebreak == true {
+                beforeFirst += 1
+            } else {
+                afterFirst += 1
+            }
+        }
+        if beforeFirst > 0, afterFirst == 0 {
+            return .beforeFirst
+        } else if afterFirst > 0, beforeFirst == 0 {
+            return .afterFirst
+        } else {
+            return .preserve
+        }
+    }
+
     func wrapMode(for scopes: String...) -> WrapMode {
         var beforeFirst = 0, afterFirst = 0
         forEachToken(where: { $0.isStartOfScope && scopes.contains($0.string) }) { i, _ in

Sources/ParsingHelpers.swift

@@ -713,7 +713,7 @@ extension Formatter {
     }
 
     func isParameterList(at i: Int) -> Bool {
-        assert(tokens[i] == .startOfScope("("))
+        assert([.startOfScope("("), .startOfScope("<")].contains(tokens[i]))
         if let endIndex = endOfScope(at: i),
             let nextToken = next(.nonSpaceOrCommentOrLinebreak, after: endIndex),
             [.operator("->", .infix), .keyword("throws"), .keyword("rethrows")].contains(nextToken) {

Tests/InferenceTests.swift

@@ -236,6 +236,20 @@ class InferenceTests: XCTestCase {
     // MARK: wrapArguments
 
     func testInferWrapBeforeFirstArgument() {
+        let input = """
+        foo(
+            bar: Int,
+            baz: String)
+        foo(
+            bar: Int,
+            baz: String
+        )
+        """
+        let options = inferFormatOptions(from: tokenize(input))
+        XCTAssertEqual(options.wrapArguments, .beforeFirst)
+    }
+
+    func testInferWrapBeforeFirstParameter() {
         let input = """
         func foo(
             bar: Int,
@@ -245,17 +259,26 @@ class InferenceTests: XCTestCase {
             baz: String)
         """
         let options = inferFormatOptions(from: tokenize(input))
-        XCTAssertEqual(options.wrapArguments, .beforeFirst)
+        XCTAssertEqual(options.wrapParameters, .beforeFirst)
     }
 
     func testInferWrapAfterFirstArgument() {
+        let input = """
+        foo(bar: Int,
+            baz: String, quux: String)
+        """
+        let options = inferFormatOptions(from: tokenize(input))
+        XCTAssertEqual(options.wrapArguments, .afterFirst)
+    }
+
+    func testInferWrapAfterFirstParameter() {
         let input = """
         func foo(bar: Int,
                  baz: String,
                  quux: String) {}
         """
         let options = inferFormatOptions(from: tokenize(input))
-        XCTAssertEqual(options.wrapArguments, .afterFirst)
+        XCTAssertEqual(options.wrapParameters, .afterFirst)
     }
 
     func testInferWrapPreserve() {

Tests/XCTestManifests.swift

@@ -252,7 +252,9 @@ extension InferenceTests {
         ("testInferUppercaseHexExponent", testInferUppercaseHexExponent),
         ("testInferUseVoid", testInferUseVoid),
         ("testInferWrapAfterFirstArgument", testInferWrapAfterFirstArgument),
+        ("testInferWrapAfterFirstParameter", testInferWrapAfterFirstParameter),
         ("testInferWrapBeforeFirstArgument", testInferWrapBeforeFirstArgument),
+        ("testInferWrapBeforeFirstParameter", testInferWrapBeforeFirstParameter),
         ("testInferWrapElementsAfterFirstArgument", testInferWrapElementsAfterFirstArgument),
         ("testInferWrapElementsAfterSecondArgument", testInferWrapElementsAfterSecondArgument),
         ("testInferWrapPreserve", testInferWrapPreserve),