IRGen: Basic support for computing function signatures with shape parameters

Author: slavapestov

lib/IRGen/GenDistributed.cpp

@@ -680,22 +680,21 @@ 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;
-    auto numDirectGenericArgs =
+    auto expandedSignature =
         expandPolymorphicSignature(IGM, targetTy, targetGenericArguments);
+    assert(expandedSignature.numShapes == 0 &&
+           "Distributed actors don't support variadic generics");
 
     // Generic arguments associated with the distributed thunk directly
     // e.g. `distributed func echo<T, U>(...)`
     assert(
         !IGM.getLLVMContext().supportsTypedPointers() ||
-        numDirectGenericArgs ==
+        expandedSignature.numTypeMetadataPtrs ==
             llvm::count_if(targetGenericArguments, [&](const llvm::Type *type) {
               return type == IGM.TypeMetadataPtrTy;
             }));
 
-    auto expectedWitnessTables =
-        targetGenericArguments.size() - numDirectGenericArgs;
-
-    for (unsigned index = 0; index < numDirectGenericArgs; ++index) {
+    for (unsigned index = 0; index < expandedSignature.numTypeMetadataPtrs; ++index) {
       auto offset =
           Size(index * IGM.DataLayout.getTypeAllocSize(IGM.TypeMetadataPtrTy));
       auto alignment =
@@ -708,7 +707,7 @@ void DistributedAccessor::emit() {
     }
 
     emitLoadOfWitnessTables(witnessTables, numWitnessTables,
-                            expectedWitnessTables, arguments);
+                            expandedSignature.numWitnessTablePtrs, arguments);
   }
 
   // Step two, let's form and emit a call to the distributed method

lib/IRGen/GenProto.cpp

@@ -3426,31 +3426,46 @@ void irgen::bindGenericRequirement(IRGenFunction &IGF,
   // Get the corresponding context type.
   auto type = getInContext(requirement.getTypeParameter());
 
-  if (requirement.isWitnessTable()) {
+  switch (requirement.getKind()) {
+  case GenericRequirement::Kind::Shape: {
+    assert(isa<ArchetypeType>(type));
+    assert(value->getType() == IGF.IGM.SizeTy);
+    auto kind = LocalTypeDataKind::forPackShapeExpression();
+    IGF.setUnscopedLocalTypeData(type, kind, value);
+    break;
+  }
+
+  case GenericRequirement::Kind::Metadata: {
+    assert(value->getType() == IGF.IGM.TypeMetadataPtrTy);
+    setTypeMetadataName(IGF.IGM, value, type);
+    IGF.bindLocalTypeDataFromTypeMetadata(type, IsExact, value, metadataState);
+    break;
+  }
+
+  case GenericRequirement::Kind::WitnessTable: {
     auto proto = requirement.getProtocol();
     assert(isa<ArchetypeType>(type));
     assert(value->getType() == IGF.IGM.WitnessTablePtrTy);
     setProtocolWitnessTableName(IGF.IGM, value, type, proto);
     auto kind = LocalTypeDataKind::forAbstractProtocolWitnessTable(proto);
     IGF.setUnscopedLocalTypeData(type, kind, value);
-  } else {
-    assert(requirement.isMetadata());
-    assert(value->getType() == IGF.IGM.TypeMetadataPtrTy);
-    setTypeMetadataName(IGF.IGM, value, type);
-    IGF.bindLocalTypeDataFromTypeMetadata(type, IsExact, value, metadataState);
+    break;
+  }
   }
 }
 
 namespace {
   /// A class for expanding a polymorphic signature.
   class ExpandPolymorphicSignature : public PolymorphicConvention {
+    unsigned numShapes = 0;
     unsigned numTypeMetadataPtrs = 0;
+    unsigned numWitnessTablePtrs = 0;
 
   public:
     ExpandPolymorphicSignature(IRGenModule &IGM, CanSILFunctionType fn)
       : PolymorphicConvention(IGM, fn) {}
 
-    unsigned
+    ExpandedSignature
     expand(SmallVectorImpl<llvm::Type *> &out,
            SmallVectorImpl<PolymorphicSignatureExpandedTypeSource> *reqs) {
       auto outStartSize = out.size();
@@ -3461,17 +3476,24 @@ namespace {
       enumerateUnfulfilledRequirements([&](GenericRequirement reqt) {
         if (reqs)
           reqs->push_back(reqt);
-        out.push_back(reqt.isWitnessTable() ? IGM.WitnessTablePtrTy
-                                            : IGM.TypeMetadataPtrTy);
-
-        if (reqt.isMetadata())
+        switch (reqt.getKind()) {
+        case GenericRequirement::Kind::Shape:
+          out.push_back(IGM.SizeTy);
+          ++numShapes;
+          break;
+        case GenericRequirement::Kind::Metadata:
+          out.push_back(IGM.TypeMetadataPtrTy);
           ++numTypeMetadataPtrs;
-        else
-          assert(reqt.isWitnessTable());
+          break;
+        case GenericRequirement::Kind::WitnessTable:
+          out.push_back(IGM.WitnessTablePtrTy);
+          ++numWitnessTablePtrs;
+          break;
+        }
       });
       assert((!reqs || reqs->size() == (out.size() - outStartSize)) &&
              "missing type source for type");
-      return numTypeMetadataPtrs;
+      return {numShapes, numTypeMetadataPtrs, numWitnessTablePtrs};
     }
 
   private:
@@ -3499,7 +3521,7 @@ namespace {
 } // end anonymous namespace
 
 /// Given a generic signature, add the argument types required in order to call it.
-unsigned irgen::expandPolymorphicSignature(
+ExpandedSignature irgen::expandPolymorphicSignature(
     IRGenModule &IGM, CanSILFunctionType polyFn,
     SmallVectorImpl<llvm::Type *> &out,
     SmallVectorImpl<PolymorphicSignatureExpandedTypeSource> *outReqs) {

lib/IRGen/GenProto.h

@@ -99,10 +99,16 @@ namespace irgen {
     return -1 - (int)index;
   }
 
+  struct ExpandedSignature {
+    unsigned numShapes;
+    unsigned numTypeMetadataPtrs;
+    unsigned numWitnessTablePtrs;
+  };
+
   /// Add the witness parameters necessary for calling a function with
   /// the given generics clause.
-  /// Returns the number of type metadata pointers added to `types`.
-  unsigned expandPolymorphicSignature(
+  /// Returns the number of lowered parameters of each kind.
+  ExpandedSignature expandPolymorphicSignature(
       IRGenModule &IGM, CanSILFunctionType type,
       SmallVectorImpl<llvm::Type *> &types,
       SmallVectorImpl<PolymorphicSignatureExpandedTypeSource> *outReqs =

lib/IRGen/LocalTypeDataKind.h

@@ -49,6 +49,7 @@ class LocalTypeDataKind {
     FormalTypeMetadata,
     RepresentationTypeMetadata,
     ValueWitnessTable,
+    Shape,
     // <- add more special cases here
 
     // The first enumerator for an individual value witness.
@@ -95,6 +96,11 @@ class LocalTypeDataKind {
     return LocalTypeDataKind(ValueWitnessDiscriminatorBase + (unsigned)witness);
   }
 
+  /// A reference to the shape expression of a pack type.
+  static LocalTypeDataKind forPackShapeExpression() {
+    return LocalTypeDataKind(Shape);
+  }
+
   /// A reference to a protocol witness table for an archetype.
   ///
   /// This only works for non-concrete types because in principle we might