IRGen: Fulfillment of pack shapes from nominal type metadata

slavapestov · rjmccall · commit 7d779a9b3621 · 2023-03-10T12:15:55.000-05:00
diff --git a/include/swift/IRGen/GenericRequirement.h b/include/swift/IRGen/GenericRequirement.h
@@ -55,10 +55,10 @@ class GenericRequirement {
   CanType type;
   ProtocolDecl *proto;
 
+public:
   GenericRequirement(Kind kind, CanType type, ProtocolDecl *proto)
     : kind(kind), type(type), proto(proto) {}
 
-public:
   Kind getKind() const {
     return kind;
   }
@@ -76,35 +76,34 @@ class GenericRequirement {
   }
 
   static GenericRequirement forShape(CanType type) {
-    assert(type->isParameterPack());
+    assert(!isa<PackType>(type));
+    assert(type->isParameterPack() || isa<PackArchetypeType>(type));
     return GenericRequirement(Kind::Shape, type, nullptr);
   }
 
   bool isMetadata() const {
     return kind == Kind::Metadata || kind == Kind::MetadataPack;
   }
 
-  static GenericRequirement forMetadata(CanType type, bool isPack) {
-    auto kind = isPack ? Kind::MetadataPack : Kind::Metadata;
-    return GenericRequirement(kind, type, nullptr);
-  }
-
   static GenericRequirement forMetadata(CanType type) {
-    return forMetadata(type, type->hasParameterPack());
+    assert(!isa<PackType>(type));
+    auto kind = ((type->isParameterPack() ||
+                  isa<PackArchetypeType>(type))
+                 ? Kind::MetadataPack : Kind::Metadata);
+    return GenericRequirement(kind, type, nullptr);
   }
 
   bool isWitnessTable() const {
     return kind == Kind::WitnessTable || kind == Kind::WitnessTablePack;
   }
 
-  static GenericRequirement forWitnessTable(CanType type, ProtocolDecl *proto,
-                                            bool isPack) {
-    auto kind = isPack ? Kind::WitnessTablePack : Kind::WitnessTable;
-    return GenericRequirement(kind, type, proto);
-  }
-
   static GenericRequirement forWitnessTable(CanType type, ProtocolDecl *proto) {
-    return forWitnessTable(type, proto, type->hasParameterPack());
+    assert(!isa<PackType>(type));
+    auto kind = ((type->isParameterPack() ||
+                  isa<PackArchetypeType>(type))
+                 ? Kind::WitnessTablePack
+                 : Kind::WitnessTable);
+    return GenericRequirement(kind, type, proto);
   }
 
   static llvm::Type *typeForKind(irgen::IRGenModule &IGM,
@@ -142,12 +141,14 @@ template <> struct DenseMapInfo<swift::GenericRequirement> {
   using GenericRequirement = swift::GenericRequirement;
   using CanTypeInfo = llvm::DenseMapInfo<swift::CanType>;
   static GenericRequirement getEmptyKey() {
-    return GenericRequirement::forMetadata(CanTypeInfo::getEmptyKey(),
-                                           /*isPack=*/false);
+    return GenericRequirement(GenericRequirement::Kind::Metadata,
+                              CanTypeInfo::getEmptyKey(),
+                              nullptr);
   }
   static GenericRequirement getTombstoneKey() {
-    return GenericRequirement::forMetadata(CanTypeInfo::getTombstoneKey(),
-                                           /*isPack=*/false);
+    return GenericRequirement(GenericRequirement::Kind::Metadata,
+                              CanTypeInfo::getTombstoneKey(),
+                              nullptr);
   }
   static llvm::hash_code getHashValue(GenericRequirement req) {
     return hash_combine(CanTypeInfo::getHashValue(req.getTypeParameter()),
diff --git a/lib/IRGen/Fulfillment.cpp b/lib/IRGen/Fulfillment.cpp
@@ -291,30 +291,44 @@ bool FulfillmentMap::searchNominalTypeMetadata(IRGenModule &IGM,
     if (!keys.hasInterestingType(arg))
       continue;
 
-    // If the fulfilled value is type metadata, refine the path.
-    if (requirement.isMetadata()) {
+    switch (requirement.getKind()) {
+    case GenericRequirement::Kind::Shape: {
+      // If the fulfilled value is a shape class, refine the path.
+      MetadataPath argPath = path;
+      argPath.addNominalTypeArgumentShapeComponent(reqtIndex);
+
+      // Add the fulfillment.
+      hadFulfillment |= addFulfillment(GenericRequirement::forShape(arg),
+                                       source, std::move(argPath),
+                                       MetadataState::Complete);
+      break;
+    }
+    case GenericRequirement::Kind::Metadata:
+    case GenericRequirement::Kind::MetadataPack: {
+      // If the fulfilled value is type metadata, refine the path.
       auto argState =
           getPresumedMetadataStateForTypeArgument(metadataState);
       MetadataPath argPath = path;
       argPath.addNominalTypeArgumentComponent(reqtIndex);
       hadFulfillment |= searchTypeMetadata(
           IGM, arg, IsExact, argState, source, std::move(argPath), keys);
-      continue;
+      break;
     }
+    case GenericRequirement::Kind::WitnessTable:
+    case GenericRequirement::Kind::WitnessTablePack: {
+      // Ignore it unless the type itself is interesting.
+      if (!keys.isInterestingType(arg))
+        continue;
 
-    // Otherwise, it's a conformance.
-    assert(requirement.isWitnessTable());
-
-    // Ignore it unless the type itself is interesting.
-    if (!keys.isInterestingType(arg))
-      continue;
-
-    // Refine the path.
-    MetadataPath argPath = path;
-    argPath.addNominalTypeArgumentConformanceComponent(reqtIndex);
+      // Refine the path.
+      MetadataPath argPath = path;
+      argPath.addNominalTypeArgumentConformanceComponent(reqtIndex);
 
-    hadFulfillment |= searchWitnessTable(IGM, arg, requirement.getProtocol(),
-                                         source, std::move(argPath), keys);
+      hadFulfillment |= searchWitnessTable(IGM, arg, requirement.getProtocol(),
+                                           source, std::move(argPath), keys);
+      break;
+    }
+    }
   }
 
   return hadFulfillment;
diff --git a/lib/IRGen/GenMeta.h b/lib/IRGen/GenMeta.h
@@ -130,6 +130,15 @@ namespace irgen {
                                            unsigned reqtIndex,
                                            llvm::Value *metadata);
 
+  /// Given a reference to nominal type metadata of the given type,
+  /// derive a reference to a the pack shape stored in the nth
+  /// requirement slot.  The type must have generic arguments.
+  llvm::Value *emitArgumentPackShapeRef(IRGenFunction &IGF,
+                                        NominalTypeDecl *theDecl,
+                                        const GenericTypeRequirements &reqts,
+                                        unsigned reqtIndex,
+                                        llvm::Value *metadata);
+
   /// Given a metatype value, read its instance type.
   llvm::Value *emitMetatypeInstanceType(IRGenFunction &IGF,
                                         llvm::Value *metatypeMetadata);
diff --git a/lib/IRGen/GenProto.cpp b/lib/IRGen/GenProto.cpp
@@ -2782,7 +2782,8 @@ MetadataResponse MetadataPath::followComponent(IRGenFunction &IGF,
                                                DynamicMetadataRequest request) {
   switch (component.getKind()) {
   case Component::Kind::NominalTypeArgument:
-  case Component::Kind::NominalTypeArgumentConformance: {
+  case Component::Kind::NominalTypeArgumentConformance:
+  case Component::Kind::NominalTypeArgumentShape: {
     assert(sourceKey.Kind == LocalTypeDataKind::forFormalTypeMetadata());
     auto type = sourceKey.Type;
     if (auto archetypeTy = dyn_cast<ArchetypeType>(type))
@@ -2820,9 +2821,24 @@ MetadataResponse MetadataPath::followComponent(IRGenFunction &IGF,
       // Do a dynamic check if necessary to satisfy the request.
       return emitCheckTypeMetadataState(IGF, request, response);
 
+    // If this is a shape class, load the value.
+    } else if (component.getKind() == Component::Kind::NominalTypeArgumentShape) {
+      assert(requirement.isShape() && "index mismatch!");
+
+      sourceKey.Kind = LocalTypeDataKind::forPackShapeExpression();
+
+      if (!source) return MetadataResponse();
+
+      auto sourceMetadata = source.getMetadata();
+      auto shape = emitArgumentPackShapeRef(IGF, nominal,
+                                            requirements, reqtIndex,
+                                            sourceMetadata);
+
+      return MetadataResponse::forComplete(shape);
+
     // Otherwise, we need to switch sourceKey.Kind to the appropriate
     // conformance kind.
-    } else {
+    } else if (component.getKind() == Component::Kind::NominalTypeArgumentConformance) {
       assert(requirement.isWitnessTable() && "index mismatch!");
       auto conformance = subs.lookupConformance(requirement.getTypeParameter(),
                                                 requirement.getProtocol());
@@ -2840,6 +2856,8 @@ MetadataResponse MetadataPath::followComponent(IRGenFunction &IGF,
 
       return MetadataResponse::forComplete(wtable);
     }
+
+    llvm_unreachable("Bad component kind");
   }
 
   case Component::Kind::OutOfLineBaseProtocol: {
@@ -3062,6 +3080,10 @@ void MetadataPath::print(llvm::raw_ostream &out) const {
       out << "nominal_type_argument_conformance["
           << component.getPrimaryIndex() << "]";
       break;
+    case Component::Kind::NominalTypeArgumentShape:
+      out << "nominal_type_argument_shape["
+          << component.getPrimaryIndex() << "]";
+      break;
     case Component::Kind::ConditionalConformance:
       out << "conditional_conformance[" << component.getPrimaryIndex() << "]";
       break;
diff --git a/lib/IRGen/LocalTypeData.cpp b/lib/IRGen/LocalTypeData.cpp
@@ -545,22 +545,13 @@ addAbstractForFulfillments(IRGenFunction &IGF, FulfillmentMap &&fulfillments,
     CanType type = fulfillment.first.getTypeParameter();
     LocalTypeDataKind localDataKind;
 
-    // For now, ignore witness-table fulfillments when they're not for
-    // archetypes.
-    if (fulfillment.first.isWitnessTable()) {
-      ProtocolDecl *protocol = fulfillment.first.getProtocol();
-      if (auto archetype = dyn_cast<ArchetypeType>(type)) {
-        auto conformsTo = archetype->getConformsTo();
-        auto it = std::find(conformsTo.begin(), conformsTo.end(), protocol);
-        if (it == conformsTo.end()) continue;
-        localDataKind = LocalTypeDataKind::forAbstractProtocolWitnessTable(*it);
-      } else {
-        continue;
-      }
-
-    } else {
-      assert(fulfillment.first.isMetadata());
-
+    switch (fulfillment.first.getKind()) {
+    case GenericRequirement::Kind::Shape: {
+      localDataKind = LocalTypeDataKind::forPackShapeExpression();
+      break;
+    }
+    case GenericRequirement::Kind::Metadata:
+    case GenericRequirement::Kind::MetadataPack: {
       // Ignore type metadata fulfillments for non-dependent types that
       // we can produce very cheaply.  We don't want to end up emitting
       // the type metadata for Int by chasing through N layers of metadata
@@ -571,6 +562,24 @@ addAbstractForFulfillments(IRGenFunction &IGF, FulfillmentMap &&fulfillments,
       }
 
       localDataKind = LocalTypeDataKind::forFormalTypeMetadata();
+      break;
+    }
+    case GenericRequirement::Kind::WitnessTable:
+    case GenericRequirement::Kind::WitnessTablePack: {
+      // For now, ignore witness-table fulfillments when they're not for
+      // archetypes.
+      ProtocolDecl *protocol = fulfillment.first.getProtocol();
+      if (auto archetype = dyn_cast<ArchetypeType>(type)) {
+        auto conformsTo = archetype->getConformsTo();
+        auto it = std::find(conformsTo.begin(), conformsTo.end(), protocol);
+        if (it == conformsTo.end()) continue;
+        localDataKind = LocalTypeDataKind::forAbstractProtocolWitnessTable(*it);
+      } else {
+        continue;
+      }
+
+      break;
+    }
     }
 
     // Find the chain for the key.
diff --git a/lib/IRGen/MetadataLayout.cpp b/lib/IRGen/MetadataLayout.cpp
@@ -235,6 +235,19 @@ llvm::Value *irgen::emitArgumentWitnessTableRef(IRGenFunction &IGF,
                                       IGF.IGM.WitnessTablePtrTy);
 }
 
+/// Given a reference to nominal type metadata of the given type,
+/// derive a reference to the pack shape for the nth argument
+/// metadata.  The type must have generic arguments.
+llvm::Value *irgen::emitArgumentPackShapeRef(IRGenFunction &IGF,
+                                             NominalTypeDecl *decl,
+                                       const GenericTypeRequirements &reqts,
+                                             unsigned reqtIndex,
+                                             llvm::Value *metadata) {
+  assert(reqts.getRequirements()[reqtIndex].isShape());
+  return emitLoadOfGenericRequirement(IGF, metadata, decl, reqtIndex,
+                                      IGF.IGM.SizeTy);
+}
+
 Address irgen::emitAddressOfFieldOffsetVector(IRGenFunction &IGF,
                                               llvm::Value *metadata,
                                               NominalTypeDecl *decl) {
diff --git a/lib/IRGen/MetadataPath.h b/lib/IRGen/MetadataPath.h
@@ -57,6 +57,9 @@ class MetadataPath {
       /// Witness table at requirement index P of a generic nominal type.
       NominalTypeArgumentConformance,
 
+      /// Pack length at requirement index P of a generic nominal type.
+      NominalTypeArgumentShape,
+
       /// Type metadata at requirement index P of a generic nominal type.
       NominalTypeArgument,
 
@@ -104,6 +107,7 @@ class MetadataPath {
       switch (getKind()) {
       case Kind::OutOfLineBaseProtocol:
       case Kind::NominalTypeArgumentConformance:
+      case Kind::NominalTypeArgumentShape:
       case Kind::NominalTypeArgument:
       case Kind::ConditionalConformance:
         return OperationCost::Load;
@@ -160,6 +164,13 @@ class MetadataPath {
                              index));
   }
 
+  /// Add a step to this path which gets the pack length stored at
+  /// requirement index n in a generic type metadata.
+  void addNominalTypeArgumentShapeComponent(unsigned index) {
+    Path.push_back(Component(Component::Kind::NominalTypeArgumentShape,
+                             index));
+  }
+
   /// Add a step to this path which gets the inherited protocol at
   /// a particular witness index.
   void addInheritedProtocolComponent(WitnessIndex index) {
