actually set the lazy resolver

Author: Marc Rasi

include/swift/AST/Attr.h

@@ -1926,6 +1926,8 @@ class DerivativeAttr final
   }
   AbstractFunctionDecl *getOriginalFunction(ASTContext &context) const;
   void setOriginalFunction(AbstractFunctionDecl *decl);
+  void setOriginalFunctionResolver(LazyMemberLoader *resolver,
+                                   uint64_t resolverContextData);
 
   AutoDiffDerivativeFunctionKind getDerivativeKind() const {
     assert(Kind && "Derivative function kind has not yet been resolved");

lib/AST/Attr.cpp

@@ -1764,6 +1764,13 @@ void DerivativeAttr::setOriginalFunction(AbstractFunctionDecl *decl) {
   OriginalFunction = decl;
 }
 
+void DerivativeAttr::setOriginalFunctionResolver(
+    LazyMemberLoader *resolver, uint64_t resolverContextData) {
+  assert(!OriginalFunction && "cannot overwrite original function");
+  OriginalFunction = resolver;
+  ResolverContextData = resolverContextData;
+}
+
 TransposeAttr::TransposeAttr(bool implicit, SourceLoc atLoc,
                              SourceRange baseRange, TypeRepr *baseTypeRepr,
                              DeclNameRefWithLoc originalName,

lib/Serialization/Deserialization.cpp

@@ -4379,7 +4379,6 @@ llvm::Error DeclDeserializer::deserializeDeclAttributes() {
 
         DeclNameRefWithLoc origName{
             DeclNameRef(MF.getDeclBaseName(origNameId)), DeclNameLoc()};
-        auto *origDecl = cast<AbstractFunctionDecl>(MF.getDecl(origDeclId));
         auto derivativeKind =
             getActualAutoDiffDerivativeFunctionKind(rawDerivativeKind);
         if (!derivativeKind)
@@ -4392,7 +4391,7 @@ llvm::Error DeclDeserializer::deserializeDeclAttributes() {
         auto *derivativeAttr =
             DerivativeAttr::create(ctx, isImplicit, SourceLoc(), SourceRange(),
                                    /*baseType*/ nullptr, origName, indices);
-        derivativeAttr->setOriginalFunction(origDecl);
+        derivativeAttr->setOriginalFunctionResolver(&MF, origDeclId);
         derivativeAttr->setDerivativeKind(*derivativeKind);
         Attr = derivativeAttr;
         break;

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

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

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

@@ -10,12 +10,4 @@ extension Struct: Differentiable {
   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

@@ -1,14 +1,17 @@
 // 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: %target-swift-frontend-typecheck -verify -I %t %s
+// "-verify-ignore-unknown" is for "<unknown>:0: note: 'init()' declared here"
+// RUN: %target-swift-frontend-typecheck -verify -verify-ignore-unknown -I %t %s
 
 // SR-12526: Fix cross-module deserialization crash involving `@derivative` attribute.
 
 import a
 import b
 
 func foo(_ s: Struct) {
+  // Without this error, SR-12526 does not trigger.
+  // expected-error @+1 {{'Struct' initializer is inaccessible due to 'internal' protection level}}
   _ = Struct()
   _ = s.method(1)
 }