IRGen: Refactor function type metadata emission

Author: slavapestov

lib/IRGen/MetadataRequest.cpp

@@ -1370,28 +1370,19 @@ static llvm::Value *getFunctionParameterRef(IRGenFunction &IGF,
   return IGF.emitAbstractTypeMetadataRef(type);
 }
 
-static MetadataResponse emitFunctionTypeMetadataRef(IRGenFunction &IGF,
-                                                    CanFunctionType type,
-                                                    DynamicMetadataRequest request) {
-  auto result =
-    IGF.emitAbstractTypeMetadataRef(type->getResult()->getCanonicalType());
-
-  auto params = type.getParams();
-  auto numParams = params.size();
-
-  // Retrieve the ABI parameter flags from the type-level parameter
-  // flags.
-  auto getABIParameterFlags = [](ParameterTypeFlags flags) {
-    return ParameterFlags()
+/// Mapping type-level parameter flags to ABI parameter flags.
+static ParameterFlags getABIParameterFlags(ParameterTypeFlags flags) {
+  return ParameterFlags()
         .withValueOwnership(flags.getValueOwnership())
         .withVariadic(flags.isVariadic())
         .withAutoClosure(flags.isAutoClosure())
         .withNoDerivative(flags.isNoDerivative())
         .withIsolated(flags.isIsolated());
-  };
+}
 
+static FunctionTypeFlags getFunctionTypeFlags(CanFunctionType type) {
   bool hasParameterFlags = false;
-  for (auto param : params) {
+  for (auto param : type.getParams()) {
     if (!getABIParameterFlags(param.getParameterFlags()).isNone()) {
       hasParameterFlags = true;
       break;
@@ -1417,80 +1408,149 @@ static MetadataResponse emitFunctionTypeMetadataRef(IRGenFunction &IGF,
     break;
   }
 
-  FunctionMetadataDifferentiabilityKind metadataDifferentiabilityKind;
-  switch (type->getDifferentiabilityKind()) {
-  case DifferentiabilityKind::NonDifferentiable:
-    metadataDifferentiabilityKind =
-        FunctionMetadataDifferentiabilityKind::NonDifferentiable;
-    break;
-  case DifferentiabilityKind::Normal:
-    metadataDifferentiabilityKind =
-        FunctionMetadataDifferentiabilityKind::Normal;
-    break;
-  case DifferentiabilityKind::Linear:
-    metadataDifferentiabilityKind =
-        FunctionMetadataDifferentiabilityKind::Linear;
-    break;
-  case DifferentiabilityKind::Forward:
-    metadataDifferentiabilityKind =
-        FunctionMetadataDifferentiabilityKind::Forward;
-    break;
-  case DifferentiabilityKind::Reverse:
-    metadataDifferentiabilityKind =
-        FunctionMetadataDifferentiabilityKind::Reverse;
-    break;
+  return FunctionTypeFlags()
+      .withConvention(metadataConvention)
+      .withAsync(type->isAsync())
+      .withConcurrent(type->isSendable())
+      .withThrows(type->isThrowing())
+      .withParameterFlags(hasParameterFlags)
+      .withEscaping(isEscaping)
+      .withDifferentiable(type->isDifferentiable())
+      .withGlobalActor(!type->getGlobalActor().isNull());
+}
+
+namespace {
+struct FunctionTypeMetadataParamInfo {
+  StackAddress parameters;
+  StackAddress paramFlags;
+  unsigned numParams;
+};
+}
+
+static FunctionTypeMetadataParamInfo
+emitFunctionTypeMetadataParams(IRGenFunction &IGF,
+                               AnyFunctionType::CanParamArrayRef params,
+                               FunctionTypeFlags flags,
+                               DynamicMetadataRequest request,
+                               SmallVectorImpl<llvm::Value *> &arguments) {
+  FunctionTypeMetadataParamInfo info;
+  info.numParams = params.size();
+
+  ConstantInitBuilder paramFlags(IGF.IGM);
+  auto flagsArr = paramFlags.beginArray();
+
+  if (!params.empty()) {
+    auto arrayTy =
+        llvm::ArrayType::get(IGF.IGM.TypeMetadataPtrTy, info.numParams);
+    info.parameters = StackAddress(IGF.createAlloca(
+        arrayTy, IGF.IGM.getTypeMetadataAlignment(), "function-parameters"));
+
+    IGF.Builder.CreateLifetimeStart(info.parameters.getAddress(),
+                                    IGF.IGM.getPointerSize() * info.numParams);
+
+    for (unsigned i : indices(params)) {
+      auto param = params[i];
+      auto paramFlags = getABIParameterFlags(param.getParameterFlags());
+
+      auto argPtr = IGF.Builder.CreateStructGEP(info.parameters.getAddress(), i,
+                                                IGF.IGM.getPointerSize());
+      auto *typeRef = getFunctionParameterRef(IGF, param);
+      IGF.Builder.CreateStore(typeRef, argPtr);
+      if (i == 0)
+        arguments.push_back(argPtr.getAddress());
+
+      flagsArr.addInt32(paramFlags.getIntValue());
+    }
+  } else {
+    auto parametersPtr =
+        llvm::ConstantPointerNull::get(
+          IGF.IGM.TypeMetadataPtrTy->getPointerTo());
+    arguments.push_back(parametersPtr);
+  }
+
+  auto *Int32Ptr = IGF.IGM.Int32Ty->getPointerTo();
+  if (flags.hasParameterFlags()) {
+    auto *flagsVar = flagsArr.finishAndCreateGlobal(
+        "parameter-flags", IGF.IGM.getPointerAlignment(),
+        /* constant */ true);
+    arguments.push_back(IGF.Builder.CreateBitCast(flagsVar, Int32Ptr));
+  } else {
+    flagsArr.abandon();
+    arguments.push_back(llvm::ConstantPointerNull::get(Int32Ptr));
   }
 
-  auto flags = FunctionTypeFlags()
-                   .withNumParameters(numParams)
-                   .withConvention(metadataConvention)
-                   .withAsync(type->isAsync())
-                   .withConcurrent(type->isSendable())
-                   .withThrows(type->isThrowing())
-                   .withParameterFlags(hasParameterFlags)
-                   .withEscaping(isEscaping)
-                   .withDifferentiable(type->isDifferentiable())
-                   .withGlobalActor(!type->getGlobalActor().isNull());
-
-  auto flagsVal = llvm::ConstantInt::get(IGF.IGM.SizeTy,
-                                         flags.getIntValue());
-  llvm::Value *diffKindVal = nullptr;
-  if (type->isDifferentiable()) {
-    assert(metadataDifferentiabilityKind.isDifferentiable());
-    diffKindVal = llvm::ConstantInt::get(
-        IGF.IGM.SizeTy, metadataDifferentiabilityKind.getIntValue());
-  } else if (type->getGlobalActor()) {
-    diffKindVal = llvm::ConstantInt::get(
-        IGF.IGM.SizeTy,
-        FunctionMetadataDifferentiabilityKind::NonDifferentiable);
-  }
-
-  auto collectParameters =
-      [&](llvm::function_ref<void(unsigned, llvm::Value *,
-                                  ParameterFlags flags)>
-              processor) {
-        for (auto index : indices(params)) {
-          auto param = params[index];
-          auto flags = param.getParameterFlags();
-
-          auto parameterFlags = getABIParameterFlags(flags);
-          processor(index, getFunctionParameterRef(IGF, param),
-                    parameterFlags);
-        }
-      };
+  return info;
+}
+
+static FunctionTypeMetadataParamInfo
+emitDynamicFunctionTypeMetadataParams(IRGenFunction &IGF,
+                                      AnyFunctionType::CanParamArrayRef params,
+                                      FunctionTypeFlags flags,
+                                      CanPackType packType,
+                                      DynamicMetadataRequest request,
+                                      SmallVectorImpl<llvm::Value *> &arguments) {
+  assert(false);
+}
+
+static void cleanupFunctionTypeMetadataParams(IRGenFunction &IGF,
+                                              FunctionTypeMetadataParamInfo info) {
+  if (info.parameters.isValid()) {
+    if (info.parameters.getExtraInfo()) {
+      IGF.emitDeallocateDynamicAlloca(info.parameters);
+    } else {
+      IGF.Builder.CreateLifetimeEnd(info.parameters.getAddress(),
+                                    IGF.IGM.getPointerSize() * info.numParams);
+    }
+  }
+}
+
+static CanPackType getInducedPackType(AnyFunctionType::CanParamArrayRef params,
+                                      ASTContext &ctx) {
+  SmallVector<CanType, 2> elts;
+  for (auto param : params)
+    elts.push_back(param.getPlainType());
+  
+  return CanPackType::get(ctx, elts);
+}
+
+static MetadataResponse emitFunctionTypeMetadataRef(IRGenFunction &IGF,
+                                                    CanFunctionType type,
+                                                    DynamicMetadataRequest request) {
+  auto result =
+    IGF.emitAbstractTypeMetadataRef(type->getResult()->getCanonicalType());
+
+  auto params = type.getParams();
+  bool hasPackExpansion = type->containsPackExpansionParam();
+
+  auto flags = getFunctionTypeFlags(type);
+  llvm::Value *flagsVal = nullptr;
+  llvm::Value *shapeExpression = nullptr;
+  CanPackType packType;
+
+  if (!hasPackExpansion) {
+    flags = flags.withNumParameters(params.size());
+    flagsVal = llvm::ConstantInt::get(IGF.IGM.SizeTy,
+                                      flags.getIntValue());
+  } else {
+    packType = getInducedPackType(type.getParams(), type->getASTContext());
+    auto *shapeExpression = IGF.emitPackShapeExpression(packType);
+
+    flagsVal = llvm::ConstantInt::get(IGF.IGM.SizeTy,
+                                      flags.getIntValue());
+    flagsVal = IGF.Builder.CreateOr(flagsVal, shapeExpression);
+  }
 
   auto constructSimpleCall =
       [&](llvm::SmallVectorImpl<llvm::Value *> &arguments)
       -> FunctionPointer {
+    assert(!flags.hasParameterFlags());
+    assert(!shapeExpression);
+
     arguments.push_back(flagsVal);
 
-    collectParameters([&](unsigned i, llvm::Value *typeRef,
-                          ParameterFlags flags) {
-      arguments.push_back(typeRef);
-      if (hasParameterFlags)
-        arguments.push_back(
-            llvm::ConstantInt::get(IGF.IGM.Int32Ty, flags.getIntValue()));
-    });
+    for (auto param : params) {
+      arguments.push_back(getFunctionParameterRef(IGF, param));
+    }
 
     arguments.push_back(result);
 
@@ -1512,13 +1572,13 @@ static MetadataResponse emitFunctionTypeMetadataRef(IRGenFunction &IGF,
     }
   };
 
-  switch (numParams) {
+  switch (params.size()) {
   case 0:
   case 1:
   case 2:
   case 3: {
-    if (!hasParameterFlags && !type->isDifferentiable() &&
-        !type->getGlobalActor()) {
+    if (!flags.hasParameterFlags() && !type->isDifferentiable() &&
+        !type->getGlobalActor() && !hasPackExpansion) {
       llvm::SmallVector<llvm::Value *, 8> arguments;
       auto metadataFn = constructSimpleCall(arguments);
       auto *call = IGF.Builder.CreateCall(metadataFn, arguments);
@@ -1537,54 +1597,60 @@ static MetadataResponse emitFunctionTypeMetadataRef(IRGenFunction &IGF,
            "0 parameter case should be specialized unless it is a "
            "differentiable function or has a global actor");
 
-    auto *const Int32Ptr = IGF.IGM.Int32Ty->getPointerTo();
     llvm::SmallVector<llvm::Value *, 8> arguments;
 
     arguments.push_back(flagsVal);
 
-    if (diffKindVal) {
-      arguments.push_back(diffKindVal);
-    }
-
-    ConstantInitBuilder paramFlags(IGF.IGM);
-    auto flagsArr = paramFlags.beginArray();
+    llvm::Value *diffKindVal = nullptr;
 
-    Address parameters;
-    if (!params.empty()) {
-      auto arrayTy =
-          llvm::ArrayType::get(IGF.IGM.TypeMetadataPtrTy, numParams);
-      parameters = IGF.createAlloca(
-          arrayTy, IGF.IGM.getTypeMetadataAlignment(), "function-parameters");
-
-      IGF.Builder.CreateLifetimeStart(parameters,
-                                      IGF.IGM.getPointerSize() * numParams);
+    {
+      FunctionMetadataDifferentiabilityKind metadataDifferentiabilityKind;
+      switch (type->getDifferentiabilityKind()) {
+      case DifferentiabilityKind::NonDifferentiable:
+        metadataDifferentiabilityKind =
+            FunctionMetadataDifferentiabilityKind::NonDifferentiable;
+        break;
+      case DifferentiabilityKind::Normal:
+        metadataDifferentiabilityKind =
+            FunctionMetadataDifferentiabilityKind::Normal;
+        break;
+      case DifferentiabilityKind::Linear:
+        metadataDifferentiabilityKind =
+            FunctionMetadataDifferentiabilityKind::Linear;
+        break;
+      case DifferentiabilityKind::Forward:
+        metadataDifferentiabilityKind =
+            FunctionMetadataDifferentiabilityKind::Forward;
+        break;
+      case DifferentiabilityKind::Reverse:
+        metadataDifferentiabilityKind =
+            FunctionMetadataDifferentiabilityKind::Reverse;
+        break;
+      }
 
-      collectParameters([&](unsigned i, llvm::Value *typeRef,
-                            ParameterFlags flags) {
-        auto argPtr = IGF.Builder.CreateStructGEP(parameters, i,
-                                                  IGF.IGM.getPointerSize());
-        IGF.Builder.CreateStore(typeRef, argPtr);
-        if (i == 0)
-          arguments.push_back(argPtr.getAddress());
+      if (type->isDifferentiable()) {
+        assert(metadataDifferentiabilityKind.isDifferentiable());
+        diffKindVal = llvm::ConstantInt::get(
+            IGF.IGM.SizeTy, metadataDifferentiabilityKind.getIntValue());
+      } else if (type->getGlobalActor()) {
+        diffKindVal = llvm::ConstantInt::get(
+            IGF.IGM.SizeTy,
+            FunctionMetadataDifferentiabilityKind::NonDifferentiable);
+      }
+    }
 
-        if (hasParameterFlags)
-          flagsArr.addInt32(flags.getIntValue());
-      });
-    } else {
-      auto parametersPtr =
-          llvm::ConstantPointerNull::get(
-            IGF.IGM.TypeMetadataPtrTy->getPointerTo());
-      arguments.push_back(parametersPtr);
+    if (diffKindVal) {
+      arguments.push_back(diffKindVal);
     }
 
-    if (hasParameterFlags) {
-      auto *flagsVar = flagsArr.finishAndCreateGlobal(
-          "parameter-flags", IGF.IGM.getPointerAlignment(),
-          /* constant */ true);
-      arguments.push_back(IGF.Builder.CreateBitCast(flagsVar, Int32Ptr));
+    FunctionTypeMetadataParamInfo info;
+    if (!hasPackExpansion) {
+      assert(!shapeExpression);
+      info = emitFunctionTypeMetadataParams(IGF, params, flags, request,
+                                            arguments);
     } else {
-      flagsArr.abandon();
-      arguments.push_back(llvm::ConstantPointerNull::get(Int32Ptr));
+      info = emitDynamicFunctionTypeMetadataParams(IGF, params, flags, packType,
+                                                   request, arguments);
     }
 
     arguments.push_back(result);
@@ -1608,9 +1674,7 @@ static MetadataResponse emitFunctionTypeMetadataRef(IRGenFunction &IGF,
     auto call = IGF.Builder.CreateCall(getMetadataFn, arguments);
     call->setDoesNotThrow();
 
-    if (parameters.isValid())
-      IGF.Builder.CreateLifetimeEnd(parameters,
-                                    IGF.IGM.getPointerSize() * numParams);
+    cleanupFunctionTypeMetadataParams(IGF, info);
 
     return MetadataResponse::forComplete(call);
   }