Sema: Don't need to pass the ConformanceChecker around so much

slavapestov · slavapestov · commit 5249432dcc39 · 2024-01-17T13:08:46.000-05:00
I want to slim down ConformanceChecker because it doesn't really fit
with request-driven resolution of witnesses. This is a first step
towards that.
diff --git a/include/swift/AST/TypeCheckRequests.h b/include/swift/AST/TypeCheckRequests.h
@@ -2829,6 +2829,25 @@ class TypeWitnessRequest
   void cacheResult(TypeWitnessAndDecl value) const;
 };
 
+class ReferencedAssociatedTypesRequest
+    : public SimpleRequest<ReferencedAssociatedTypesRequest,
+                           TinyPtrVector<AssociatedTypeDecl *>(ValueDecl *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  TinyPtrVector<AssociatedTypeDecl *>
+  evaluate(Evaluator &evaluator, ValueDecl *req) const;
+
+public:
+  // Caching.
+  bool isCached() const { return true; }
+};
+
 class ValueWitnessRequest
     : public SimpleRequest<ValueWitnessRequest,
                            Witness(NormalProtocolConformance *, ValueDecl *),
diff --git a/include/swift/AST/TypeCheckerTypeIDZone.def b/include/swift/AST/TypeCheckerTypeIDZone.def
@@ -426,6 +426,9 @@ SWIFT_REQUEST(TypeChecker, AssociatedConformanceRequest,
               ProtocolConformanceRef(NormalProtocolConformance *,
                                      Type, ProtocolDecl *, unsigned),
               SeparatelyCached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, ReferencedAssociatedTypesRequest,
+              TinyPtrVector<AssociatedTypeDecl *>(ValueDecl *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ValueWitnessRequest,
               Witness(NormalProtocolConformance *, ValueDecl *),
               SeparatelyCached, NoLocationInfo)
diff --git a/lib/Sema/TypeCheckProtocol.cpp b/lib/Sema/TypeCheckProtocol.cpp
@@ -1359,9 +1359,9 @@ WitnessChecker::WitnessChecker(ASTContext &ctx, ProtocolDecl *proto,
                                Type adoptee, DeclContext *dc)
     : Context(ctx), Proto(proto), Adoptee(adoptee), DC(dc) {}
 
-void
-WitnessChecker::lookupValueWitnessesViaImplementsAttr(
-    ValueDecl *req, SmallVector<ValueDecl *, 4> &witnesses) {
+static void
+lookupValueWitnessesViaImplementsAttr(
+    DeclContext *DC, ValueDecl *req, SmallVector<ValueDecl *, 4> &witnesses) {
 
   auto name = req->createNameRef();
   auto *nominal = DC->getSelfNominalTypeDecl();
@@ -1411,15 +1411,15 @@ static bool contextMayExpandOperator(
 }
 
 SmallVector<ValueDecl *, 4>
-WitnessChecker::lookupValueWitnesses(ValueDecl *req, bool *ignoringNames) {
+swift::lookupValueWitnesses(DeclContext *DC, ValueDecl *req, bool *ignoringNames) {
   assert(!isa<AssociatedTypeDecl>(req) && "Not for lookup for type witnesses*");
   assert(req->isProtocolRequirement());
 
   SmallVector<ValueDecl *, 4> witnesses;
 
   // Do an initial check to see if there are any @_implements remappings
   // for this requirement.
-  lookupValueWitnessesViaImplementsAttr(req, witnesses);
+  lookupValueWitnessesViaImplementsAttr(DC, req, witnesses);
 
   auto reqName = req->createNameRef();
   auto reqBaseName = reqName.withoutArgumentLabels();
@@ -1441,7 +1441,7 @@ WitnessChecker::lookupValueWitnesses(ValueDecl *req, bool *ignoringNames) {
     for (auto candidate : lookup) {
       auto decl = candidate.getValueDecl();
       if (!isa<ProtocolDecl>(decl->getDeclContext()) &&
-          swift::isMemberOperator(cast<FuncDecl>(decl), Adoptee)) {
+          swift::isMemberOperator(cast<FuncDecl>(decl), DC->getSelfInterfaceType())) {
         witnesses.push_back(decl);
       }
     }
@@ -1515,7 +1515,7 @@ bool WitnessChecker::findBestWitness(
     SmallVector<ValueDecl *, 4> witnesses;
     switch (attempt) {
     case Regular:
-      witnesses = lookupValueWitnesses(requirement, ignoringNames);
+      witnesses = lookupValueWitnesses(DC, requirement, ignoringNames);
       break;
     case OperatorsFromOverlay: {
       // If we have a Clang declaration, the matching operator might be in the
@@ -2924,29 +2924,26 @@ ConformanceChecker::~ConformanceChecker() {
             !getASTContext().hasDelayedConformanceErrors(Conformance));
 }
 
-ArrayRef<AssociatedTypeDecl *>
-ConformanceChecker::getReferencedAssociatedTypes(ValueDecl *req) {
-  // Check whether we've already cached this information.
-  auto known = ReferencedAssociatedTypes.find(req);
-  if (known != ReferencedAssociatedTypes.end())
-    return known->second;
 
+TinyPtrVector<AssociatedTypeDecl *>
+ReferencedAssociatedTypesRequest::evaluate(Evaluator &eval,
+                                           ValueDecl *req) const {
   // Collect the set of associated types rooted on Self in the
   // signature. Note that for references to nested types, we only
   // want to consider the outermost dependent member type.
   //
   // For example, a requirement typed '(Iterator.Element) -> ()'
   // is not considered to reference the associated type 'Iterator'.
-  auto &assocTypes = ReferencedAssociatedTypes[req];
+  TinyPtrVector<AssociatedTypeDecl *> assocTypes;
 
   class Walker : public TypeWalker {
     ProtocolDecl *Proto;
-    llvm::SmallVectorImpl<AssociatedTypeDecl *> &assocTypes;
+    llvm::TinyPtrVector<AssociatedTypeDecl *> &assocTypes;
     llvm::SmallPtrSet<AssociatedTypeDecl *, 4> knownAssocTypes;
 
   public:
     Walker(ProtocolDecl *Proto,
-           llvm::SmallVectorImpl<AssociatedTypeDecl *> &assocTypes)
+           llvm::TinyPtrVector<AssociatedTypeDecl *> &assocTypes)
       : Proto(Proto), assocTypes(assocTypes) {}
 
     Action walkToTypePre(Type type) override {
@@ -2963,7 +2960,7 @@ ConformanceChecker::getReferencedAssociatedTypes(ValueDecl *req) {
     }
   };
 
-  Walker walker(Proto, assocTypes);
+  Walker walker(cast<ProtocolDecl>(req->getDeclContext()), assocTypes);
 
   // This dance below is to avoid calling getCanonicalType() on a
   // GenericFunctionType, which creates a GenericSignatureBuilder, which
@@ -4301,7 +4298,10 @@ ConformanceChecker::resolveWitnessViaLookup(ValueDecl *requirement) {
 
   // If any of the type witnesses was erroneous, don't bother to check
   // this value witness: it will fail.
-  for (auto assocType : getReferencedAssociatedTypes(requirement)) {
+  auto referenced = evaluateOrDefault(getASTContext().evaluator,
+                                      ReferencedAssociatedTypesRequest{requirement},
+                                      TinyPtrVector<AssociatedTypeDecl *>());
+  for (auto assocType : referenced) {
     if (Conformance->getTypeWitness(assocType)->hasError()) {
       return ResolveWitnessResult::ExplicitFailed;
     }
diff --git a/lib/Sema/TypeCheckProtocol.h b/lib/Sema/TypeCheckProtocol.h
@@ -636,6 +636,19 @@ struct RequirementMatch {
   swift::Witness getWitness(ASTContext &ctx) const;
 };
 
+/// Gather the value witnesses for the given requirement.
+///
+/// \param DC A nominal type or extension context where the conformance
+/// was declared.
+/// \param req A member of a protocol that DC conforms to.
+/// \param ignoringNames If non-null and there are no value
+/// witnesses with the correct full name, the results will reflect
+/// lookup for just the base name and the pointee will be set to
+/// \c true.
+SmallVector<ValueDecl *, 4> lookupValueWitnesses(DeclContext *DC,
+                                                 ValueDecl *req,
+                                                 bool *ignoringNames);
+
 struct RequirementCheck;
 
 class WitnessChecker {
@@ -676,15 +689,6 @@ class WitnessChecker {
     return RequiredAccessScopeAndUsableFromInline.value().second;
   }
 
-  /// Gather the value witnesses for the given requirement.
-  ///
-  /// \param ignoringNames If non-null and there are no value
-  /// witnesses with the correct full name, the results will reflect
-  /// lookup for just the base name and the pointee will be set to
-  /// \c true.
-  SmallVector<ValueDecl *, 4> lookupValueWitnesses(ValueDecl *req,
-                                                   bool *ignoringNames);
-
   void lookupValueWitnessesViaImplementsAttr(ValueDecl *req,
                                              SmallVector<ValueDecl *, 4>
                                              &witnesses);
@@ -774,11 +778,6 @@ class ConformanceChecker : public WitnessChecker {
   /// Witnesses that are currently being resolved.
   llvm::SmallPtrSet<ValueDecl *, 4> ResolvingWitnesses;
 
-  /// Caches the set of associated types that are referenced in each
-  /// requirement.
-  llvm::DenseMap<ValueDecl *, llvm::SmallVector<AssociatedTypeDecl *, 2>>
-    ReferencedAssociatedTypes;
-
   /// Keep track of missing witnesses, either type or value, for later
   /// diagnosis emits. This may contain witnesses that are external to the
   /// protocol under checking.
@@ -804,10 +803,6 @@ class ConformanceChecker : public WitnessChecker {
   /// Whether objcMethodRequirements has been computed.
   bool computedObjCMethodRequirements = false;
 
-  /// Retrieve the associated types that are referenced by the given
-  /// requirement with a base of 'Self'.
-  ArrayRef<AssociatedTypeDecl *> getReferencedAssociatedTypes(ValueDecl *req);
-
   /// Record a (non-type) witness for the given requirement.
   void recordWitness(ValueDecl *requirement, const RequirementMatch &match);
 
@@ -1132,13 +1127,11 @@ class AssociatedTypeInference {
 
   /// Infer associated type witnesses for the given value requirement.
   InferredAssociatedTypesByWitnesses inferTypeWitnessesViaValueWitnesses(
-                   ConformanceChecker &checker,
                    const llvm::SetVector<AssociatedTypeDecl *> &allUnresolved,
                    ValueDecl *req);
 
   /// Infer associated type witnesses for the given associated type.
   InferredAssociatedTypesByWitnesses inferTypeWitnessesViaAssociatedType(
-                   ConformanceChecker &checker,
                    const llvm::SetVector<AssociatedTypeDecl *> &allUnresolved,
                    AssociatedTypeDecl *assocType);
 
@@ -1147,7 +1140,6 @@ class AssociatedTypeInference {
   /// \param assocTypes The set of associated types we're interested in.
   InferredAssociatedTypes
   inferTypeWitnessesViaValueWitnesses(
-    ConformanceChecker &checker,
     const llvm::SetVector<AssociatedTypeDecl *> &assocTypes);
 
   /// Compute a "fixed" type witness for an associated type, e.g.,
@@ -1255,8 +1247,7 @@ class AssociatedTypeInference {
   ///   involving multiple files/modules, and not in tests within the Swift
   ///   project itself.
   bool canAttemptEagerTypeWitnessDerivation(
-      ConformanceChecker &checker,
-      AssociatedTypeDecl *assocType);
+      DeclContext *DC, AssociatedTypeDecl *assocType);
 
 public:
   /// Describes a mapping from associated type declarations to their
diff --git a/lib/Sema/TypeCheckProtocolInference.cpp b/lib/Sema/TypeCheckProtocolInference.cpp
@@ -25,6 +25,7 @@
 #include "swift/AST/SubstitutionMap.h"
 #include "swift/AST/TypeMatcher.h"
 #include "swift/AST/Types.h"
+#include "swift/AST/TypeCheckRequests.h"
 #include "swift/Basic/Defer.h"
 #include "swift/ClangImporter/ClangModule.h"
 #include "llvm/ADT/Statistic.h"
@@ -374,7 +375,6 @@ static bool isExtensionUsableForInference(const ExtensionDecl *extension,
 
 InferredAssociatedTypesByWitnesses
 AssociatedTypeInference::inferTypeWitnessesViaValueWitnesses(
-                    ConformanceChecker &checker,
                     const llvm::SetVector<AssociatedTypeDecl *> &allUnresolved,
                     ValueDecl *req) {
   // Conformances constructed by the ClangImporter should have explicit type
@@ -392,7 +392,7 @@ AssociatedTypeInference::inferTypeWitnessesViaValueWitnesses(
   InferredAssociatedTypesByWitnesses result;
 
   for (auto witness :
-       checker.lookupValueWitnesses(req, /*ignoringNames=*/nullptr)) {
+       lookupValueWitnesses(dc, req, /*ignoringNames=*/nullptr)) {
     LLVM_DEBUG(llvm::dbgs() << "Inferring associated types from decl:\n";
                witness->dump(llvm::dbgs()));
 
@@ -582,9 +582,7 @@ next_witness:;
 
 InferredAssociatedTypes
 AssociatedTypeInference::inferTypeWitnessesViaValueWitnesses(
-  ConformanceChecker &checker,
-  const llvm::SetVector<AssociatedTypeDecl *> &assocTypes)
-{
+  const llvm::SetVector<AssociatedTypeDecl *> &assocTypes) {
   InferredAssociatedTypes result;
   for (auto member : proto->getMembers()) {
     auto req = dyn_cast<ValueDecl>(member);
@@ -598,8 +596,7 @@ AssociatedTypeInference::inferTypeWitnessesViaValueWitnesses(
       if (assocTypes.count(assocType) == 0)
         continue;
 
-      auto reqInferred = inferTypeWitnessesViaAssociatedType(checker,
-                                                             assocTypes,
+      auto reqInferred = inferTypeWitnessesViaAssociatedType(assocTypes,
                                                              assocType);
       if (!reqInferred.empty())
         result.push_back({req, std::move(reqInferred)});
@@ -623,7 +620,9 @@ AssociatedTypeInference::inferTypeWitnessesViaValueWitnesses(
     // Check whether any of the associated types we care about are
     // referenced in this value requirement.
     {
-      const auto referenced = checker.getReferencedAssociatedTypes(req);
+      auto referenced = evaluateOrDefault(ctx.evaluator,
+                                          ReferencedAssociatedTypesRequest{req},
+                                          TinyPtrVector<AssociatedTypeDecl *>());
       if (llvm::find_if(referenced, [&](AssociatedTypeDecl *const assocType) {
                           return assocTypes.count(assocType);
                         }) == referenced.end())
@@ -633,7 +632,7 @@ AssociatedTypeInference::inferTypeWitnessesViaValueWitnesses(
     // Infer associated types from the potential value witnesses for
     // this requirement.
     auto reqInferred =
-      inferTypeWitnessesViaValueWitnesses(checker, assocTypes, req);
+      inferTypeWitnessesViaValueWitnesses(assocTypes, req);
     if (reqInferred.empty())
       continue;
 
@@ -738,7 +737,6 @@ static Type removeSelfParam(ValueDecl *value, Type type) {
 
 InferredAssociatedTypesByWitnesses
 AssociatedTypeInference::inferTypeWitnessesViaAssociatedType(
-                   ConformanceChecker &checker,
                    const llvm::SetVector<AssociatedTypeDecl *> &allUnresolved,
                    AssociatedTypeDecl *assocType) {
   // Form the default name _Default_Foo.
@@ -2513,13 +2511,13 @@ bool AssociatedTypeInference::diagnoseAmbiguousSolutions(
 }
 
 bool AssociatedTypeInference::canAttemptEagerTypeWitnessDerivation(
-    ConformanceChecker &checker,
-    AssociatedTypeDecl *assocType) {
+    DeclContext *DC, AssociatedTypeDecl *assocType) {
 
   /// Rather than locating the TypeID via the default implementation of
   /// Identifiable, we need to find the type based on the associated ActorSystem
-  if (checker.Adoptee->isDistributedActor() &&
-      assocType->getProtocol()->isSpecificProtocol(KnownProtocolKind::Identifiable)) {
+  if (auto *nominal = DC->getSelfNominalTypeDecl())
+    if (nominal->isDistributedActor() &&
+        assocType->getProtocol()->isSpecificProtocol(KnownProtocolKind::Identifiable)) {
     return true;
   }
 
@@ -2550,7 +2548,7 @@ auto AssociatedTypeInference::solve(ConformanceChecker &checker)
     if (conformance->hasTypeWitness(assocType))
       continue;
 
-    if (canAttemptEagerTypeWitnessDerivation(checker, assocType)) {
+    if (canAttemptEagerTypeWitnessDerivation(dc, assocType)) {
       auto derivedType = computeDerivedTypeWitness(assocType);
       if (derivedType.first) {
         checker.recordTypeWitness(assocType,
@@ -2602,8 +2600,7 @@ auto AssociatedTypeInference::solve(ConformanceChecker &checker)
     return result;
 
   // Infer potential type witnesses from value witnesses.
-  inferred = inferTypeWitnessesViaValueWitnesses(checker,
-                                                 unresolvedAssocTypes);
+  inferred = inferTypeWitnessesViaValueWitnesses(unresolvedAssocTypes);
   LLVM_DEBUG(llvm::dbgs() << "Candidates for inference:\n";
              dumpInferredAssociatedTypes(inferred));
 
@@ -3036,7 +3033,11 @@ void ConformanceChecker::resolveSingleWitness(ValueDecl *requirement) {
 
   // If any of the type witnesses was erroneous, don't bother to check
   // this value witness: it will fail.
-  for (auto assocType : getReferencedAssociatedTypes(requirement)) {
+  auto assocTypes = evaluateOrDefault(getASTContext().evaluator,
+                                      ReferencedAssociatedTypesRequest{requirement},
+                                      TinyPtrVector<AssociatedTypeDecl *>());
+
+  for (auto assocType : assocTypes) {
     if (Conformance->getTypeWitness(assocType)->hasError()) {
       Conformance->setInvalid();
       return;
