Merge pull request #61699 from aschwaighofer/fix_gen_distributed_opaque_ptr

IRGen: Fix the computation of the number of type metadata pointers for distributed thunks under opaque pointer IR

Author: aschwaighofer

lib/IRGen/GenDistributed.cpp

@@ -681,14 +681,17 @@ void DistributedAccessor::emit() {
     // We need this to determine the expected number of witness tables
     // to load from the buffer provided by the caller.
     llvm::SmallVector<llvm::Type *, 4> targetGenericArguments;
-    expandPolymorphicSignature(IGM, targetTy, targetGenericArguments);
+    auto numDirectGenericArgs =
+        expandPolymorphicSignature(IGM, targetTy, targetGenericArguments);
 
     // Generic arguments associated with the distributed thunk directly
     // e.g. `distributed func echo<T, U>(...)`
-    auto numDirectGenericArgs =
-        llvm::count_if(targetGenericArguments, [&](const llvm::Type *type) {
-          return type == IGM.TypeMetadataPtrTy;
-        });
+    assert(
+        !IGM.getLLVMContext().supportsTypedPointers() ||
+        numDirectGenericArgs ==
+            llvm::count_if(targetGenericArguments, [&](const llvm::Type *type) {
+              return type == IGM.TypeMetadataPtrTy;
+            }));
 
     auto expectedWitnessTables =
         targetGenericArguments.size() - numDirectGenericArgs;

lib/IRGen/GenProto.cpp

@@ -3456,12 +3456,15 @@ void irgen::bindGenericRequirement(IRGenFunction &IGF,
 namespace {
   /// A class for expanding a polymorphic signature.
   class ExpandPolymorphicSignature : public PolymorphicConvention {
+    unsigned numTypeMetadataPtrs = 0;
+
   public:
     ExpandPolymorphicSignature(IRGenModule &IGM, CanSILFunctionType fn)
       : PolymorphicConvention(IGM, fn) {}
 
-    void expand(SmallVectorImpl<llvm::Type *> &out,
-                SmallVectorImpl<PolymorphicSignatureExpandedTypeSource> *reqs) {
+    unsigned
+    expand(SmallVectorImpl<llvm::Type *> &out,
+           SmallVectorImpl<PolymorphicSignatureExpandedTypeSource> *reqs) {
       auto outStartSize = out.size();
       (void)outStartSize;
       for (auto &source : getSources())
@@ -3472,9 +3475,13 @@ namespace {
           reqs->push_back(reqt);
         out.push_back(reqt.Protocol ? IGM.WitnessTablePtrTy
                                     : IGM.TypeMetadataPtrTy);
+
+        if (!reqt.Protocol)
+          ++numTypeMetadataPtrs;
       });
       assert((!reqs || reqs->size() == (out.size() - outStartSize)) &&
              "missing type source for type");
+      return numTypeMetadataPtrs;
     }
 
   private:
@@ -3488,6 +3495,7 @@ namespace {
       case MetadataSource::Kind::GenericLValueMetadata:
         if (reqs)
           reqs->push_back(source);
+        ++numTypeMetadataPtrs;
         return out.push_back(IGM.TypeMetadataPtrTy);
       case MetadataSource::Kind::SelfMetadata:
       case MetadataSource::Kind::SelfWitnessTable:
@@ -3501,11 +3509,11 @@ namespace {
 } // end anonymous namespace
 
 /// Given a generic signature, add the argument types required in order to call it.
-void irgen::expandPolymorphicSignature(
+unsigned irgen::expandPolymorphicSignature(
     IRGenModule &IGM, CanSILFunctionType polyFn,
     SmallVectorImpl<llvm::Type *> &out,
     SmallVectorImpl<PolymorphicSignatureExpandedTypeSource> *outReqs) {
-  ExpandPolymorphicSignature(IGM, polyFn).expand(out, outReqs);
+  return ExpandPolymorphicSignature(IGM, polyFn).expand(out, outReqs);
 }
 
 void irgen::expandTrailingWitnessSignature(IRGenModule &IGM,

lib/IRGen/GenProto.h

@@ -101,7 +101,8 @@ namespace irgen {
 
   /// Add the witness parameters necessary for calling a function with
   /// the given generics clause.
-  void expandPolymorphicSignature(
+  /// Returns the number of type metadata pointers added to `types`.
+  unsigned expandPolymorphicSignature(
       IRGenModule &IGM, CanSILFunctionType type,
       SmallVectorImpl<llvm::Type *> &types,
       SmallVectorImpl<PolymorphicSignatureExpandedTypeSource> *outReqs =