[IRGen] Forwarded wtable packs when possible.

Previously wtable packs were never forwarded because pack types
consisting of a single element which was itself a pack archetype were
never canonicalized into that.

In the forwarding scenario, gepped and loaded from wtable when lowering
open_pack_element.

Author: nate-chandler

lib/IRGen/GenPack.cpp

@@ -456,8 +456,6 @@ static llvm::Value *emitPackExpansionElementWitnessTable(
           ->mapPackTypeIntoElementContext(patternTy->mapTypeOutOfContext())
           ->getCanonicalType();
 
-  // FIXME: Handle witness table packs for associatedtype's conformances.
-
   // Emit the element witness table.
   auto *wtable = emitWitnessTableRef(IGF, instantiatedPatternTy,
                                      /*srcMetadataCache=*/nullptr, conformance);
@@ -567,18 +565,38 @@ void irgen::cleanupWitnessTablePack(IRGenFunction &IGF, StackAddress pack,
   }
 }
 
+static llvm::Value *tryGetLocalPackTypeData(IRGenFunction &IGF,
+                                            CanPackType packType,
+                                            LocalTypeDataKind localDataKind) {
+  if (auto *wtable = IGF.tryGetLocalTypeData(packType, localDataKind))
+    return wtable;
+
+  if (auto packArchetypeType = getForwardedPackArchetypeType(packType)) {
+    if (auto *wtable =
+            IGF.tryGetLocalTypeData(packArchetypeType, localDataKind))
+      return wtable;
+  }
+
+  return nullptr;
+}
+
 llvm::Value *irgen::emitWitnessTablePackRef(IRGenFunction &IGF,
                                             CanPackType packType,
                                             PackConformance *conformance) {
   assert(Lowering::TypeConverter::protocolRequiresWitnessTable(
              conformance->getProtocol()) &&
          "looking up witness table for protocol that doesn't have one");
 
+  if (auto *wtable = tryGetLocalPackTypeData(
+          IGF, packType,
+          LocalTypeDataKind::forAbstractProtocolWitnessTable(
+              conformance->getProtocol())))
+    return wtable;
+
   auto localDataKind =
       LocalTypeDataKind::forProtocolWitnessTablePack(conformance);
 
-  auto wtable = IGF.tryGetLocalTypeData(packType, localDataKind);
-  if (wtable)
+  if (auto *wtable = tryGetLocalPackTypeData(IGF, packType, localDataKind))
     return wtable;
 
   auto pack = emitWitnessTablePack(IGF, packType, conformance);
@@ -600,20 +618,26 @@ llvm::Value *irgen::emitTypeMetadataPackElementRef(
       tryGetLocalPackTypeMetadata(IGF, packType, request);
   llvm::SmallVector<llvm::Value *> materializedWtablePacks;
   for (auto protocol : protocols) {
-    auto wtable = IGF.tryGetLocalTypeData(
-        packType, LocalTypeDataKind::forAbstractProtocolWitnessTable(protocol));
-    materializedWtablePacks.push_back(wtable);
+    auto *wtablePack = tryGetLocalPackTypeData(
+        IGF, packType,
+        LocalTypeDataKind::forAbstractProtocolWitnessTable(protocol));
+    materializedWtablePacks.push_back(wtablePack);
   }
   if (materializedMetadataPack &&
       llvm::all_of(materializedWtablePacks,
-                   [](auto *wtable) { return wtable; })) {
+                   [](auto *wtablePack) { return wtablePack; })) {
     auto *gep = IGF.Builder.CreateInBoundsGEP(
         IGF.IGM.TypeMetadataPtrTy, materializedMetadataPack.getMetadata(),
         index);
     auto addr =
         Address(gep, IGF.IGM.TypeMetadataPtrTy, IGF.IGM.getPointerAlignment());
     auto *metadata = IGF.Builder.CreateLoad(addr);
-    for (auto *wtable : materializedWtablePacks) {
+    for (auto *wtablePack : materializedWtablePacks) {
+      auto *gep = IGF.Builder.CreateInBoundsGEP(IGF.IGM.WitnessTablePtrTy,
+                                                wtablePack, index);
+      auto addr = Address(gep, IGF.IGM.WitnessTablePtrTy,
+                          IGF.IGM.getPointerAlignment());
+      auto *wtable = IGF.Builder.CreateLoad(addr);
       wtables.push_back(wtable);
     }
     return metadata;

test/IRGen/run_variadic_generics.sil

@@ -344,6 +344,17 @@ entry(%intIndex : $Builtin.Word):
   return %t : $()
 }
 
+// Verify that we just gep into parameter packs for metadata and witness tables when that's all that the pack consists of.
+// CHECK-LL-LABEL: define {{.*}}@direct_access_from_parameter_with_conformance(
+// CHECK-LL-SAME:        i{{(32|64)}} [[INDEX:%[^,]+]], 
+// CHECK-LL-SAME:        i{{(32|64)}} {{%[^,]+}},
+// CHECK-LL-SAME:        %swift.type** [[METADATA_PACK:%[^,]+]],
+// CHECK-LL-SAME:        i8*** [[WTABLE_PACK:%[^,]+]])
+// CHECK-LL:         [[METADATA_ADDRESS:%[^,]+]] = getelementptr inbounds %swift.type*, %swift.type** [[METADATA_PACK]], i{{(32|64)}} [[INDEX]]
+// CHECK-LL:         [[METADATA:%[^,]+]] = load %swift.type*, %swift.type** [[METADATA_ADDRESS]]
+// CHECK-LL:         [[WTABLE_ADDRESS:%[^,]+]] = getelementptr inbounds i8**, i8*** [[WTABLE_PACK]], i{{(32|64)}} [[INDEX]]
+// CHECK-LL:         [[WTABLE:%[^,]+]] = load i8**, i8*** [[WTABLE_ADDRESS]], align 8
+// CHECK-LL:         call swiftcc void @printGenericType(%swift.type* [[METADATA]], %swift.type* [[METADATA]])
 sil @direct_access_from_parameter_with_conformance : $<T_1... : P> (Builtin.Word) -> () {
 entry(%intIndex : $Builtin.Word):
   %innerIndex = dynamic_pack_index %intIndex of $Pack{repeat each T_1}

test/IRGen/variadic_generic_functions.sil

@@ -32,16 +32,36 @@ sil @c : $() -> () {
   return %ret : $()
 }
 
-// CHECK: define {{.*}}void @f(i{{(64|32)}} %0, %swift.type** %T, i8*** %T.P)
+// CHECK: define {{.*}}void @g(i{{(64|32)}} %0, %swift.type** %T, i8*** %T.P)
+sil @g : $<T... : P> () -> () {
+    %f = function_ref @f : $@convention(thin) <T... : P> () -> ()
+    apply %f<Pack{repeat each T, S, repeat each T, repeat each T}>() : $@convention(thin) <T... : P> () -> ()
+    %ret = tuple ()
+    return %ret : $()
+}
+
+// Verify that polymorphic parameters are just forwarded along.
+// CHECK-LABEL: define {{.*}}void @f(i{{(32|64)}} %0, %swift.type** %T, i8*** %T.P)
+// CHECK:         call swiftcc void @fc(i{{(32|64)}} %0, %swift.type** %T, i8*** %T.P)
 sil @f : $<T... : P> () -> () {
+    %fc = function_ref @fc : $@convention(thin) <T... : P> () -> ()
+    apply %fc<Pack{repeat each T}>() : $@convention(thin) <T... : P> () -> ()
     %ret = tuple ()
     return %ret : $()
 }
+sil @fc : $<T... : P> () -> () {}
 
-// CHECK: define {{.*}}void @g(i{{(64|32)}} %0, %swift.type** %T, i8*** %T.P)
-sil @g : $<T... : P> () -> () {
-    %f = function_ref @f : $@convention(thin) <T... : P> () -> ()
-    apply %f<Pack{repeat each T, S, repeat each T, repeat each T}>() : $@convention(thin) <T... : P> () -> ()
+protocol PA {
+  associatedtype A
+}
+
+// CHECK-LABEL: define {{.*}}@f1(i{{(32|64)}} %0, %swift.type** %T, i8*** %T.PA, i8*** %T.A.P)
+// CHECK:         call swiftcc void @f1c(i{{(32|64)}} %0, %swift.type** %T, i8*** %T.PA, i8*** %T.A.P)
+sil @f1 : $<T... : PA where each T.A : P> () -> () {
+    %f1c = function_ref @f1c : $@convention(thin) <T... : PA where each T.A : P> () -> ()
+    apply %f1c<Pack{repeat each T}>() : $@convention(thin) <T... : PA where each T.A : P> () -> ()
     %ret = tuple ()
     return %ret : $()
 }
+
+sil @f1c : $<T... : PA where each T.A : P> () -> () {}