[IRGen] Track repeated requirements for async functions in NecessaryBindings.

Previously, NecessaryBindings used a SetVector to store the
GenericRequirements that it accumulates.  That was a problem for the
layout of async contexts where it is possible for the same generic
argument to be repeated.  Meanwhile, that is correct for partial
application forwarders where there is a single thunk for every partial
apply and consequently the bindings can be packed in without
duplication.

Here, a SetVector is used only when the NecessaryBindings are for
partial apply forwarders.  When the NecessaryBindings are instead for
async functions, a SmallVector is used.

Author: nate-chandler

lib/IRGen/GenProto.cpp

@@ -2688,21 +2688,21 @@ void irgen::emitPolymorphicParametersFromArray(IRGenFunction &IGF,
 
 Size NecessaryBindings::getBufferSize(IRGenModule &IGM) const {
   // We need one pointer for each archetype or witness table.
-  return IGM.getPointerSize() * Requirements.size();
+  return IGM.getPointerSize() * size();
 }
 
 void NecessaryBindings::restore(IRGenFunction &IGF, Address buffer,
                                 MetadataState metadataState) const {
-  bindFromGenericRequirementsBuffer(IGF, Requirements.getArrayRef(), buffer,
+  bindFromGenericRequirementsBuffer(IGF, getRequirements(), buffer,
                                     metadataState,
-                                    [&](CanType type) { return type;});
+                                    [&](CanType type) { return type; });
 }
 
 template <typename Transform>
 static void save(const NecessaryBindings &bindings, IRGenFunction &IGF,
                  Address buffer, Transform transform) {
   emitInitOfGenericRequirementsBuffer(
-      IGF, bindings.getRequirements().getArrayRef(), buffer,
+      IGF, bindings.getRequirements(), buffer,
       [&](GenericRequirement requirement) -> llvm::Value * {
         CanType type = requirement.TypeParameter;
         if (auto protocol = requirement.Protocol) {
@@ -2770,14 +2770,13 @@ void NecessaryBindings::addTypeMetadata(CanType type) {
   // Generic types are trickier, because they can require conformances.
 
   // Otherwise, just record the need for this metadata.
-  Requirements.insert({type, nullptr});
+  addRequirement({type, nullptr});
 }
 
 /// Add all the abstract conditional conformances in the specialized
 /// conformance to the \p requirements.
-static void addAbstractConditionalRequirements(
-    SpecializedProtocolConformance *specializedConformance,
-    llvm::SetVector<GenericRequirement> &requirements) {
+void NecessaryBindings::addAbstractConditionalRequirements(
+    SpecializedProtocolConformance *specializedConformance) {
   auto subMap = specializedConformance->getSubstitutionMap();
   auto condRequirements = specializedConformance->getConditionalRequirements();
   for (auto req : condRequirements) {
@@ -2789,7 +2788,7 @@ static void addAbstractConditionalRequirements(
     auto archetype = dyn_cast<ArchetypeType>(ty);
     if (!archetype)
       continue;
-    requirements.insert({ty, proto});
+    addRequirement({ty, proto});
   }
   // Recursively add conditional requirements.
   for (auto &conf : subMap.getConformances()) {
@@ -2799,7 +2798,7 @@ static void addAbstractConditionalRequirements(
         dyn_cast<SpecializedProtocolConformance>(conf.getConcrete());
     if (!specializedConf)
       continue;
-    addAbstractConditionalRequirements(specializedConf, requirements);
+    addAbstractConditionalRequirements(specializedConf);
   }
 }
 
@@ -2811,8 +2810,8 @@ void NecessaryBindings::addProtocolConformance(CanType type,
         dyn_cast<SpecializedProtocolConformance>(concreteConformance);
     // The partial apply forwarder does not have the context to reconstruct
     // abstract conditional conformance requirements.
-    if (forPartialApply && specializedConf) {
-      addAbstractConditionalRequirements(specializedConf, Requirements);
+    if (forPartialApply() && specializedConf) {
+      addAbstractConditionalRequirements(specializedConf);
     } else if (forAsyncFunction()) {
       ProtocolDecl *protocol = conf.getRequirement();
       GenericRequirement requirement;
@@ -2821,15 +2820,15 @@ void NecessaryBindings::addProtocolConformance(CanType type,
       std::pair<GenericRequirement, ProtocolConformanceRef> pair{requirement,
                                                                  conf};
       Conformances.insert(pair);
-      Requirements.insert({type, concreteConformance->getProtocol()});
+      addRequirement({type, concreteConformance->getProtocol()});
     }
     return;
   }
   assert(isa<ArchetypeType>(type) || forAsyncFunction());
 
   // TODO: pass something about the root conformance necessary to
   // reconstruct this.
-  Requirements.insert({type, conf.getAbstract()});
+  addRequirement({type, conf.getAbstract()});
 }
 
 llvm::Value *irgen::emitWitnessTableRef(IRGenFunction &IGF,
@@ -3033,7 +3032,8 @@ NecessaryBindings NecessaryBindings::computeBindings(
     bool forPartialApplyForwarder, bool considerParameterSources) {
 
   NecessaryBindings bindings;
-  bindings.forPartialApply = forPartialApplyForwarder;
+  bindings.kind =
+      forPartialApplyForwarder ? Kind::PartialApply : Kind::AsyncFunction;
 
   // Bail out early if we don't have polymorphic parameters.
   if (!hasPolymorphicParameters(origType))

lib/IRGen/NecessaryBindings.h

@@ -32,8 +32,9 @@ namespace swift {
   enum class MetadataState : size_t;
   class ProtocolDecl;
   class ProtocolConformanceRef;
+  class SpecializedProtocolConformance;
 
-namespace irgen {
+  namespace irgen {
   class Address;
   class IRGenFunction;
   class IRGenModule;
@@ -42,14 +43,32 @@ namespace irgen {
 /// NecessaryBindings - The set of metadata that must be saved in
 /// order to perform some set of operations on a type.
 class NecessaryBindings {
-  llvm::SetVector<GenericRequirement> Requirements;
+  enum class Kind {
+    /// Are the bindings to be computed for a partial apply forwarder.
+    /// In the case this is true we need to store/restore the conformance of a
+    /// specialized type with conditional conformance because the conditional
+    /// requirements are not available in the partial apply forwarder.
+    PartialApply,
+    AsyncFunction,
+  };
+  Kind kind;
+  llvm::SetVector<GenericRequirement> RequirementsSet;
+  llvm::SmallVector<GenericRequirement, 2> RequirementsVector;
   llvm::DenseMap<GenericRequirement, ProtocolConformanceRef> Conformances;
 
-  /// Are the bindings to be computed for a partial apply forwarder.
-  /// In the case this is true we need to store/restore the conformance of a
-  /// specialized type with conditional conformance because the conditional
-  /// requirements are not available in the partial apply forwarder.
-  bool forPartialApply = false;
+  void addRequirement(GenericRequirement requirement) {
+    switch (kind) {
+    case Kind::PartialApply:
+      RequirementsSet.insert(requirement);
+      break;
+    case Kind::AsyncFunction:
+      RequirementsVector.push_back(requirement);
+      break;
+    }
+  }
+
+  void addAbstractConditionalRequirements(
+      SpecializedProtocolConformance *specializedConformance);
 
 public:
   NecessaryBindings() = default;
@@ -70,24 +89,27 @@ class NecessaryBindings {
 
   /// Get the requirement from the bindings at index i.
   const GenericRequirement &operator[](size_t i) const {
-    return Requirements[i];
+    switch (kind) {
+    case Kind::PartialApply:
+      return RequirementsSet[i];
+    case Kind::AsyncFunction:
+      return RequirementsVector[i];
+    }
   }
 
   ProtocolConformanceRef
   getConformance(const GenericRequirement &requirement) const {
     return Conformances.lookup(requirement);
   }
 
-  size_t size() const {
-    return Requirements.size();
-  }
+  size_t size() const { return getRequirements().size(); }
 
   /// Add whatever information is necessary to reconstruct a witness table
   /// reference for the given type.
   void addProtocolConformance(CanType type, ProtocolConformanceRef conf);
 
   /// Is the work to do trivial?
-  bool empty() const { return Requirements.empty(); }
+  bool empty() const { return getRequirements().empty(); }
 
   /// Returns the required size of the bindings.
   /// Pointer alignment is sufficient.
@@ -101,11 +123,17 @@ class NecessaryBindings {
   /// Restore the necessary bindings from the given buffer.
   void restore(IRGenFunction &IGF, Address buffer, MetadataState state) const;
 
-  const llvm::SetVector<GenericRequirement> &getRequirements() const {
-    return Requirements;
+  const llvm::ArrayRef<GenericRequirement> getRequirements() const {
+    switch (kind) {
+    case Kind::PartialApply:
+      return RequirementsSet.getArrayRef();
+    case Kind::AsyncFunction:
+      return RequirementsVector;
+    }
   }
 
-  bool forAsyncFunction() { return !forPartialApply; }
+  bool forPartialApply() { return kind == Kind::PartialApply; }
+  bool forAsyncFunction() { return kind == Kind::AsyncFunction; }
 
 private:
   static NecessaryBindings computeBindings(IRGenModule &IGM,