Parse @transpose(of:wrt:)

Author: CodaFi

Sources/SwiftParser/Attributes.swift

@@ -41,6 +41,8 @@ extension Parser {
       return RawSyntax(self.parseDifferentiableAttribute())
     case .derivative:
       return RawSyntax(self.parseDerivativeAttribute())
+    case .transpose:
+      return RawSyntax(self.parseTransposeAttribute())
     case .objc:
       return RawSyntax(self.parseObjectiveCAttribute())
     case ._specialize:
@@ -321,6 +323,29 @@ extension Parser {
       arena: self.arena)
   }
 
+  mutating func parseTransposeAttribute() -> RawAttributeSyntax {
+    let (unexpectedBeforeAtSign, atSign) = self.expect(.atSign)
+    let (unexpectedBeforeTranspose, transpose) = self.expectContextualKeyword("transpose")
+
+    let (unexpectedBeforeLeftParen, leftParen) = self.expect(.leftParen)
+    let argument = self.parseDerivativeAttributeArguments()
+    let (unexpectedBeforeRightParen, rightParen) = self.expect(.rightParen)
+
+    return RawAttributeSyntax(
+      unexpectedBeforeAtSign,
+      atSignToken: atSign,
+      unexpectedBeforeTranspose,
+      attributeName: transpose,
+      unexpectedBeforeLeftParen,
+      leftParen: leftParen,
+      argument: RawSyntax(argument),
+      unexpectedBeforeRightParen,
+      rightParen: rightParen,
+      tokenList: nil,
+      arena: self.arena)
+  }
+
+
   mutating func parseDerivativeAttributeArguments() -> RawDerivativeRegistrationAttributeArgumentsSyntax {
     let (unexpectedBeforeOfLabel, ofLabel) = self.expectContextualKeyword("of")
     let (unexpectedBetweenOfLabelAndColon, colon) = self.expect(.colon)

Tests/SwiftParserTest/Attributes.swift

@@ -142,4 +142,30 @@ final class AttributeTests: XCTestCase {
       ) {}
       """)
   }
+
+  func testTransposeAttribute() {
+    AssertParse(
+      """
+      @transpose(of: +)
+      func addTranspose(_ v: Float) -> (Float, Float) {
+        return (v, v)
+      }
+      """)
+
+    AssertParse(
+      """
+      @transpose(of: -, wrt: (0, 1))
+      func subtractTranspose(_ v: Float) -> (Float, Float) {
+        return (v, -v)
+      }
+      """)
+
+    AssertParse(
+      """
+      @transpose(of: Float.-, wrt: (0, 1))
+      func subtractTranspose(_ v: Float) -> (Float, Float) {
+        return (v, -v)
+      }
+      """)
+  }
 }