Merge pull request #30826 from dan-zheng/SR-12526

[AutoDiff] Add SR-12526 negative test.

Author: dan-zheng

lib/Sema/TypeCheckAttr.cpp

@@ -4423,14 +4423,6 @@ static bool typeCheckDerivativeAttr(ASTContext &Ctx, Decl *D,
     return true;
   }
 
-  // Reject different-file derivative registration.
-  // TODO(TF-1021): Lift same-file derivative registration restriction.
-  if (originalAFD->getParentSourceFile() != derivative->getParentSourceFile()) {
-    diags.diagnose(attr->getLocation(),
-                   diag::derivative_attr_not_in_same_file_as_original);
-    return true;
-  }
-
   // Reject duplicate `@derivative` attributes.
   auto &derivativeAttrs = Ctx.DerivativeAttrs[std::make_tuple(
       originalAFD, resolvedDiffParamIndices, kind)];

test/AutoDiff/Sema/DerivativeRegistrationCrossModule/Inputs/a.swift

@@ -0,0 +1,7 @@
+public struct Struct {
+  public func method(_ x: Float) -> Float { x }
+
+  public static func +(_ lhs: Self, rhs: Self) -> Self {
+    lhs
+  }
+}

test/AutoDiff/Sema/DerivativeRegistrationCrossModule/Inputs/b.swift

@@ -0,0 +1,21 @@
+import _Differentiation
+import a
+
+extension Struct: Differentiable {
+  public struct TangentVector: Differentiable & AdditiveArithmetic {}
+  public mutating func move(along _: TangentVector) {}
+
+  @usableFromInline
+  @derivative(of: method, wrt: x)
+  func vjpMethod(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+    (x, { $0 })
+  }
+
+  @usableFromInline
+  @derivative(of: +)
+  static func vjpAdd(_ lhs: Self, rhs: Self) -> (
+    value: Self, pullback: (TangentVector) -> (TangentVector, TangentVector)
+  ) {
+    (lhs + rhs, { v in (v, v) })
+  }
+}

test/AutoDiff/Sema/DerivativeRegistrationCrossModule/main.swift

@@ -0,0 +1,14 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -emit-module -primary-file %S/Inputs/a.swift -emit-module-path %t/a.swiftmodule
+// RUN: %target-swift-frontend -emit-module -primary-file %S/Inputs/b.swift -emit-module-path %t/b.swiftmodule -I %t
+// RUN: not --crash %target-swift-frontend-typecheck -verify -I %t %s
+
+// SR-12526: Fix cross-module deserialization crash involving `@derivative` attribute.
+
+import a
+import b
+
+func foo(_ s: Struct) {
+  _ = Struct()
+  _ = s.method(1)
+}

test/AutoDiff/Sema/derivative_attr_type_checking.swift

@@ -743,11 +743,9 @@ extension InoutParameters {
   }
 }
 
-// Test cross-file derivative registration. Currently unsupported.
-// TODO(TF-1021): Lift this restriction.
+// Test cross-file derivative registration.
 
 extension FloatingPoint where Self: Differentiable {
-  // expected-error @+1 {{derivative not in the same file as the original function}}
   @derivative(of: rounded)
   func vjpRounded() -> (
     value: Self,