Merge pull request #791 from CodaFi/derivative!

Parse @Derivative(of:wrt:) and @transpose(of:wrt:)

Author: CodaFi

Sources/SwiftParser/Attributes.swift

@@ -39,6 +39,10 @@ extension Parser {
       return RawSyntax(self.parseAvailabilityAttribute())
     case .differentiable:
       return RawSyntax(self.parseDifferentiableAttribute())
+    case .derivative:
+      return RawSyntax(self.parseDerivativeAttribute())
+    case .transpose:
+      return RawSyntax(self.parseTransposeAttribute())
     case .objc:
       return RawSyntax(self.parseObjectiveCAttribute())
     case ._specialize:
@@ -296,6 +300,84 @@ extension Parser {
   }
 }
 
+extension Parser {
+  mutating func parseDerivativeAttribute() -> RawAttributeSyntax {
+    let (unexpectedBeforeAtSign, atSign) = self.expect(.atSign)
+    let (unexpectedBeforeDerivative, derivative) = self.expectContextualKeyword("derivative")
+
+    let (unexpectedBeforeLeftParen, leftParen) = self.expect(.leftParen)
+    let argument = self.parseDerivativeAttributeArguments()
+    let (unexpectedBeforeRightParen, rightParen) = self.expect(.rightParen)
+
+    return RawAttributeSyntax(
+      unexpectedBeforeAtSign,
+      atSignToken: atSign,
+      unexpectedBeforeDerivative,
+      attributeName: derivative,
+      unexpectedBeforeLeftParen,
+      leftParen: leftParen,
+      argument: RawSyntax(argument),
+      unexpectedBeforeRightParen,
+      rightParen: rightParen,
+      tokenList: nil,
+      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)
+    let originalDeclName = self.parseQualifiedDeclarationName()
+    let period = self.consume(if: .period)
+    let accessor: RawTokenSyntax?
+    if period != nil {
+      accessor = self.parseAnyIdentifier()
+    } else {
+      accessor = nil
+    }
+    let comma = self.consume(if: .comma)
+    let diffParams: RawDifferentiabilityParamsClauseSyntax?
+    if comma != nil {
+      diffParams = self.parseDifferentiabilityParameters()
+    } else {
+      diffParams = nil
+    }
+    return RawDerivativeRegistrationAttributeArgumentsSyntax(
+      unexpectedBeforeOfLabel,
+      ofLabel: ofLabel,
+      unexpectedBetweenOfLabelAndColon,
+      colon: colon,
+      originalDeclName: originalDeclName,
+      period: period,
+      accessorKind: accessor,
+      comma: comma,
+      diffParams: diffParams,
+      arena: self.arena)
+  }
+}
+
 extension Parser {
   mutating func parseObjectiveCAttribute() -> RawAttributeSyntax {
     let (unexpectedBeforeAtSign, atSign) = self.expect(.atSign)

Sources/SwiftParser/Names.swift

@@ -127,6 +127,42 @@ extension Parser {
   }
 }
 
+extension Parser {
+  mutating func parseQualifiedDeclarationName() -> RawQualifiedDeclNameSyntax {
+    let type: RawTypeSyntax?
+    let dot: RawTokenSyntax?
+    if self.lookahead().canParseBaseTypeForQualifiedDeclName() {
+      type = self.parseTypeIdentifier()
+      dot = self.consumePrefix(".", as: .period)
+    } else {
+      type = nil
+      dot = nil
+    }
+
+    let (name, args) = self.parseDeclNameRef([
+      .zeroArgCompoundNames,
+      .keywordsUsingSpecialNames,
+      .operators,
+    ])
+    return RawQualifiedDeclNameSyntax(
+      baseType: type,
+      dot: dot,
+      name: name,
+      arguments: args,
+      arena: self.arena)
+  }
+}
+
+extension Parser.Lookahead {
+  func canParseBaseTypeForQualifiedDeclName() -> Bool {
+    var lookahead = self.lookahead()
+    guard lookahead.canParseSimpleTypeIdentifier() else {
+      return false
+    }
+    return lookahead.currentToken.starts(with: ".")
+  }
+}
+
 extension Parser.Lookahead {
   func canParseArgumentLabelList() -> Bool {
     var lookahead = self.lookahead()

Tests/SwiftParserTest/Attributes.swift

@@ -103,4 +103,69 @@ final class AttributeTests: XCTestCase {
       @_Concurrency.MainActor(unsafe) public struct Image : SwiftUI.View {}
       """)
   }
+
+  func testDerivativeAttribute() {
+    AssertParse(
+      """
+      @inlinable
+      @differentiable(reverse, wrt: self)
+      public func differentiableMap<Result: Differentiable>(
+        _ body: @differentiable(reverse) (Element) -> Result
+      ) -> [Result] {
+        map(body)
+      }
+      """)
+
+    AssertParse(
+      """
+      @inlinable
+      @differentiable(reverse, wrt: (self, initialResult))
+      public func differentiableReduce<Result: Differentiable>(
+        _ initialResult: Result,
+        _ nextPartialResult: @differentiable(reverse) (Result, Element) -> Result
+      ) -> Result {
+        reduce(initialResult, nextPartialResult)
+      }
+      """)
+
+    AssertParse(
+      """
+      @inlinable
+      @derivative(of: differentiableReduce)
+      internal func _vjpDifferentiableReduce<Result: Differentiable>(
+        _ initialResult: Result,
+        _ nextPartialResult: @differentiable(reverse) (Result, Element) -> Result
+      ) -> (
+        value: Result,
+        pullback: (Result.TangentVector)
+          -> (Array.TangentVector, Result.TangentVector)
+      ) {}
+      """)
+  }
+
+  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)
+      }
+      """)
+  }
 }