Runtime: On-heap packs know their own length

Author: slavapestov

include/swift/ABI/Metadata.h

@@ -4944,6 +4944,14 @@ class TargetPackPointer {
   PackLifetime getLifetime() const {
     return (bool)(Ptr & 1) ? PackLifetime::OnHeap : PackLifetime::OnStack;
   }
+
+  // Get the number of elements in the pack, only valid for on-heap packs.
+  size_t getNumElements() const {
+    if (getLifetime() == PackLifetime::OnHeap)
+      return *(reinterpret_cast<const size_t *>(Ptr & ~1) - 1);
+
+    fatalError(0, "Cannot get length of on-stack pack");
+  }
 };
 
 /// A pointer to a metadata pack.

stdlib/public/runtime/Metadata.cpp

@@ -1275,7 +1275,9 @@ swift_allocateMetadataPack(const Metadata * const *ptr, size_t count) {
   PackCacheEntry<Metadata>::Key key{ptr, count};
   auto bytes = MetadataPacks.getOrInsert(key).first->getElements();
 
-  return MetadataPackPointer(bytes, PackLifetime::OnHeap).getPointer();
+  MetadataPackPointer pack(bytes, PackLifetime::OnHeap);
+  assert(pack.getNumElements() == count);
+  return pack.getPointer();
 }
 
 /// The uniquing structure for witness table packs.
@@ -1292,7 +1294,9 @@ swift_allocateWitnessTablePack(const WitnessTable * const *ptr, size_t count) {
   PackCacheEntry<WitnessTable>::Key key{ptr, count};
   auto bytes = WitnessTablePacks.getOrInsert(key).first->getElements();
 
-  return WitnessTablePackPointer(bytes, PackLifetime::OnHeap).getPointer();
+  WitnessTablePackPointer pack(bytes, PackLifetime::OnHeap);
+  assert(pack.getNumElements() == count);
+  return pack.getPointer();
 }
 
 /***************************************************************************/

stdlib/public/runtime/MetadataLookup.cpp

@@ -1200,6 +1200,14 @@ _gatherGenericParameters(const ContextDescriptor *context,
     auto generics = context->getGenericContext();
     assert(generics);
     {
+      // Add a placeholder length for each shape class.
+      auto packShapeHeader = generics->getGenericPackShapeHeader();
+      if (packShapeHeader.NumShapeClasses > 0) {
+        assert(allGenericArgsVec.empty());
+        allGenericArgsVec.resize(packShapeHeader.NumShapeClasses);
+      }
+
+      // If we have the wrong number of generic arguments, fail.
       auto genericParams = generics->getGenericParams();
       unsigned n = genericParams.size();
       if (allGenericArgs.size() != n) {
@@ -1210,6 +1218,8 @@ _gatherGenericParameters(const ContextDescriptor *context,
                  "generic args, expected " + std::to_string(n);
         });
       }
+
+      // Add metadata for each canonical generic parameter.
       for (unsigned i = 0; i != n; ++i) {
         const auto &param = genericParams[i];
         if (param.getKind() != GenericParamKind::Type) {
@@ -1223,9 +1233,23 @@ _gatherGenericParameters(const ContextDescriptor *context,
         if (param.hasKeyArgument())
           allGenericArgsVec.push_back(allGenericArgs[i]);
       }
-    }
 
-    // If we have the wrong number of generic arguments, fail.
+      // Fill in the length for each shape class.
+      auto packShapeDescriptors = generics->getGenericPackShapeDescriptors();
+      for (auto packShapeDescriptor : packShapeDescriptors) {
+        if (packShapeDescriptor.Kind != GenericPackKind::Metadata)
+          continue;
+
+        assert(packShapeDescriptor.Index < allGenericArgsVec.size());
+        assert(packShapeDescriptor.ShapeClass < packShapeHeader.NumShapeClasses);
+
+        MetadataPackPointer pack(allGenericArgsVec[packShapeDescriptor.Index]);
+        assert(pack.getLifetime() == PackLifetime::OnHeap);
+
+        allGenericArgsVec[packShapeDescriptor.ShapeClass] =
+            reinterpret_cast<const void *>(pack.getNumElements());
+      }
+    }
 
     // Check whether the generic requirements are satisfied, collecting
     // any extra arguments we need for the instantiation function.