Sema: Requestify Obj-C requirements map computation.

tshortli · tshortli · commit c0303e0e7d11 · 2024-03-11T15:37:00.000-07:00
Avoids repeatedly rebuilding requirement maps during witness resolution.
diff --git a/include/swift/AST/Decl.h b/include/swift/AST/Decl.h
@@ -5074,6 +5074,35 @@ enum class KnownDerivableProtocolKind : uint8_t {
 
 using PrimaryAssociatedTypeName = std::pair<Identifier, SourceLoc>;
 
+/// A wrapper for a dictionary that maps Obj-C protocol requirement selectors to
+/// a list of function decls.
+class ObjCRequirementMap {
+public:
+  using FunctionList = TinyPtrVector<AbstractFunctionDecl *>;
+
+private:
+  using MethodKey = std::pair<ObjCSelector, char>;
+  llvm::SmallDenseMap<MethodKey, FunctionList, 4> storage;
+
+  static MethodKey getObjCMethodKey(AbstractFunctionDecl *func);
+
+public:
+  void addRequirement(AbstractFunctionDecl *requirement) {
+    storage[getObjCMethodKey(requirement)].push_back(requirement);
+  }
+
+  /// Retrieve the Objective-C requirements in this protocol that have the
+  /// given Objective-C method key.
+  FunctionList getRequirements(AbstractFunctionDecl *requirement) const {
+    auto key = getObjCMethodKey(requirement);
+    auto known = storage.find(key);
+    if (known == storage.end())
+      return {};
+
+    return known->second;
+  }
+};
+
 /// ProtocolDecl - A declaration of a protocol, for example:
 ///
 ///   protocol Drawable {
@@ -5303,6 +5332,10 @@ class ProtocolDecl final : public NominalTypeDecl {
   /// i.e., for a protocol P, returns the kind if inverse constraint ~P exists.
   std::optional<InvertibleProtocolKind> getInvertibleProtocolKind() const;
 
+  /// Returns a dictionary that maps Obj-C protocol requirement selectors to a
+  /// list of function decls.
+  ObjCRequirementMap getObjCRequiremenMap() const;
+
 private:
   void computeKnownProtocolKind() const;
 
diff --git a/include/swift/AST/TypeCheckRequests.h b/include/swift/AST/TypeCheckRequests.h
@@ -4828,6 +4828,23 @@ class LocalTypeDeclsRequest
   bool isCached() const { return true; }
 };
 
+class ObjCRequirementMapRequest
+    : public SimpleRequest<ObjCRequirementMapRequest,
+                           ObjCRequirementMap(const ProtocolDecl *proto),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  ObjCRequirementMap evaluate(Evaluator &evaluator,
+                              const ProtocolDecl *proto) 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"
diff --git a/include/swift/AST/TypeCheckerTypeIDZone.def b/include/swift/AST/TypeCheckerTypeIDZone.def
@@ -559,3 +559,7 @@ SWIFT_REQUEST(TypeChecker, UniqueUnderlyingTypeSubstitutionsRequest,
 SWIFT_REQUEST(TypeChecker, LocalTypeDeclsRequest,
               ArrayRef<TypeDecl *>(SourceFile *),
               Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, ObjCRequirementMapRequest,
+              ObjCRequirementMap(const ProtocolDecl *proto),
+              Cached, NoLocationInfo)
+
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -6467,6 +6467,15 @@ ProtocolDecl::getInvertibleProtocolKind() const {
   return std::nullopt;
 }
 
+ObjCRequirementMap ProtocolDecl::getObjCRequiremenMap() const {
+  ObjCRequirementMap defaultMap;
+  if (!isObjC())
+    return defaultMap;
+
+  return evaluateOrDefault(getASTContext().evaluator,
+                           ObjCRequirementMapRequest{this}, defaultMap);
+}
+
 ArrayRef<ProtocolDecl *> ProtocolDecl::getInheritedProtocols() const {
   auto *mutThis = const_cast<ProtocolDecl *>(this);
   return evaluateOrDefault(getASTContext().evaluator,
diff --git a/lib/Sema/TypeCheckProtocol.cpp b/lib/Sema/TypeCheckProtocol.cpp
@@ -1840,54 +1840,39 @@ RequirementCheck WitnessChecker::checkWitness(ValueDecl *requirement,
 # pragma mark Witness resolution
 
 /// Retrieve the Objective-C method key from the given function.
-namespace {
-  using ObjCMethodKey = std::pair<ObjCSelector, char>;
-  using ObjCRequirementMap = llvm::SmallDenseMap<ObjCMethodKey,
-                                                 TinyPtrVector<AbstractFunctionDecl *>, 4>;
-}
-
-/// Retrieve the Objective-C method key from the given function.
-static ObjCMethodKey getObjCMethodKey(AbstractFunctionDecl *func) {
+ObjCRequirementMap::MethodKey
+ObjCRequirementMap::getObjCMethodKey(AbstractFunctionDecl *func) {
   return std::make_pair(func->getObjCSelector(), func->isInstanceMember());
 }
 
 /// Precompute map for getObjCRequirements().
-static ObjCRequirementMap getObjCRequirementMap(ProtocolDecl *proto) {
-  ObjCRequirementMap map;
-
-  if (!proto->isObjC())
-    return map;
+ObjCRequirementMap
+ObjCRequirementMapRequest::evaluate(Evaluator &evaluator,
+                                    const ProtocolDecl *proto) const {
+  // This map only applies to Obj-C protocols so it's wasteful to evaluate this
+  // request and cache the result for non-Obj-C protocols.
+  assert(proto->isObjC());
 
+  ObjCRequirementMap map;
   for (auto requirement : proto->getProtocolRequirements()) {
     auto funcRequirement = dyn_cast<AbstractFunctionDecl>(requirement);
     if (!funcRequirement)
       continue;
 
-    map[getObjCMethodKey(funcRequirement)].push_back(funcRequirement);
+    map.addRequirement(funcRequirement);
   }
 
   return map;
 }
 
-/// Retrieve the Objective-C requirements in this protocol that have the
-/// given Objective-C method key.
-static ArrayRef<AbstractFunctionDecl *>
-getObjCRequirements(const ObjCRequirementMap &map, ObjCMethodKey key) {
-  auto known = map.find(key);
-  if (known == map.end())
-    return { };
-
-  return known->second;
-}
-
 /// @returns a non-null requirement if the given requirement is part of a
 /// group of ObjC requirements that share the same ObjC method key.
 /// The first such requirement that the predicate function returns true for
 /// is the requirement required by this function. Otherwise, nullptr is
 /// returned.
 static ValueDecl *getObjCRequirementSibling(
-    ProtocolDecl *proto, ValueDecl *requirement, const ObjCRequirementMap &map,
-    llvm::function_ref<bool(AbstractFunctionDecl*)> predicate) {
+    ProtocolDecl *proto, ValueDecl *requirement,
+    llvm::function_ref<bool(AbstractFunctionDecl *)> predicate) {
   if (!proto->isObjC())
     return nullptr;
 
@@ -1896,8 +1881,8 @@ static ValueDecl *getObjCRequirementSibling(
 
   // We only care about functions
   if (auto fnRequirement = dyn_cast<AbstractFunctionDecl>(requirement)) {
-    auto fnSelector = getObjCMethodKey(fnRequirement);
-    auto similarRequirements = getObjCRequirements(map, fnSelector);
+    auto map = proto->getObjCRequiremenMap();
+    auto similarRequirements = map.getRequirements(fnRequirement);
     // ... whose selector is one that maps to multiple requirement declarations.
     for (auto candidate : similarRequirements) {
       if (candidate == fnRequirement)
@@ -1948,9 +1933,8 @@ class MultiConformanceChecker {
     NormalProtocolConformance *conformance);
 
   /// Determine whether the given requirement was left unsatisfied.
-  bool isUnsatisfiedReq(
-      NormalProtocolConformance *conformance, ValueDecl *req,
-      const ObjCRequirementMap &map);
+  bool isUnsatisfiedReq(NormalProtocolConformance *conformance, ValueDecl *req);
+
 public:
   MultiConformanceChecker(ASTContext &ctx) : Context(ctx) {}
 
@@ -1977,9 +1961,8 @@ class MultiConformanceChecker {
 
 }
 
-bool MultiConformanceChecker::
-isUnsatisfiedReq(NormalProtocolConformance *conformance, ValueDecl *req,
-                 const ObjCRequirementMap &map) {
+bool MultiConformanceChecker::isUnsatisfiedReq(
+    NormalProtocolConformance *conformance, ValueDecl *req) {
   if (conformance->isInvalid()) return false;
   if (isa<TypeDecl>(req)) return false;
 
@@ -1993,7 +1976,7 @@ isUnsatisfiedReq(NormalProtocolConformance *conformance, ValueDecl *req,
     // If another @objc requirement refers to the same Objective-C
     // method, this requirement isn't unsatisfied.
     if (getObjCRequirementSibling(
-            proto, req, map, [conformance](AbstractFunctionDecl *cand) {
+            proto, req, [conformance](AbstractFunctionDecl *cand) {
               return static_cast<bool>(conformance->getWitness(cand));
             })) {
       return false;
@@ -2034,12 +2017,11 @@ void MultiConformanceChecker::checkAllConformances() {
     if (!anyInvalid) {
       // Check whether there are any unsatisfied requirements.
       auto proto = conformance->getProtocol();
-      ObjCRequirementMap map = getObjCRequirementMap(proto);
 
       for (auto *req : proto->getProtocolRequirements()) {
         // If the requirement is unsatisfied, we might want to warn
         // about near misses; record it.
-        if (isUnsatisfiedReq(conformance, req, map)) {
+        if (isUnsatisfiedReq(conformance, req)) {
           UnsatisfiedReqs.push_back(req);
           continue;
         }
@@ -3726,10 +3708,8 @@ hasSatisfiedObjCSiblingRequirement(ProtocolDecl *proto,
   if (!proto->isObjC())
     return false;
 
-  auto map = getObjCRequirementMap(proto);
-
   if (getObjCRequirementSibling(
-          proto, fnRequirement, map,
+          proto, fnRequirement,
           [proto, conformance](AbstractFunctionDecl *candidate) {
             // FIXME: This performs a recursive lookup in the lazy case, so
             // we have to dodge the cycle.
@@ -3962,6 +3942,27 @@ getAdopteeSelfSameTypeConstraint(ClassDecl *selfClass, ValueDecl *witness) {
   return target;
 }
 
+static bool allowOptionalWitness(ProtocolDecl *proto,
+                                 NormalProtocolConformance *conformance,
+                                 ValueDecl *requirement) {
+  auto Attrs = requirement->getAttrs();
+
+  // An optional requirement is trivially satisfied with an empty requirement.
+  if (Attrs.hasAttribute<OptionalAttr>())
+    return true;
+
+  // An 'unavailable' requirement is treated like an optional requirement.
+  if (Attrs.isUnavailable(proto->getASTContext()))
+    return true;
+
+  // A requirement with a satisfied Obj-C alternative requirement is effectively
+  // optional.
+  if (hasSatisfiedObjCSiblingRequirement(proto, conformance, requirement))
+    return true;
+
+  return false;
+}
+
 void ConformanceChecker::checkNonFinalClassWitness(ValueDecl *requirement,
                                                    ValueDecl *witness) {
   auto *classDecl = DC->getSelfClassDecl();
@@ -4388,7 +4389,7 @@ ConformanceChecker::resolveWitnessViaLookup(ValueDecl *requirement) {
   //
   // Treat 'unavailable' implicitly as if it were 'optional'.
   // The compiler will reject actual uses.
-  if (allowOptionalWitness(requirement)) {
+  if (allowOptionalWitness(Proto, Conformance, requirement)) {
     return ResolveWitnessResult::Missing;
   }
 
@@ -4561,7 +4562,7 @@ ResolveWitnessResult ConformanceChecker::resolveWitnessViaDefault(
                        ValueDecl *requirement) {
   assert(!isa<AssociatedTypeDecl>(requirement) && "Use resolveTypeWitnessVia*");
 
-  if (allowOptionalWitness(requirement)) {
+  if (allowOptionalWitness(Proto, Conformance, requirement)) {
     recordOptionalWitness(requirement);
     return ResolveWitnessResult::Success;
   }
@@ -4632,25 +4633,6 @@ void ConformanceChecker::resolveSingleWitness(ValueDecl *requirement) {
   }
 }
 
-bool ConformanceChecker::allowOptionalWitness(ValueDecl *requirement) {
-  auto Attrs = requirement->getAttrs();
-
-  // An optional requirement is trivially satisfied with an empty requirement.
-  if (Attrs.hasAttribute<OptionalAttr>())
-    return true;
-
-  // An 'unavailable' requirement is treated like an optional requirement.
-  if (Attrs.isUnavailable(getASTContext()))
-    return true;
-
-  // A requirement with a satisfied Obj-C alternative requirement is effectively
-  // optional.
-  if (hasSatisfiedObjCSiblingRequirement(Proto, Conformance, requirement))
-    return true;
-
-  return false;
-}
-
 /// FIXME: It feels like this could be part of findExistentialSelfReferences().
 static std::optional<Requirement>
 hasInvariantSelfRequirement(const ProtocolDecl *proto,
@@ -4973,7 +4955,6 @@ hasInvalidTypeInConformanceContext(const ValueDecl *requirement,
 
 void ConformanceChecker::resolveValueWitnesses() {
   bool usesPreconcurrencyConformance = false;
-  auto objcRequirementMap = getObjCRequirementMap(Proto);
 
   for (auto *requirement : Proto->getProtocolRequirements()) {
     // Associated type requirements handled elsewhere.
@@ -5119,7 +5100,7 @@ void ConformanceChecker::resolveValueWitnesses() {
     // async-looking ObjC protocol method requirement into two Swift protocol
     // requirements: an async version and a sync version. Exactly one of the two
     // must be witnessed by the conformer.
-    if (getObjCRequirementSibling(Proto, requirement, objcRequirementMap,
+    if (getObjCRequirementSibling(Proto, requirement,
                                   [this](AbstractFunctionDecl *cand) {
                                     return static_cast<bool>(
                                         this->Conformance->getWitness(cand));
diff --git a/lib/Sema/TypeCheckProtocol.h b/lib/Sema/TypeCheckProtocol.h
@@ -168,9 +168,6 @@ class ConformanceChecker : public WitnessChecker {
   /// particular requirement and adoptee is required, before the
   /// conformance has been completed checked.
   void resolveSingleWitness(ValueDecl *requirement);
-
-private:
-  bool allowOptionalWitness(ValueDecl *requirement);
 };
 
 /// Match the given witness to the given requirement.
