AST: Introduce GenericSignature::forEachParam()

This replaces the inefficient pattern:

  for (auto param : sig->getGenericParams()) {
    if (sig->isCanonicalTypeInContext(param)) {
      ...
    } else {
      ...
    }
  }

Author: slavapestov

include/swift/AST/GenericSignature.h

@@ -190,9 +190,10 @@ class alignas(1 << TypeAlignInBits) GenericSignature final
   Optional<ProtocolConformanceRef>
   lookupConformance(CanType depTy, ProtocolDecl *proto) const;
 
-  /// Return a vector of all generic parameters that are not subject to
-  /// a concrete same-type constraint.
-  SmallVector<GenericTypeParamType *, 2> getSubstitutableParams() const;
+  /// Iterate over all generic parameters, passing a flag to the callback
+  /// indicating if the generic parameter is canonical or not.
+  void forEachParam(
+    llvm::function_ref<void(GenericTypeParamType *, bool)> callback) const;
 
   /// Check if the generic signature makes all generic parameters
   /// concrete.

include/swift/SIL/TypeSubstCloner.h

@@ -297,12 +297,15 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
     if (SubsMap.empty())
       return false;
 
-    for (auto ParamType : Sig->getSubstitutableParams()) {
+    bool Result = false;
+    Sig->forEachParam([&](GenericTypeParamType *ParamType, bool Canonical) {
+      if (!Canonical)
+        return;
       if (!Type(ParamType).subst(SubsMap)->isEqual(ParamType))
-        return true;
-    }
+        Result = true;
+    });
 
-    return false;
+    return Result;
   }
 
   enum { ForInlining = true };

lib/AST/ASTContext.cpp

@@ -4384,14 +4384,15 @@ void SubstitutionMap::Storage::Profile(
   // Profile those replacement types that corresponding to canonical generic
   // parameters within the generic signature.
   id.AddInteger(replacementTypes.size());
-  auto genericParams = genericSig->getGenericParams();
-  for (unsigned i : indices(genericParams)) {
-    auto gp = genericParams[i];
-    if (genericSig->isCanonicalTypeInContext(gp->getCanonicalType()))
+
+  unsigned i = 0;
+  genericSig->forEachParam([&](GenericTypeParamType *gp, bool canonical) {
+    if (canonical)
       id.AddPointer(replacementTypes[i].getPointer());
     else
       id.AddPointer(nullptr);
-  }
+    i++;
+  });
 
   // Conformances.
   id.AddInteger(conformances.size());

lib/AST/GenericSignature.cpp

@@ -97,45 +97,42 @@ GenericSignature::getInnermostGenericParams() const {
   return params;
 }
 
-
-SmallVector<GenericTypeParamType *, 2>
-GenericSignature::getSubstitutableParams() const {
+void GenericSignature::forEachParam(
+    llvm::function_ref<void(GenericTypeParamType *, bool)> callback) const {
   // Figure out which generic parameters are concrete or same-typed to another
-  // generic parameter.
+  // type parameter.
   auto genericParams = getGenericParams();
-  auto genericParamsAreNotSubstitutable =
-    SmallVector<bool, 4>(genericParams.size(), false);
+  auto genericParamsAreCanonical =
+    SmallVector<bool, 4>(genericParams.size(), true);
+
   for (auto req : getRequirements()) {
     if (req.getKind() != RequirementKind::SameType) continue;
 
     GenericTypeParamType *gp;
     if (auto secondGP = req.getSecondType()->getAs<GenericTypeParamType>()) {
-      // If two generic parameters are same-typed, then the left-hand one
-      // is canonical.
+      // If two generic parameters are same-typed, then the right-hand one
+      // is non-canonical.
+      assert(req.getFirstType()->is<GenericTypeParamType>());
       gp = secondGP;
     } else {
-      // If an associated type is same-typed, it doesn't constrain the generic
-      // parameter itself.
-      if (req.getSecondType()->isTypeParameter()) continue;
-
-      // Otherwise, the generic parameter is concrete.
+      // Otherwise, the right-hand side is an associated type or concrete type,
+      // and the left-hand one is non-canonical.
       gp = req.getFirstType()->getAs<GenericTypeParamType>();
       if (!gp) continue;
+
+      // If an associated type is same-typed, it doesn't constrain the generic
+      // parameter itself. That is, if T == U.Foo, then T is canonical, whereas
+      // U.Foo is not.
+      if (req.getSecondType()->isTypeParameter()) continue;
     }
 
     unsigned index = GenericParamKey(gp).findIndexIn(genericParams);
-    genericParamsAreNotSubstitutable[index] = true;
+    genericParamsAreCanonical[index] = false;
   }
 
-  // Collect the generic parameters that are substitutable.
-  SmallVector<GenericTypeParamType *, 2> result;
-  for (auto index : indices(genericParams)) {
-    auto gp = genericParams[index];
-    if (!genericParamsAreNotSubstitutable[index])
-      result.push_back(gp);
-  }
-
-  return result;
+  // Call the callback with each parameter and the result of the above analysis.
+  for (auto index : indices(genericParams))
+    callback(genericParams[index], genericParamsAreCanonical[index]);
 }
 
 bool GenericSignature::areAllParamsConcrete() const {

lib/AST/SubstitutionMap.cpp

@@ -200,18 +200,19 @@ SubstitutionMap SubstitutionMap::get(GenericSignature *genericSig,
   // Form the replacement types.
   SmallVector<Type, 4> replacementTypes;
   replacementTypes.reserve(genericSig->getGenericParams().size());
-  for (auto gp : genericSig->getGenericParams()) {
+
+  genericSig->forEachParam([&](GenericTypeParamType *gp, bool canonical) {
     // Don't eagerly form replacements for non-canonical generic parameters.
-    if (!genericSig->isCanonicalTypeInContext(gp->getCanonicalType())) {
+    if (!canonical) {
       replacementTypes.push_back(Type());
-      continue;
+      return;
     }
 
     // Record the replacement.
     Type replacement = Type(gp).subst(subs, lookupConformance,
                                       SubstFlags::UseErrorType);
     replacementTypes.push_back(replacement);
-  }
+  });
 
   // Form the stored conformances.
   SmallVector<ProtocolConformanceRef, 4> conformances;

lib/IRGen/GenMeta.cpp

@@ -313,15 +313,15 @@ namespace {
       GenericSignature *sig = asImpl().getGenericSignature();
       assert(sig);
       auto canSig = sig->getCanonicalSignature();
-      
-      for (auto param : canSig->getGenericParams()) {
+
+      canSig->forEachParam([&](GenericTypeParamType *param, bool canonical) {
         // Currently, there are only type parameters. The parameter is a key
         // argument if it's canonical in its generic context.
         asImpl().addGenericParameter(GenericParamKind::Type,
-                 /*key argument*/ canSig->isCanonicalTypeInContext(param),
-                 /*extra argument*/ false);
-      }
-      
+                                     /*key argument*/ canonical,
+                                     /*extra argument*/ false);
+      });
+
       // Pad the structure up to four bytes for the following requirements.
       unsigned padding = (unsigned) -canSig->getGenericParams().size() & 3;
       for (unsigned i = 0; i < padding; ++i)

lib/IRGen/GenProto.cpp

@@ -238,9 +238,10 @@ irgen::enumerateGenericSignatureRequirements(CanGenericSignature signature,
   if (!signature) return;
 
   // Get all of the type metadata.
-  for (auto gp : signature->getSubstitutableParams()) {
-    callback({CanType(gp), nullptr});
-  }
+  signature->forEachParam([&](GenericTypeParamType *gp, bool canonical) {
+    if (canonical)
+      callback({CanType(gp), nullptr});
+  });
 
   // Get the protocol conformances.
   for (auto &reqt : signature->getRequirements()) {

lib/RemoteAST/RemoteAST.cpp

@@ -178,7 +178,12 @@ class RemoteASTTypeBuilder {
 
     // Build a SubstitutionMap.
     auto *genericSig = decl->getGenericSignature();
-    auto genericParams = genericSig->getSubstitutableParams();
+
+    SmallVector<GenericTypeParamType *, 4> genericParams;
+    genericSig->forEachParam([&](GenericTypeParamType *gp, bool canonical) {
+      if (canonical)
+        genericParams.push_back(gp);
+    });
     if (genericParams.size() != args.size())
       return Type();
 

lib/SILOptimizer/IPO/EagerSpecializer.cpp

@@ -347,7 +347,11 @@ void EagerDispatch::emitDispatchTo(SILFunction *NewFunc) {
   auto GenericSig =
     GenericFunc->getLoweredFunctionType()->getGenericSignature();
   auto SubMap = ReInfo.getClonerParamSubstitutionMap();
-  for (auto ParamTy : GenericSig->getSubstitutableParams()) {
+
+  GenericSig->forEachParam([&](GenericTypeParamType *ParamTy, bool Canonical) {
+    if (!Canonical)
+      return;
+
     auto Replacement = Type(ParamTy).subst(SubMap);
     assert(!Replacement->hasTypeParameter());
 
@@ -368,7 +372,8 @@ void EagerDispatch::emitDispatchTo(SILFunction *NewFunc) {
                                   Replacement, LayoutInfo);
       }
     }
-  }
+  });
+
   static_cast<void>(FailedTypeCheckBB);
 
   if (OldReturnBB == &EntryBB) {

lib/SILOptimizer/Utils/Generics.cpp

@@ -434,7 +434,11 @@ bool ReabstractionInfo::prepareAndCheck(ApplySite Apply, SILFunction *Callee,
   bool HasConcreteGenericParams = false;
   bool HasNonArchetypeGenericParams = false;
   HasUnboundGenericParams = false;
-  for (auto GP : CalleeGenericSig->getSubstitutableParams()) {
+
+  CalleeGenericSig->forEachParam([&](GenericTypeParamType *GP, bool Canonical) {
+    if (!Canonical)
+      return;
+
     // Check only the substitutions for the generic parameters.
     // Ignore any dependent types, etc.
     auto Replacement = Type(GP).subst(CalleeParamSubMap);
@@ -458,11 +462,10 @@ bool ReabstractionInfo::prepareAndCheck(ApplySite Apply, SILFunction *Callee,
             HasNonArchetypeGenericParams = true;
         }
       }
-      continue;
+    } else {
+      HasConcreteGenericParams = true;
     }
-
-    HasConcreteGenericParams = true;
-  }
+  });
 
   if (HasUnboundGenericParams) {
     // Bail if we cannot specialize generic substitutions, but all substitutions
@@ -1561,14 +1564,16 @@ void FunctionSignaturePartialSpecializer::
   // Simply create a set of same-type requirements based on concrete
   // substitutions.
   SmallVector<Requirement, 4> Requirements;
-  for (auto GP : CalleeGenericSig->getSubstitutableParams()) {
+  CalleeGenericSig->forEachParam([&](GenericTypeParamType *GP, bool Canonical) {
+    if (!Canonical)
+      return;
     auto Replacement = Type(GP).subst(CalleeInterfaceToCallerArchetypeMap);
     if (Replacement->hasArchetype())
-      continue;
+      return;
     // Replacement is concrete. Add a same type requirement.
     Requirement Req(RequirementKind::SameType, GP, Replacement);
     Requirements.push_back(Req);
-  }
+  });
 
   // Create a new generic signature by taking the existing one
   // and adding new requirements to it. No need to introduce