[AutoDiff] add some @derivative(of:) tests

include/swift/AST/AutoDiff.h

@@ -491,12 +491,18 @@ template<typename T> struct DenseMapInfo;
 template<> struct DenseMapInfo<AutoDiffConfig> {
   static AutoDiffConfig getEmptyKey() {
     auto *ptr = llvm::DenseMapInfo<void *>::getEmptyKey();
+    // The `derivativeGenericSignature` component must be `nullptr` so that
+    // `getHashValue` and `isEqual` do not try to `getCanonicalSignature()` on
+    // an invalid pointer.
     return {static_cast<IndexSubset *>(ptr), static_cast<IndexSubset *>(ptr),
             nullptr};
   }
 
   static AutoDiffConfig getTombstoneKey() {
     auto *ptr = llvm::DenseMapInfo<void *>::getTombstoneKey();
+    // The `derivativeGenericSignature` component must be `nullptr` so that
+    // `getHashValue` and `isEqual` do not try to `getCanonicalSignature()` on
+    // an invalid pointer.
     return {static_cast<IndexSubset *>(ptr), static_cast<IndexSubset *>(ptr),
             nullptr};
   }

test/AutoDiff/Serialization/derivative_attr.swift

@@ -0,0 +1,27 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend %s -emit-module -parse-as-library -o %t
+// RUN: llvm-bcanalyzer %t/derivative_attr.swiftmodule | %FileCheck %s -check-prefix=BCANALYZER
+// RUN: %target-sil-opt -disable-sil-linking -enable-sil-verify-all %t/derivative_attr.swiftmodule -o - | %FileCheck %s
+
+// BCANALYZER-NOT: UnknownCode
+
+struct DerivativeOfSpecialMethods: Differentiable {
+  var x: Float
+
+  init(_ x: Float) { self.x = x }
+  subscript(_ i: Int) -> Float { x }
+
+  // CHECK: @derivative(of: init, wrt: x)
+  // CHECK-NEXT: static func dInit(_ x: Float)
+  @derivative(of: init(_:))
+  static func dInit(_ x: Float) -> (value: Self, pullback: (TangentVector) -> Float) {
+    fatalError()
+  }
+
+  // CHECK: @derivative(of: subscript, wrt: self)
+  // CHECK-NEXT: func dSubscript(_ i: Int)
+  @derivative(of: subscript(_:))
+  func dSubscript(_ i: Int) -> (value: Float, pullback: (Float) -> TangentVector) {
+    fatalError()
+  }
+}

test/AutoDiff/derivative_attr_type_checking.swift

@@ -532,3 +532,33 @@ struct StoredProperty: Differentiable {
     (stored, { _ in .zero })
   }
 }
+
+// When the generic signature was not considered while calculating the actual pullback type, the
+// typechecker did not realize that `T.TangentVector == Float`, and therefore it complained that
+// "'pullback' does not have expected type '(Float) -> (Float)'".
+// Users were able to work around this by setting the pullback type to `(Float) -> Float`.
+func genericSignatureConsidered<T>(_ x: T) -> T { fatalError() }
+@derivative(of: genericSignatureConsidered)
+func dGenericSignatureConsidered<T>(_ x: T)
+  -> (value: T, pullback: (T.TangentVector) -> T.TangentVector)
+  where T: Differentiable, T.TangentVector == Float
+{
+  fatalError()
+}
+
+// When the generic signature was not considered while calculating the actual pullback type,
+// the typechecker complained that the pullback type was not correct, and there was no pullback type
+// that users could specify to satisfy the typechecker.
+struct Wrapper<T: AdditiveArithmetic & Equatable>: AdditiveArithmetic, Equatable {
+  var t: T
+  init(_ t: T) { self.t = t }
+}
+extension Wrapper: Differentiable where T: Differentiable, T == T.TangentVector {
+  typealias TangentVector = Wrapper<T.TangentVector>
+}
+extension Wrapper where T: Differentiable, T == T.TangentVector {
+  @derivative(of: init(_:))
+  static func dInit(_ t: T) -> (value: Self, pullback: (Wrapper<T>.TangentVector) -> (T)) {
+    fatalError()
+  }
+}

test/AutoDiff/derivative_registration.swift

@@ -156,4 +156,24 @@ DerivativeRegistrationTests.testWithLeakChecking("DerivativeGenericSignature") {
   expectEqual(1000, dx)
 }
 
+// When non-canonicalized generic signatures are used to compare derivative configurations, the
+// `@differentiable` and `@derivative` attributes create separate derivatives, and we get a
+// duplicate symbol error in TBDGen.
+public protocol RefinesDifferentiable: Differentiable {}
+extension Float: RefinesDifferentiable {}
+@differentiable(where T: Differentiable, T: RefinesDifferentiable)
+public func nonCanonicalizedGenSigComparison<T>(_ t: T) -> T { t }
+@derivative(of: nonCanonicalizedGenSigComparison)
+public func dNonCanonicalizedGenSigComparison<T: RefinesDifferentiable>(_ t: T)
+  -> (value: T, pullback: (T.TangentVector) -> T.TangentVector)
+{
+  (t, { _ in T.TangentVector.zero })
+}
+DerivativeRegistrationTests.testWithLeakChecking("NonCanonicalizedGenericSignatureComparison") {
+  let dx = gradient(at: Float(0), in: nonCanonicalizedGenSigComparison)
+  // Expect that we use the custom registered derivative, not a generated derivative (which would
+  // give a gradient of 1).
+  expectEqual(0, dx)
+}
+
 runAllTests()