[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

@@ -240,6 +240,7 @@ class ModuleDecl
     : public DeclContext, public TypeDecl, public ASTAllocated<ModuleDecl> {
   friend class DirectOperatorLookupRequest;
   friend class DirectPrecedenceGroupLookupRequest;
+  friend class CustomDerivativesLookupRequest;
 
   /// The ABI name of the module, if it differs from the module name.
   mutable Identifier ModuleABIName;

include/swift/AST/NameLookupRequests.h

@@ -969,6 +969,22 @@ class LookupIntrinsicRequest
   bool isCached() const { return true; }
 };
 
+using CustomDerivativesLookupResult = SmallVector<AbstractFunctionDecl *, 0>;
+
+class CustomDerivativesLookupRequest
+    : public SimpleRequest<CustomDerivativesLookupRequest,
+                           CustomDerivativesLookupResult(ModuleDecl *),
+                           RequestFlags::Uncached /*TODO*/> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  CustomDerivativesLookupResult evaluate(Evaluator &evaluator,
+                                         ModuleDecl *module) const;
+};
+
 using ObjCCategoryNameMap =
   llvm::DenseMap<Identifier, llvm::TinyPtrVector<ExtensionDecl *>>;
 

include/swift/AST/NameLookupTypeIDZone.def

@@ -118,3 +118,5 @@ SWIFT_REQUEST(NameLookup, LookupIntrinsicRequest,
               FuncDecl *(ModuleDecl *, Identifier), Cached, NoLocationInfo)
 SWIFT_REQUEST(NameLookup, ObjCCategoryNameMapRequest,
               ObjCCategoryNameMap(ClassDecl *, ExtensionDecl *), Cached, NoLocationInfo)
+SWIFT_REQUEST(NameLookup, CustomDerivativesLookupRequest,
+              CustomDerivativesLookupResult(ModuleDecl *), Uncached, 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>
@@ -175,7 +176,7 @@ class swift::SourceLookupCache {
 
   template<typename Range>
   void addToUnqualifiedLookupCache(Range decls, bool onlyOperators);
-  template<typename Range>
+  template <typename Range>
   void addToMemberCache(Range decls);
 
   using AuxiliaryDeclMap = llvm::DenseMap<DeclName, TinyPtrVector<MissingDecl *>>;
@@ -205,6 +206,8 @@ class swift::SourceLookupCache {
   /// guaranteed to be meaningful.
   void getPrecedenceGroups(SmallVectorImpl<PrecedenceGroupDecl *> &results);
 
+  CustomDerivativesLookupResult getCustomDerivativeDecls();
+
   /// Look up an operator declaration.
   ///
   /// \param name The operator name ("+", ">>", etc.)
@@ -275,6 +278,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 +580,10 @@ void SourceLookupCache::getOperatorDecls(
     results.append(ops.second.begin(), ops.second.end());
 }
 
+CustomDerivativesLookupResult SourceLookupCache::getCustomDerivativeDecls() {
+  return CustomDerivatives;
+}
+
 void SourceLookupCache::lookupOperator(Identifier name, OperatorFixity fixity,
                                        TinyPtrVector<OperatorDecl *> &results) {
   auto ops = Operators.find(name);
@@ -1578,6 +1590,13 @@ void SourceFile::lookupOperatorDirect(
   getCache().lookupOperator(name, fixity, results);
 }
 
+CustomDerivativesLookupResult
+CustomDerivativesLookupRequest::evaluate(Evaluator &evaluator,
+                                         ModuleDecl *module) const {
+  assert(isParsedModule(module));
+  return module->getSourceLookupCache().getCustomDerivativeDecls();
+}
+
 void DirectPrecedenceGroupLookupRequest::writeDependencySink(
     evaluator::DependencyCollector &reqTracker,
     const TinyPtrVector<PrecedenceGroupDecl *> &groups) const {

lib/Sema/TypeChecker.cpp

@@ -34,6 +34,7 @@
 #include "swift/AST/Initializer.h"
 #include "swift/AST/ModuleLoader.h"
 #include "swift/AST/NameLookup.h"
+#include "swift/AST/NameLookupRequests.h"
 #include "swift/AST/PrettyStackTrace.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/SourceFile.h"
@@ -404,34 +405,17 @@ 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);
+    CustomDerivativesLookupRequest request(SF.getParentModule());
+    CustomDerivativesLookupResult result =
+        evaluateOrDefault(SF.getASTContext().evaluator, request, {});
+    for (AbstractFunctionDecl *AFD : result)
+      for (auto *derAttr : AFD->getAttrs().getAttributes<DerivativeAttr>())
+        (void)derAttr->getOriginalFunction(SF.getASTContext());
     return;
   }
   case SourceFileKind::SIL: