Merge pull request #63733 from nate-chandler/variadic-generics/irgen/open_pack_element_requirements

nate-chandler · web-flow · commit 57ceb506df8b · 2023-02-16T20:07:35.000-08:00
[IRGen] Bound all requirements of open_pack_element.
diff --git a/include/swift/AST/GenericEnvironment.h b/include/swift/AST/GenericEnvironment.h
@@ -200,6 +200,9 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
   void forEachPackElementArchetype(
           llvm::function_ref<void(ElementArchetypeType*)> function) const;
 
+  void forEachPackElementGenericTypeParam(
+      llvm::function_ref<void(GenericTypeParamType *)> function) const;
+
   using PackElementBindingCallback =
     llvm::function_ref<void(ElementArchetypeType *elementType,
                             PackType *packSubstitution)>;
diff --git a/lib/AST/GenericEnvironment.cpp b/lib/AST/GenericEnvironment.cpp
@@ -141,13 +141,12 @@ void GenericEnvironment::forEachPackElementArchetype(
   }
 }
 
-void GenericEnvironment::forEachPackElementBinding(
-    PackElementBindingCallback function) const {
+void GenericEnvironment::forEachPackElementGenericTypeParam(
+    llvm::function_ref<void(GenericTypeParamType *)> function) const {
   auto sig = getGenericSignature();
   auto shapeClass = getOpenedElementShapeClass();
   auto packElements = sig.getInnermostGenericParams();
   auto packElementDepth = packElements.front()->getDepth();
-  auto elementIt = packElements.begin();
 
   // Each parameter pack in the outer generic parameters has
   // a corresponding pack element parameter at the innermost
@@ -164,13 +163,23 @@ void GenericEnvironment::forEachPackElementBinding(
     if (!sig->haveSameShape(genericParam, shapeClass->mapTypeOutOfContext()))
       continue;
 
+    function(genericParam);
+  }
+}
+
+void GenericEnvironment::forEachPackElementBinding(
+    PackElementBindingCallback function) const {
+  auto sig = getGenericSignature();
+  auto packElements = sig.getInnermostGenericParams();
+  auto elementIt = packElements.begin();
+  forEachPackElementGenericTypeParam([&](auto *genericParam) {
     assert(elementIt != packElements.end());
     auto *elementArchetype =
         mapTypeIntoContext(*elementIt++)->castTo<ElementArchetypeType>();
-    auto *packSubstitution =
-        maybeApplyOuterContextSubstitutions(genericParam)->castTo<PackType>();
+    auto *packSubstitution = maybeApplyOuterContextSubstitutions(genericParam)
+                                 ->template castTo<PackType>();
     function(elementArchetype, packSubstitution);
-  }
+  });
 
   assert(elementIt == packElements.end());
 }
diff --git a/lib/IRGen/GenPack.cpp b/lib/IRGen/GenPack.cpp
@@ -262,16 +262,13 @@ static llvm::Value *bindWitnessTableAtIndex(IRGenFunction &IGF,
 }
 
 struct OpenedElementContext {
-  GenericEnvironment *packEnvironment;
-  CanGenericSignature packSignature;
-
-  GenericEnvironment *elementEnvironment;
-  CanGenericSignature elementSignature;
+  GenericEnvironment *environment;
+  CanGenericSignature signature;
 
   static OpenedElementContext
   createForPackExpansion(IRGenFunction &IGF, CanPackExpansionType expansionTy) {
     // Get the outer generic signature and environment.
-    auto *genericEnv = cast<PackArchetypeType>(expansionTy.getCountType())
+    auto *genericEnv = cast<ArchetypeType>(expansionTy.getCountType())
                            ->getGenericEnvironment();
     auto subMap = genericEnv->getForwardingSubstitutionMap();
 
@@ -282,15 +279,15 @@ struct OpenedElementContext {
     auto *elementEnv = GenericEnvironment::forOpenedElement(
         elementSig, UUID::fromTime(), expansionTy.getCountType(), subMap);
 
-    return {genericEnv, genericSig, elementEnv, elementSig};
+    return {elementEnv, elementSig};
   }
 };
 
 static void bindElementSignatureRequirementsAtIndex(
     IRGenFunction &IGF, OpenedElementContext const &context, llvm::Value *index,
     DynamicMetadataRequest request) {
   enumerateGenericSignatureRequirements(
-      context.elementSignature, [&](GenericRequirement requirement) {
+      context.signature, [&](GenericRequirement requirement) {
         switch (requirement.getKind()) {
         case GenericRequirement::Kind::Shape:
         case GenericRequirement::Kind::Metadata:
@@ -299,12 +296,12 @@ static void bindElementSignatureRequirementsAtIndex(
         case GenericRequirement::Kind::MetadataPack: {
           auto ty = requirement.getTypeParameter();
           auto patternPackArchetype = cast<PackArchetypeType>(
-              context.packEnvironment->mapTypeIntoContext(ty)
+              context.environment->maybeApplyOuterContextSubstitutions(ty)
                   ->getCanonicalType());
           auto response =
               IGF.emitTypeMetadataRef(patternPackArchetype, request);
           auto elementArchetype =
-              context.elementEnvironment
+              context.environment
                   ->mapPackTypeIntoElementContext(
                       patternPackArchetype->getInterfaceType())
                   ->getCanonicalType();
@@ -319,20 +316,20 @@ static void bindElementSignatureRequirementsAtIndex(
           auto ty = requirement.getTypeParameter();
           auto proto = requirement.getProtocol();
           auto patternPackArchetype = cast<PackArchetypeType>(
-              context.packEnvironment->mapTypeIntoContext(ty)
+              context.environment->maybeApplyOuterContextSubstitutions(ty)
                   ->getCanonicalType());
           auto elementArchetype =
-              context.elementEnvironment
+              context.environment
                   ->mapPackTypeIntoElementContext(
                       patternPackArchetype->getInterfaceType())
                   ->getCanonicalType();
           llvm::Value *_metadata = nullptr;
           auto packConformance =
-              context.packSignature->lookupConformance(ty, proto);
+              context.signature->lookupConformance(ty, proto);
           auto *wtablePack = emitWitnessTableRef(IGF, patternPackArchetype,
                                                  &_metadata, packConformance);
           auto elementConformance =
-              context.elementSignature->lookupConformance(ty, proto);
+              context.signature->lookupConformance(ty, proto);
           auto *wtable = bindWitnessTableAtIndex(
               IGF, elementArchetype, elementConformance, wtablePack, index);
           assert(wtable);
@@ -350,7 +347,7 @@ static llvm::Value *emitPackExpansionElementMetadata(
 
   // Replace pack archetypes with element archetypes in the pattern type.
   auto instantiatedPatternTy =
-      context.elementEnvironment
+      context.environment
           ->mapPackTypeIntoElementContext(patternTy->mapTypeOutOfContext())
           ->getCanonicalType();
 
@@ -526,11 +523,11 @@ static llvm::Value *emitPackExpansionElementWitnessTable(
 
   // Replace pack archetypes with element archetypes in the pattern type.
   auto instantiatedPatternTy =
-      context.elementEnvironment
+      context.environment
           ->mapPackTypeIntoElementContext(patternTy->mapTypeOutOfContext())
           ->getCanonicalType();
   auto instantiatedConformance =
-      context.elementEnvironment->getGenericSignature()->lookupConformance(
+      context.environment->getGenericSignature()->lookupConformance(
           instantiatedPatternTy, conformance.getRequirement());
 
   // Emit the element witness table.
@@ -646,7 +643,7 @@ llvm::Value *irgen::emitTypeMetadataPackElementRef(
     ArrayRef<ProtocolDecl *> protocols, llvm::Value *index,
     DynamicMetadataRequest request,
     llvm::SmallVectorImpl<llvm::Value *> &wtables) {
-  // If the packs have already been materialized, just gep into it.
+  // If the packs have already been materialized, just gep into them.
   auto materializedMetadataPack =
       tryGetLocalPackTypeMetadata(IGF, packType, request);
   llvm::SmallVector<llvm::Value *> materializedWtablePacks;
@@ -692,9 +689,8 @@ llvm::Value *irgen::emitTypeMetadataPackElementRef(
     auto response = IGF.emitTypeMetadataRef(ty, request);
     auto *metadata = response.getMetadata();
     for (auto protocol : protocols) {
-      llvm::Value *_metadata = nullptr;
       auto *wtable =
-          emitWitnessTableRef(IGF, ty, /*srcMetadataCache=*/&_metadata,
+          emitWitnessTableRef(IGF, ty, /*srcMetadataCache=*/&metadata,
                               ProtocolConformanceRef(protocol));
       wtables.push_back(wtable);
     }
@@ -918,6 +914,71 @@ llvm::Value *irgen::emitTypeMetadataPackElementRef(
   return metadataPhi;
 }
 
+void irgen::bindOpenedElementArchetypesAtIndex(IRGenFunction &IGF,
+                                               GenericEnvironment *environment,
+                                               llvm::Value *index) {
+  assert(environment->getKind() == GenericEnvironment::Kind::OpenedElement);
+
+  // Record the generic type parameters of interest.
+  llvm::SmallPtrSet<CanType, 2> openablePackParams;
+  environment->forEachPackElementGenericTypeParam([&](auto *genericParam) {
+    openablePackParams.insert(genericParam->getCanonicalType());
+  });
+
+  // Find the archetypes and conformances which must be bound.
+  llvm::SmallSetVector<CanType, 2> types;
+  llvm::DenseMap<CanType, llvm::SmallVector<ProtocolDecl *, 2>>
+      protocolsForType;
+  auto isDerivedFromPackElementGenericTypeParam = [&](CanType ty) -> bool {
+    // Is this type itself an openable pack parameter OR a dependent type of
+    // one?
+    return openablePackParams.contains(
+        ty->getRootGenericParam()->getCanonicalType());
+  };
+  enumerateGenericSignatureRequirements(
+      environment->getGenericSignature().getCanonicalSignature(),
+      [&](GenericRequirement requirement) {
+        switch (requirement.getKind()) {
+        case GenericRequirement::Kind::MetadataPack: {
+          auto ty = requirement.getTypeParameter();
+          if (!isDerivedFromPackElementGenericTypeParam(ty))
+            return;
+          types.insert(ty);
+          protocolsForType.insert({ty, {}});
+          break;
+        }
+        case GenericRequirement::Kind::WitnessTablePack: {
+          auto ty = requirement.getTypeParameter();
+          if (!isDerivedFromPackElementGenericTypeParam(ty))
+            return;
+          types.insert(ty);
+          auto iterator = protocolsForType.insert({ty, {}}).first;
+          iterator->getSecond().push_back(requirement.getProtocol());
+          break;
+        }
+        case GenericRequirement::Kind::Shape:
+        case GenericRequirement::Kind::Metadata:
+        case GenericRequirement::Kind::WitnessTable:
+          break;
+        }
+      });
+
+  // For each archetype to be bound, find the corresponding conformances and
+  // bind the metadata and wtables.
+  for (auto ty : types) {
+    auto protocols = protocolsForType.find(ty)->getSecond();
+    auto archetype = cast<ElementArchetypeType>(
+        environment->mapPackTypeIntoElementContext(ty)->getCanonicalType());
+    auto pack =
+        cast<PackType>(environment->maybeApplyOuterContextSubstitutions(ty)
+                           ->getCanonicalType());
+    llvm::SmallVector<llvm::Value *, 2> wtables;
+    auto *metadata = emitTypeMetadataPackElementRef(
+        IGF, pack, protocols, index, MetadataState::Complete, wtables);
+    IGF.bindArchetype(archetype, metadata, MetadataState::Complete, wtables);
+  }
+}
+
 void irgen::cleanupTypeMetadataPack(IRGenFunction &IGF,
                                     StackAddress pack,
                                     Optional<unsigned> elementCount) {
diff --git a/lib/IRGen/GenPack.h b/lib/IRGen/GenPack.h
@@ -53,6 +53,10 @@ emitTypeMetadataPackRef(IRGenFunction &IGF,
                         CanPackType packType,
                         DynamicMetadataRequest request);
 
+void bindOpenedElementArchetypesAtIndex(IRGenFunction &IGF,
+                                        GenericEnvironment *env,
+                                        llvm::Value *index);
+
 llvm::Value *
 emitTypeMetadataPackElementRef(IRGenFunction &IGF, CanPackType packType,
                                ArrayRef<ProtocolDecl *> protocols,
diff --git a/lib/IRGen/IRGenSIL.cpp b/lib/IRGen/IRGenSIL.cpp
@@ -6902,17 +6902,8 @@ void IRGenSILFunction::visitScalarPackIndexInst(ScalarPackIndexInst *i) {
 void IRGenSILFunction::visitOpenPackElementInst(swift::OpenPackElementInst *i) {
   llvm::Value *index = getLoweredSingletonExplosion(i->getIndexOperand());
 
-  i->getOpenedGenericEnvironment()->forEachPackElementBinding(
-      [&](auto *archetype, auto *pack) {
-        auto protocols = archetype->getConformsTo();
-        llvm::SmallVector<llvm::Value *, 2> wtables;
-        auto *metadata = emitTypeMetadataPackElementRef(
-            *this, CanPackType(pack), protocols, index, MetadataState::Complete,
-            wtables);
-        bindArchetype(CanElementArchetypeType(archetype), metadata,
-                      MetadataState::Complete, wtables);
-      });
-
+  auto *env = i->getOpenedGenericEnvironment();
+  bindOpenedElementArchetypesAtIndex(*this, env, index);
   // The result is just used for type dependencies.
 }
 
diff --git a/test/IRGen/run_variadic_generics.sil b/test/IRGen/run_variadic_generics.sil
