[AutoDiff] Enhance performance of custom derivatives lookup

In swiftlang#58965, lookup for custom derivatives in non-primary source files was
introduced. It required traversing all delayed parsed function bodies of a
file if the file was compiled with differential programming enabled (even for
functions with no `@derivative` attribute).

This patch introduces `CustomDerivativesLookupRequest` to address the issue.

Resolves swiftlang#60102

Author: kovdan01

include/swift/AST/Module.h

@@ -25,6 +25,7 @@
 #include "swift/AST/Import.h"
 #include "swift/AST/LookupKinds.h"
 #include "swift/AST/Type.h"
+#include "swift/AST/TypeCheckRequests.h"
 #include "swift/Basic/Assertions.h"
 #include "swift/Basic/BasicSourceInfo.h"
 #include "swift/Basic/CXXStdlibKind.h"
@@ -240,6 +241,7 @@ class ModuleDecl
     : public DeclContext, public TypeDecl, public ASTAllocated<ModuleDecl> {
   friend class DirectOperatorLookupRequest;
   friend class DirectPrecedenceGroupLookupRequest;
+  friend class CustomDerivativesRequest;
 
   /// The ABI name of the module, if it differs from the module name.
   mutable Identifier ModuleABIName;

include/swift/AST/TypeCheckRequests.h

@@ -5165,6 +5165,23 @@ class GenericTypeParamDeclGetValueTypeRequest
   bool isCached() const { return true; }
 };
 
+class CustomDerivativesRequest
+    : public SimpleRequest<CustomDerivativesRequest,
+                           evaluator::SideEffect(ModuleDecl *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  evaluator::SideEffect evaluate(Evaluator &evaluator,
+                                 ModuleDecl *module) const;
+
+public:
+  bool isCached() const { return true; }
+};
+
 #define SWIFT_TYPEID_ZONE TypeChecker
 #define SWIFT_TYPEID_HEADER "swift/AST/TypeCheckerTypeIDZone.def"
 #include "swift/Basic/DefineTypeIDZone.h"

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -605,6 +605,9 @@ SWIFT_REQUEST(TypeChecker, ParamCaptureInfoRequest,
 SWIFT_REQUEST(TypeChecker, IsUnsafeRequest,
               bool(Decl *),
               SeparatelyCached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, CustomDerivativesRequest,
+              CustomDerivativesResult(ModuleDecl *),
+	      Cached, NoLocationInfo)
 
 SWIFT_REQUEST(TypeChecker, GenericTypeParamDeclGetValueTypeRequest,
               Type(GenericTypeParamDecl *), Cached, NoLocationInfo)

lib/AST/Module.cpp

@@ -166,6 +166,7 @@ class swift::SourceLookupCache {
   ValueDeclMap TopLevelValues;
   ValueDeclMap ClassMembers;
   bool MemberCachePopulated = false;
+  CustomDerivativesLookupResult CustomDerivatives;
   DeclName UniqueMacroNamePlaceholder;
 
   template<typename T>
@@ -205,6 +206,9 @@ class swift::SourceLookupCache {
   /// guaranteed to be meaningful.
   void getPrecedenceGroups(SmallVectorImpl<PrecedenceGroupDecl *> &results);
 
+  // TODO: is it valid to return const reference from here?
+  llvm::SmallVector<AbstractFunctionDecl *, 0> getCustomDerivativeDecls();
+
   /// Look up an operator declaration.
   ///
   /// \param name The operator name ("+", ">>", etc.)
@@ -275,6 +279,11 @@ void SourceLookupCache::addToUnqualifiedLookupCache(Range items,
         if (!onlyOperators && VD->getAttrs().hasAttribute<CustomAttr>()) {
           MayHaveAuxiliaryDecls.push_back(VD);
         }
+
+        if (!onlyOperators)
+          if (auto *AFD = dyn_cast<AbstractFunctionDecl>(VD))
+            if (AFD->getAttrs().hasAttribute<DerivativeAttr>())
+              CustomDerivatives.push_back(AFD);
       }
     }
 
@@ -572,6 +581,11 @@ void SourceLookupCache::getOperatorDecls(
     results.append(ops.second.begin(), ops.second.end());
 }
 
+llvm::SmallVector<AbstractFunctionDecl *, 0>
+SourceLookupCache::getCustomDerivativeDecls() {
+  return CustomDerivatives;
+}
+
 void SourceLookupCache::lookupOperator(Identifier name, OperatorFixity fixity,
                                        TinyPtrVector<OperatorDecl *> &results) {
   auto ops = Operators.find(name);
@@ -4008,6 +4022,23 @@ bool IsNonUserModuleRequest::evaluate(Evaluator &evaluator, ModuleDecl *mod) con
       (!sdkOrPlatform.empty() && pathStartsWith(sdkOrPlatform, modulePath));
 }
 
+evaluator::SideEffect
+CustomDerivativesRequest::evaluate(Evaluator &evaluator,
+                                   ModuleDecl *module) const {
+  assert(isParsedModule(module));
+  llvm::SmallVector<AbstractFunctionDecl *, 0> decls =
+      module->getSourceLookupCache().getCustomDerivativeDecls();
+  for (const AbstractFunctionDecl *afd : decls) {
+    for (const auto *derAttr :
+         afd->getAttrs().getAttributes<DerivativeAttr>()) {
+      // Resolve derivative function configurations from `@derivative`
+      // attributes by type-checking them.
+      (void)derAttr->getOriginalFunction(SF.getASTContext());
+    }
+  }
+  return {};
+}
+
 version::Version ModuleDecl::getLanguageVersionBuiltWith() const {
   for (auto *F : getFiles()) {
     auto *LD = dyn_cast<LoadedFile>(F);

lib/Sema/TypeChecker.cpp

@@ -404,34 +404,13 @@ void swift::loadDerivativeConfigurations(SourceFile &SF) {
   FrontendStatsTracer tracer(Ctx.Stats,
                              "load-derivative-configurations");
 
-  class DerivativeFinder : public ASTWalker {
-  public:
-    DerivativeFinder() {}
-
-    MacroWalking getMacroWalkingBehavior() const override {
-      return MacroWalking::Expansion;
-    }
-
-    PreWalkAction walkToDeclPre(Decl *D) override {
-      if (auto *afd = dyn_cast<AbstractFunctionDecl>(D)) {
-        for (auto *derAttr : afd->getAttrs().getAttributes<DerivativeAttr>()) {
-          // Resolve derivative function configurations from `@derivative`
-          // attributes by type-checking them.
-          (void)derAttr->getOriginalFunction(D->getASTContext());
-        }
-      }
-
-      return Action::Continue();
-    }
-  };
-
   switch (SF.Kind) {
   case SourceFileKind::DefaultArgument:
   case SourceFileKind::Library:
   case SourceFileKind::MacroExpansion:
   case SourceFileKind::Main: {
-    DerivativeFinder finder;
-    SF.walkContext(finder);
+    CustomDerivativesRequest request(SF.getParentModule());
+    evaluateOrDefault(SF.getASTContext().evaluator, request, {});
     return;
   }
   case SourceFileKind::SIL: