Format differentiation attributes.

This change adds correct formatting for `@differentiable`,
`@derivative`, and `@transpose`.

Fixes SR-12002.

Author: allevato

Sources/SwiftFormatPrettyPrint/TokenStreamCreator.swift

@@ -1933,6 +1933,56 @@ private final class TokenStreamCreator: SyntaxVisitor {
     return .visitChildren
   }
 
+  override func visit(_ node: DifferentiableAttributeArgumentsSyntax) -> SyntaxVisitorContinueKind {
+    // This node encapsulates the entire list of arguments in a `@differentiable(...)` attribute.
+    if let vjp = node.maybeVJP {
+      before(vjp.firstToken, tokens: .break(.same), .open)
+      after(vjp.lastToken, tokens: .close)
+    }
+    if let jvp = node.maybeJVP {
+      before(jvp.firstToken, tokens: .break(.same), .open)
+      after(jvp.lastToken, tokens: .close)
+    }
+    if let whereClause = node.whereClause {
+      before(whereClause.firstToken, tokens: .break(.same), .open)
+      after(whereClause.lastToken, tokens: .close)
+    }
+    return .visitChildren
+  }
+
+  override func visit(_ node: DifferentiableAttributeFuncSpecifierSyntax)
+    -> SyntaxVisitorContinueKind
+  {
+    // This node encapsulates the `vjp:` or `jvp:` label and decl name in a `@differentiable`
+    // attribute.
+    after(node.colon, tokens: .break(.continue, newlines: .elective(ignoresDiscretionary: true)))
+    return .visitChildren
+  }
+
+  override func visit(_ node: DerivativeRegistrationAttributeArgumentsSyntax)
+    -> SyntaxVisitorContinueKind
+  {
+    // This node encapsulates the entire list of arguments in a `@derivative(...)` or
+    // `@transpose(...)` attribute.
+    before(node.ofLabel, tokens: .open)
+    after(node.colon, tokens: .break(.continue, newlines: .elective(ignoresDiscretionary: true)))
+    after(node.comma, tokens: .close)
+
+    if let diffParams = node.diffParams {
+      before(diffParams.firstToken, tokens: .break(.same), .open)
+      after(diffParams.lastToken, tokens: .close)
+    }
+
+    return .visitChildren
+  }
+
+  override func visit(_ node: DifferentiationParamsClauseSyntax) -> SyntaxVisitorContinueKind {
+    // This node encapsulates the `wrt:` label and value/variable in a `@differentiable`,
+    // `@derivative`, or `@transpose` attribute.
+    after(node.colon, tokens: .break(.continue, newlines: .elective(ignoresDiscretionary: true)))
+    return .visitChildren
+  }
+
   // MARK: - Nodes representing unknown or malformed syntax
 
   override func visit(_ node: UnknownDeclSyntax) -> SyntaxVisitorContinueKind {

Tests/SwiftFormatPrettyPrintTests/DifferentiationAttributeTests.swift

@@ -0,0 +1,117 @@
+class DifferentiationAttributeTests: PrettyPrintTestCase {
+  func testDifferentiable() {
+    let input =
+      """
+      @differentiable(wrt: x, vjp: d where T: D)
+      func foo<T>(_ x: T) -> T {}
+
+      @differentiable(wrt: x, vjp: deriv where T: D)
+      func foo<T>(_ x: T) -> T {}
+
+      @differentiable(wrt: x, vjp: derivativeFoo where T: Differentiable)
+      func foo<T>(_ x: T) -> T {}
+
+      @differentiable(wrt: theVariableNamedX, vjp: derivativeFoo where T: Differentiable)
+      func foo<T>(_ theVariableNamedX: T) -> T {}
+      """
+
+    let expected =
+      """
+      @differentiable(wrt: x, vjp: d where T: D)
+      func foo<T>(_ x: T) -> T {}
+
+      @differentiable(
+        wrt: x, vjp: deriv where T: D
+      )
+      func foo<T>(_ x: T) -> T {}
+
+      @differentiable(
+        wrt: x, vjp: derivativeFoo
+        where T: Differentiable
+      )
+      func foo<T>(_ x: T) -> T {}
+
+      @differentiable(
+        wrt: theVariableNamedX,
+        vjp: derivativeFoo
+        where T: Differentiable
+      )
+      func foo<T>(_ theVariableNamedX: T) -> T {}
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 43)
+  }
+
+  func testDerivative() {
+    let input =
+      """
+      @derivative(of: foo, wrt: x)
+      func deriv<T>(_ x: T) {}
+
+      @derivative(of: foobar, wrt: x)
+      func deriv<T>(_ x: T) {}
+
+      @derivative(of: foobarbaz, wrt: theVariableNamedX)
+      func deriv<T>(_ theVariableNamedX: T) {}
+      """
+
+    let expected =
+      """
+      @derivative(of: foo, wrt: x)
+      func deriv<T>(_ x: T) {}
+
+      @derivative(
+        of: foobar, wrt: x
+      )
+      func deriv<T>(_ x: T) {}
+
+      @derivative(
+        of: foobarbaz,
+        wrt: theVariableNamedX
+      )
+      func deriv<T>(
+        _ theVariableNamedX: T
+      ) {}
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 28)
+  }
+
+  func testTranspose() {
+    let input =
+      """
+      @transpose(of: foo, wrt: 0)
+      func trans<T>(_ v: T) {}
+
+      @transpose(of: foobar, wrt: 0)
+      func trans<T>(_ v: T) {}
+
+      @transpose(of: someReallyLongName, wrt: 0)
+      func trans<T>(_ theVariableNamedV: T) {}
+      """
+
+    let expected =
+      """
+      @transpose(of: foo, wrt: 0)
+      func trans<T>(_ v: T) {}
+
+      @transpose(
+        of: foobar, wrt: 0
+      )
+      func trans<T>(_ v: T) {}
+
+      @transpose(
+        of: someReallyLongName,
+        wrt: 0
+      )
+      func trans<T>(
+        _ theVariableNamedV: T
+      ) {}
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 27)
+  }
+}

Tests/SwiftFormatPrettyPrintTests/XCTestManifests.swift

@@ -185,6 +185,17 @@ extension DictionaryDeclTests {
     ]
 }
 
+extension DifferentiationAttributeTests {
+    // DO NOT MODIFY: This is autogenerated, use:
+    //   `swift test --generate-linuxmain`
+    // to regenerate.
+    static let __allTests__DifferentiationAttributeTests = [
+        ("testDerivative", testDerivative),
+        ("testDifferentiable", testDifferentiable),
+        ("testTranspose", testTranspose),
+    ]
+}
+
 extension EnumDeclTests {
     // DO NOT MODIFY: This is autogenerated, use:
     //   `swift test --generate-linuxmain`
@@ -765,6 +776,7 @@ public func __allTests() -> [XCTestCaseEntry] {
         testCase(CommentTests.__allTests__CommentTests),
         testCase(DeinitializerDeclTests.__allTests__DeinitializerDeclTests),
         testCase(DictionaryDeclTests.__allTests__DictionaryDeclTests),
+        testCase(DifferentiationAttributeTests.__allTests__DifferentiationAttributeTests),
         testCase(EnumDeclTests.__allTests__EnumDeclTests),
         testCase(ExtensionDeclTests.__allTests__ExtensionDeclTests),
         testCase(ForInStmtTests.__allTests__ForInStmtTests),