Cut down conformance checking time by reusing RequirementEnvironments (#18598)

Trade an extra copy for saving the cost of reinitialization when there
are a lot of witnesses to check. A small but measurable win when
type-checking the standard library.

Also add a test to make sure we still get the requirement environment
right for covariant 'Self'.

Author: jrose-apple

lib/Sema/TypeCheckProtocol.cpp

@@ -646,12 +646,32 @@ swift::matchWitness(
   return finalize(anyRenaming, optionalAdjustments);
 }
 
-RequirementMatch swift::matchWitness(TypeChecker &tc,
-                                     ProtocolDecl *proto,
-                                     ProtocolConformance *conformance,
-                                     DeclContext *dc,
-                                     ValueDecl *req,
-                                     ValueDecl *witness) {
+/// Checks \p reqEnvCache for a requirement environment appropriate for
+/// \p reqSig and \p covariantSelf. If one isn't there, it gets created from
+/// the rest of the parameters.
+///
+/// Note that this means RequirementEnvironmentCaches must not be shared across
+/// multiple protocols or conformances.
+static const RequirementEnvironment &getOrCreateRequirementEnvironment(
+    WitnessChecker::RequirementEnvironmentCache &reqEnvCache,
+    DeclContext *dc, GenericSignature *reqSig, ProtocolDecl *proto,
+    ClassDecl *covariantSelf, ProtocolConformance *conformance) {
+  WitnessChecker::RequirementEnvironmentCacheKey cacheKey(reqSig,
+                                                          covariantSelf);
+  auto cacheIter = reqEnvCache.find(cacheKey);
+  if (cacheIter == reqEnvCache.end()) {
+    RequirementEnvironment reqEnv(dc, reqSig, proto, covariantSelf,
+                                  conformance);
+    cacheIter = reqEnvCache.insert({cacheKey, std::move(reqEnv)}).first;
+  }
+  return cacheIter->getSecond();
+}
+
+RequirementMatch
+swift::matchWitness(TypeChecker &tc,
+                    WitnessChecker::RequirementEnvironmentCache &reqEnvCache,
+                    ProtocolDecl *proto, ProtocolConformance *conformance,
+                    DeclContext *dc, ValueDecl *req, ValueDecl *witness) {
   using namespace constraints;
 
   // Initialized by the setup operation.
@@ -719,17 +739,18 @@ RequirementMatch swift::matchWitness(TypeChecker &tc,
     }
   }
 
-  Optional<RequirementEnvironment> reqEnvironment(
-      RequirementEnvironment(dc, reqSig, proto, covariantSelf, conformance));
+  const RequirementEnvironment &reqEnvironment =
+      getOrCreateRequirementEnvironment(reqEnvCache, dc, reqSig, proto,
+                                        covariantSelf, conformance);
 
   // Set up the constraint system for matching.
   auto setup = [&]() -> std::tuple<Optional<RequirementMatch>, Type, Type> {
     // Construct a constraint system to use to solve the equality between
     // the required type and the witness type.
     cs.emplace(tc, dc, ConstraintSystemOptions());
 
-    auto reqGenericEnv = reqEnvironment->getSyntheticEnvironment();
-    auto reqSubMap = reqEnvironment->getRequirementToSyntheticMap();
+    auto reqGenericEnv = reqEnvironment.getSyntheticEnvironment();
+    auto reqSubMap = reqEnvironment.getRequirementToSyntheticMap();
 
     Type selfTy = proto->getSelfInterfaceType().subst(reqSubMap);
     if (reqGenericEnv)
@@ -820,7 +841,7 @@ RequirementMatch swift::matchWitness(TypeChecker &tc,
                               : anyRenaming ? MatchKind::RenamedMatch
                                             : MatchKind::ExactMatch,
                             witnessType,
-                            std::move(reqEnvironment),
+                            reqEnvironment,
                             optionalAdjustments);
 
     // Compute the set of substitutions we'll need for the witness.
@@ -1020,7 +1041,7 @@ bool WitnessChecker::findBestWitness(
         }
       }
 
-      auto match = matchWitness(TC, Proto, conformance, DC,
+      auto match = matchWitness(TC, ReqEnvironmentCache, Proto, conformance, DC,
                                 requirement, witness);
       if (match.isViable()) {
         ++numViable;
@@ -3038,7 +3059,8 @@ ResolveWitnessResult ConformanceChecker::resolveWitnessViaDerivation(
     return ResolveWitnessResult::ExplicitFailed;
 
   // Try to match the derived requirement.
-  auto match = matchWitness(TC, Proto, Conformance, DC, requirement, derived);
+  auto match = matchWitness(TC, ReqEnvironmentCache, Proto, Conformance, DC,
+                            requirement, derived);
   if (match.isViable()) {
     recordWitness(requirement, match);
     return ResolveWitnessResult::Success;
@@ -4445,9 +4467,10 @@ static void diagnosePotentialWitness(TypeChecker &tc,
               proto->getFullName());
 
   // Describe why the witness didn't satisfy the requirement.
+  WitnessChecker::RequirementEnvironmentCache oneUseCache;
   auto dc = conformance->getDeclContext();
-  auto match = matchWitness(tc, conformance->getProtocol(), conformance,
-                            dc, req, witness);
+  auto match = matchWitness(tc, oneUseCache, conformance->getProtocol(),
+                            conformance, dc, req, witness);
   if (match.Kind == MatchKind::ExactMatch &&
       req->isObjC() && !witness->isObjC()) {
     // Special case: note to add @objc.
@@ -4921,6 +4944,7 @@ swift::findWitnessedObjCRequirements(const ValueDecl *witness,
     }
   }
 
+  WitnessChecker::RequirementEnvironmentCache reqEnvCache;
   ASTContext &ctx = nominal->getASTContext();
   for (auto proto : nominal->getAllProtocols()) {
     // We only care about Objective-C protocols.
@@ -4976,7 +5000,7 @@ swift::findWitnessedObjCRequirements(const ValueDecl *witness,
         auto lazyResolver = ctx.getLazyResolver();
         assert(lazyResolver && "Need a type checker to match witnesses");
         auto &tc = *static_cast<TypeChecker *>(lazyResolver);
-        if (matchWitness(tc, proto, *conformance,
+        if (matchWitness(tc, reqEnvCache, proto, *conformance,
                          witnessToMatch->getDeclContext(), req,
                          const_cast<ValueDecl *>(witnessToMatch))
               .Kind == MatchKind::ExactMatch) {
@@ -5214,8 +5238,9 @@ void TypeChecker::recordKnownWitness(NormalProtocolConformance *conformance,
   // (because property behaviors rely on renaming).
   validateDecl(witness);
   auto dc = conformance->getDeclContext();
-  auto match = matchWitness(*this, conformance->getProtocol(), conformance,
-                            dc, req, witness);
+  WitnessChecker::RequirementEnvironmentCache oneUseCache;
+  auto match = matchWitness(*this, oneUseCache, conformance->getProtocol(),
+                            conformance, dc, req, witness);
   if (match.Kind != MatchKind::ExactMatch &&
       match.Kind != MatchKind::RenamedMatch) {
     diagnose(witness, diag::property_behavior_conformance_broken,

lib/Sema/TypeCheckProtocol.h

@@ -333,14 +333,14 @@ class OptionalAdjustment {
 /// \brief Describes a match between a requirement and a witness.
 struct RequirementMatch {
   RequirementMatch(ValueDecl *witness, MatchKind kind,
-                   Optional<RequirementEnvironment> &&env = None)
+                   Optional<RequirementEnvironment> env = None)
     : Witness(witness), Kind(kind), WitnessType(), ReqEnv(std::move(env)) {
     assert(!hasWitnessType() && "Should have witness type");
   }
 
   RequirementMatch(ValueDecl *witness, MatchKind kind,
                    Type witnessType,
-                   Optional<RequirementEnvironment> &&env = None,
+                   Optional<RequirementEnvironment> env = None,
                    ArrayRef<OptionalAdjustment> optionalAdjustments = {})
     : Witness(witness), Kind(kind), WitnessType(witnessType),
       ReqEnv(std::move(env)),
@@ -430,6 +430,12 @@ struct RequirementMatch {
 struct RequirementCheck;
 
 class WitnessChecker {
+public:
+  using RequirementEnvironmentCacheKey =
+      std::pair<const GenericSignature *, const ClassDecl *>;
+  using RequirementEnvironmentCache =
+      llvm::DenseMap<RequirementEnvironmentCacheKey, RequirementEnvironment>;
+
 protected:
   TypeChecker &TC;
   ProtocolDecl *Proto;
@@ -441,6 +447,8 @@ class WitnessChecker {
   // @_implements(Protocol, DeclName)
   llvm::DenseMap<DeclName, llvm::TinyPtrVector<ValueDecl *>> ImplementsTable;
 
+  RequirementEnvironmentCache ReqEnvironmentCache;
+
   WitnessChecker(TypeChecker &tc, ProtocolDecl *proto,
                  Type adoptee, DeclContext *dc);
 
@@ -839,12 +847,11 @@ RequirementMatch matchWitness(
                      RequirementMatch(bool, ArrayRef<OptionalAdjustment>)
                    > finalize);
 
-RequirementMatch matchWitness(TypeChecker &tc,
-                              ProtocolDecl *proto,
-                              ProtocolConformance *conformance,
-                              DeclContext *dc,
-                              ValueDecl *req,
-                              ValueDecl *witness);
+RequirementMatch
+  matchWitness(TypeChecker &tc,
+               WitnessChecker::RequirementEnvironmentCache &reqEnvCache,
+               ProtocolDecl *proto, ProtocolConformance *conformance,
+               DeclContext *dc, ValueDecl *req, ValueDecl *witness);
 
 /// If the given type is a direct reference to an associated type of
 /// the given protocol, return the referenced associated type.

test/SILGen/witnesses_class.swift

@@ -102,3 +102,17 @@ class UsesDefaults<X : Barable> : HasDefaults {}
 // CHECK: [[FN:%.*]] = function_ref @$S15witnesses_class11HasDefaultsPAAE23hasDefaultGenericTakesTyy1TQz_qd__tAA7FooableRd__lF : $@convention(method) <τ_0_0 where τ_0_0 : HasDefaults><τ_1_0 where τ_1_0 : Fooable> (@in_guaranteed τ_0_0.T, @guaranteed τ_1_0, @in_guaranteed τ_0_0) -> ()
 // CHECK: apply [[FN]]<UsesDefaults<τ_0_0>, τ_1_0>(
 // CHECK: return
+
+protocol ReturnsCovariantSelf {
+  func returnsCovariantSelf() -> Self
+}
+extension ReturnsCovariantSelf {
+  func returnsCovariantSelf() -> Self { return self }
+}
+
+class NonMatchingMember: ReturnsCovariantSelf {
+  func returnsCovariantSelf() {}
+}
+
+// CHECK-LABEL: sil private [transparent] [thunk] @$S15witnesses_class17NonMatchingMemberCAA20ReturnsCovariantSelfA2aDP07returnsgH0xyFTW : $@convention(witness_method: ReturnsCovariantSelf) <τ_0_0 where τ_0_0 : NonMatchingMember> (@in_guaranteed τ_0_0) -> @out τ_0_0 {
+