[NFC] Requestify ABI role computation

Author: beccadax

include/swift/AST/ASTContext.h

@@ -415,19 +415,6 @@ class ASTContext final {
       llvm::SmallPtrSet<DerivativeAttr *, 1>>
       DerivativeAttrs;
 
-  /// For each ABI-only declaration (created with the `@abi` attribute), points
-  /// to its API-only counterpart. That is, this table maps each
-  /// `ABIAttr::abiDecl` to the declaration the `ABIAttr` is attached to.
-  ///
-  /// \seeAlso \c recordABIAttr()
-  llvm::DenseMap<Decl *, Decl *> ABIDeclCounterparts;
-
-  /// Register \c this->abiDecl 's  relationship with \p owner in
-  /// `ABIDeclCounterparts` . This is necessary for
-  /// \c ABIRoleInfo::ABIRoleInfo() to determine that \c this->abiDecl
-  /// is ABI-only and locate its API counterpart.
-  void recordABIAttr(ABIAttr *attr, Decl *owner);
-
   /// The Swift module currently being compiled.
   ModuleDecl *MainModule = nullptr;
 

include/swift/AST/Decl.h

@@ -1023,6 +1023,12 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl>, public Swi
   /// attribute macro expansion.
   DeclAttributes getSemanticAttrs() const;
 
+  /// Register the relationship between \c this and \p attr->abiDecl , assuming
+  /// that \p attr is attached to \c this . This is necessary for
+  /// \c ABIRoleInfo::ABIRoleInfo() to determine that \c attr->abiDecl
+  /// is ABI-only and locate its API counterpart.
+  void recordABIAttr(ABIAttr *attr);
+
   /// Set this declaration's attributes to the specified attribute list,
   /// applying any post-processing logic appropriate for attributes parsed
   /// from source code.

include/swift/AST/NameLookupRequests.h

@@ -25,6 +25,7 @@
 #include "swift/AST/TypeOrExtensionDecl.h"
 #include "swift/Basic/Statistic.h"
 #include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/TinyPtrVector.h"
 
 namespace swift {
@@ -1002,6 +1003,32 @@ class ObjCCategoryNameMapRequest
   bool isCached() const { return true; }
 };
 
+struct DeclABIRoleInfoResult {
+  llvm::PointerIntPair<Decl *, 2, uint8_t> storage;
+  DeclABIRoleInfoResult(Decl *counterpart, uint8_t roleValue)
+    : storage(counterpart, roleValue) {}
+};
+
+/// Return the ABI role info (role and counterpart) of a given declaration.
+class DeclABIRoleInfoRequest
+    : public SimpleRequest<DeclABIRoleInfoRequest,
+                           DeclABIRoleInfoResult(Decl *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+  void recordABIOnly(Evaluator &evaluator, Decl *counterpart);
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  DeclABIRoleInfoResult evaluate(Evaluator &evaluator, Decl *decl) const;
+
+public:
+  bool isCached() const { return true; }
+};
+
 #define SWIFT_TYPEID_ZONE NameLookup
 #define SWIFT_TYPEID_HEADER "swift/AST/NameLookupTypeIDZone.def"
 #include "swift/Basic/DefineTypeIDZone.h"

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, DeclABIRoleInfoRequest,
+              DeclABIRoleInfoResult(Decl *), Cached, NoLocationInfo)

lib/AST/Attr.cpp

@@ -2899,30 +2899,6 @@ LifetimeAttr *LifetimeAttr::create(ASTContext &context, SourceLoc atLoc,
   return new (context) LifetimeAttr(atLoc, baseRange, implicit, entry);
 }
 
-void ASTContext::recordABIAttr(ABIAttr *attr, Decl *owner) {
-  // The ABIAttr on a VarDecl ought to point to its PBD.
-  if (auto VD = dyn_cast<VarDecl>(owner)) {
-    if (auto PBD = VD->getParentPatternBinding()) {
-      owner = PBD;
-    }
-  }
-
-  auto recordABIDecl = [&](Decl *decl) {
-    ABIDeclCounterparts.insert({ decl, owner });
-  };
-
-  if (auto abiPBD = dyn_cast<PatternBindingDecl>(attr->abiDecl)) {
-    // Add to *every* VarDecl in the ABI PBD, even ones that don't properly
-    // match anything in the API PBD.
-    for (auto i : range(abiPBD->getNumPatternEntries())) {
-      abiPBD->getPattern(i)->forEachVariable(recordABIDecl);
-    }
-    return;
-  }
-
-  recordABIDecl(attr->abiDecl);
-}
-
 void swift::simple_display(llvm::raw_ostream &out, const DeclAttribute *attr) {
   if (attr)
     attr->print(out);

lib/AST/Decl.cpp

@@ -416,7 +416,7 @@ void Decl::attachParsedAttrs(DeclAttributes attrs) {
   for (auto *attr : attrs.getAttributes<DerivativeAttr>())
     attr->setOriginalDeclaration(this);
   for (auto attr : attrs.getAttributes<ABIAttr, /*AllowInvalid=*/true>())
-    this->getASTContext().recordABIAttr(attr, this);
+    recordABIAttr(attr);
 
   getAttrs() = attrs;
 }
@@ -4271,22 +4271,69 @@ void ValueDecl::setRenamedDecl(const AvailableAttr *attr,
                                                 std::move(renameDecl));
 }
 
-abi_role_detail::Storage abi_role_detail::computeStorage(Decl *decl) {
-  ASSERT(decl);
+void Decl::recordABIAttr(ABIAttr *attr) {
+  Decl *owner = this;
 
-  auto &ctx = decl->getASTContext();
+  // The ABIAttr on a VarDecl ought to point to its PBD.
+  if (auto VD = dyn_cast<VarDecl>(owner)) {
+    if (auto PBD = VD->getParentPatternBinding()) {
+      owner = PBD;
+    }
+  }
+
+  auto record = [&](Decl *decl) {
+    auto &evaluator = owner->getASTContext().evaluator;
+    DeclABIRoleInfoRequest(decl).recordABIOnly(evaluator, owner);
+  };
+
+  if (auto abiPBD = dyn_cast<PatternBindingDecl>(attr->abiDecl)) {
+    // Add to *every* VarDecl in the ABI PBD, even ones that don't properly
+    // match anything in the API PBD.
+    for (auto i : range(abiPBD->getNumPatternEntries())) {
+      abiPBD->getPattern(i)->forEachVariable(record);
+    }
+    return;
+  }
+
+  record(attr->abiDecl);
+}
+
+void DeclABIRoleInfoRequest::recordABIOnly(Evaluator &evaluator,
+                                           Decl *counterpart) {
+  if (evaluator.hasCachedResult(*this))
+    return;
+  DeclABIRoleInfoResult result{counterpart, ABIRole::ProvidesABI};
+  evaluator.cacheOutput(*this, std::move(result));
+}
+
+DeclABIRoleInfoResult
+DeclABIRoleInfoRequest::evaluate(Evaluator &evaluator, Decl *decl) const {
+  // NOTE: ABI decl -> API decl is manually cached through `recordABIOnly()`,
+  // so this code does not have to handle that case.
+
+  ASSERT(decl);
 
   Decl *counterpartDecl = decl;
   ABIRole::Value flags = ABIRole::Either;
 
   if (auto attr = decl->getAttrs().getAttribute<ABIAttr>()) {
     flags = ABIRole::ProvidesAPI;
     counterpartDecl = attr->abiDecl;
-  } else if (auto inverse = ctx.ABIDeclCounterparts.lookup(decl)) {
-    flags = ABIRole::ProvidesABI;
-    counterpartDecl = inverse;
   }
 
+  return DeclABIRoleInfoResult(counterpartDecl, (uint8_t)flags);
+}
+
+abi_role_detail::Storage abi_role_detail::computeStorage(Decl *decl) {
+  ASSERT(decl);
+
+  auto &ctx = decl->getASTContext();
+
+  auto result = evaluateOrDefault(ctx.evaluator, DeclABIRoleInfoRequest{decl},
+                                  { decl, ABIRole::Either });
+  Decl *counterpartDecl = result.storage.getPointer();
+  auto flags = (ABIRole::Value)result.storage.getInt();
+
   // If we did find an `@abi` attribute, resolve PBD pointers to their VarDecl.
   if (flags != ABIRole::Either) {
     if (auto VD = dyn_cast<VarDecl>(decl))

lib/Serialization/Deserialization.cpp

@@ -6513,7 +6513,7 @@ llvm::Error DeclDeserializer::finishRecursiveAttrs(Decl *decl, DeclAttribute *at
     if (!abiDeclOrError)
       return abiDeclOrError.takeError();
     unresolvedABIAttr->first->abiDecl = abiDeclOrError.get();
-    ctx.recordABIAttr(unresolvedABIAttr->first, decl);
+    decl->recordABIAttr(unresolvedABIAttr->first);
   }
   if (ABIDeclCounterpartID != 0) {
     // This decl is the `abiDecl` of an `ABIAttr`. Force the decl that `ABIAttr`