GenericSpecializer: drop metatype arguments in specialized functions

And replace them with explicit `metatype` instruction in the entry block.
This allows such metatype instructions to be deleted if they are dead.

This was already done for performance-annotated functions. But now do this for all functions.

It is essential that performance-annotated functions are specialized in the same way as other functions.
Because otherwise it can happen that the same specialization has different performance characteristics in different modules.
And it's up to the linker to select one of those ODR functions when linking.

Also, dropping metatype arguments is good for performance and code size in general.

This change also contains a few bug fixes for dropping metatype arguments.

rdar://110509780

Author: eeckstein

include/swift/SILOptimizer/Utils/Generics.h

@@ -157,10 +157,6 @@ class ReabstractionInfo {
     LoadableAndTrivial
   };
 
-  unsigned param2ArgIndex(unsigned ParamIdx) const  {
-    return ParamIdx + NumFormalIndirectResults;
-  }
-
   // Create a new substituted type with the updated signature.
   CanSILFunctionType createSubstitutedType(SILFunction *OrigF,
                                            SubstitutionMap SubstMap,
@@ -199,8 +195,8 @@ class ReabstractionInfo {
                     ApplySite Apply, SILFunction *Callee,
                     SubstitutionMap ParamSubs,
                     IsSerialized_t Serialized,
-                    bool ConvertIndirectToDirect = true,
-                    bool dropMetatypeArgs = false,
+                    bool ConvertIndirectToDirect,
+                    bool dropMetatypeArgs,
                     OptRemark::Emitter *ORE = nullptr);
 
   /// Constructs the ReabstractionInfo for generic function \p Callee with
@@ -214,7 +210,11 @@ class ReabstractionInfo {
   IsSerialized_t isSerialized() const {
     return Serialized;
   }
-  
+
+  unsigned param2ArgIndex(unsigned ParamIdx) const  {
+    return ParamIdx + NumFormalIndirectResults;
+  }
+
   /// Returns true if the specialized function needs an alternative mangling.
   /// See hasConvertedResilientParams.
   bool needAlternativeMangling() const {
@@ -314,6 +314,8 @@ class ReabstractionInfo {
   CanSILFunctionType createSpecializedType(CanSILFunctionType SubstFTy,
                                            SILModule &M) const;
 
+  CanSILFunctionType createThunkType(PartialApplyInst *forPAI) const;
+
   SILFunction *getNonSpecializedFunction() const { return Callee; }
 
   /// Map type into a context of the specialized function.

lib/SILOptimizer/IPO/CapturePropagation.cpp

@@ -492,7 +492,7 @@ static SILFunction *getSpecializedWithDeadParams(
         FuncBuilder.getModule().getSwiftModule(),
         FuncBuilder.getModule().isWholeModule(), ApplySite(), Specialized,
         PAI->getSubstitutionMap(), Specialized->isSerialized(),
-        /* ConvertIndirectToDirect */ false);
+        /* ConvertIndirectToDirect */ false, /*dropMetatypeArgs=*/ false);
     GenericFuncSpecializer FuncSpecializer(FuncBuilder,
                                            Specialized,
                                            ReInfo.getClonerParamSubstitutionMap(),

lib/SILOptimizer/IPO/UsePrespecialized.cpp

@@ -91,7 +91,9 @@ bool UsePrespecialized::replaceByPrespecialized(SILFunction &F) {
       continue;
 
     ReabstractionInfo ReInfo(M.getSwiftModule(), M.isWholeModule(), AI,
-                             ReferencedF, Subs, IsNotSerialized);
+                             ReferencedF, Subs, IsNotSerialized,
+                             /*ConvertIndirectToDirect=*/ true,
+                             /*dropMetatypeArgs=*/ false);
 
     if (!ReInfo.canBeSpecialized())
       continue;

lib/SILOptimizer/Utils/Generics.cpp

@@ -997,6 +997,35 @@ ReabstractionInfo::createSubstitutedType(SILFunction *OrigF,
   return NewFnTy;
 }
 
+CanSILFunctionType ReabstractionInfo::createThunkType(PartialApplyInst *forPAI) const {
+  if (!hasDroppedMetatypeArgs())
+    return SubstitutedType;
+
+  llvm::SmallVector<SILParameterInfo, 8> newParams;
+  auto params = SubstitutedType->getParameters();
+  unsigned firstAppliedParamIdx = params.size() - forPAI->getArguments().size();
+
+  for (unsigned paramIdx = 0; paramIdx < params.size(); ++paramIdx) {
+    if (paramIdx >= firstAppliedParamIdx && isDroppedMetatypeArg(param2ArgIndex(paramIdx)))
+      continue;
+    newParams.push_back(params[paramIdx]);
+  }
+
+  auto newFnTy = SILFunctionType::get(
+      SubstitutedType->getInvocationGenericSignature(),
+      SubstitutedType->getExtInfo(), SubstitutedType->getCoroutineKind(),
+      SubstitutedType->getCalleeConvention(), newParams,
+      SubstitutedType->getYields(), SubstitutedType->getResults(),
+      SubstitutedType->getOptionalErrorResult(),
+      SubstitutedType->getPatternSubstitutions(), SubstitutionMap(),
+      SubstitutedType->getASTContext(),
+      SubstitutedType->getWitnessMethodConformanceOrInvalid());
+
+  // This is an interface type. It should not have any archetypes.
+  assert(!newFnTy->hasArchetype());
+  return newFnTy;
+}
+
 /// Convert the substituted function type into a specialized function type based
 /// on the ReabstractionInfo.
 CanSILFunctionType ReabstractionInfo::
@@ -1067,8 +1096,10 @@ createSpecializedType(CanSILFunctionType SubstFTy, SILModule &M) const {
                      : CanGenericSignature();
 
   SILFunctionType::ExtInfo extInfo = SubstFTy->getExtInfo();
+  ProtocolConformanceRef conf = SubstFTy->getWitnessMethodConformanceOrInvalid();
   if (extInfo.hasSelfParam() && removedSelfParam) {
     extInfo = extInfo.withRepresentation(SILFunctionTypeRepresentation::Thin);
+    conf = ProtocolConformanceRef();
     assert(!extInfo.hasSelfParam());
   }
 
@@ -1077,7 +1108,7 @@ createSpecializedType(CanSILFunctionType SubstFTy, SILModule &M) const {
       SubstFTy->getCoroutineKind(), SubstFTy->getCalleeConvention(),
       SpecializedParams, SpecializedYields, SpecializedResults,
       SubstFTy->getOptionalErrorResult(), SubstitutionMap(), SubstitutionMap(),
-      M.getASTContext(), SubstFTy->getWitnessMethodConformanceOrInvalid());
+      M.getASTContext(), conf);
 }
 
 /// Create a new generic signature from an existing one by adding
@@ -2546,14 +2577,16 @@ class ReabstractionThunkGenerator {
 protected:
   FullApplySite createReabstractionThunkApply(SILBuilder &Builder);
   SILArgument *convertReabstractionThunkArguments(
-      SILBuilder &Builder, SmallVectorImpl<unsigned> &ArgsNeedingEndBorrows);
+      SILBuilder &Builder, SmallVectorImpl<unsigned> &ArgsNeedingEndBorrows,
+      CanSILFunctionType thunkType);
 };
 
 } // anonymous namespace
 
 SILFunction *ReabstractionThunkGenerator::createThunk() {
+  CanSILFunctionType thunkType = ReInfo.createThunkType(OrigPAI);
   SILFunction *Thunk = FunctionBuilder.getOrCreateSharedFunction(
-      Loc, ThunkName, ReInfo.getSubstitutedType(), IsBare, IsTransparent,
+      Loc, ThunkName, thunkType, IsBare, IsTransparent,
       ReInfo.isSerialized(), ProfileCounter(), IsThunk, IsNotDynamic,
       IsNotDistributed, IsNotRuntimeAccessible);
   // Re-use an existing thunk.
@@ -2594,7 +2627,7 @@ SILFunction *ReabstractionThunkGenerator::createThunk() {
   // Handle lowered addresses.
   SmallVector<unsigned, 4> ArgsThatNeedEndBorrow;
   SILArgument *ReturnValueAddr =
-      convertReabstractionThunkArguments(Builder, ArgsThatNeedEndBorrow);
+      convertReabstractionThunkArguments(Builder, ArgsThatNeedEndBorrow, thunkType);
 
   FullApplySite ApplySite = createReabstractionThunkApply(Builder);
 
@@ -2650,50 +2683,50 @@ FullApplySite ReabstractionThunkGenerator::createReabstractionThunkApply(
   return FullApplySite(TAI);
 }
 
+static SILFunctionArgument *addFunctionArgument(SILFunction *function,
+                                                SILType argType,
+                                                SILArgument *copyAttributesFrom) {
+  SILBasicBlock *entryBB = function->getEntryBlock();
+  auto *src = cast<SILFunctionArgument>(copyAttributesFrom);
+  auto *arg = entryBB->createFunctionArgument(argType, src->getDecl());
+  arg->setNoImplicitCopy(src->isNoImplicitCopy());
+  arg->setLifetimeAnnotation(src->getLifetimeAnnotation());
+  arg->setClosureCapture(src->isClosureCapture());
+  return arg;
+}
+
 /// Create SIL arguments for a reabstraction thunk with lowered addresses. This
 /// may involve replacing indirect arguments with loads and stores. Return the
 /// SILArgument for the address of an indirect result, or nullptr.
 ///
 /// FIXME: Remove this if we don't need to create reabstraction thunks after
 /// address lowering.
 SILArgument *ReabstractionThunkGenerator::convertReabstractionThunkArguments(
-    SILBuilder &Builder, SmallVectorImpl<unsigned> &ArgsThatNeedEndBorrow) {
+    SILBuilder &Builder, SmallVectorImpl<unsigned> &ArgsThatNeedEndBorrow,
+    CanSILFunctionType thunkType
+) {
   SILFunction *Thunk = &Builder.getFunction();
   CanSILFunctionType SpecType = SpecializedFunc->getLoweredFunctionType();
-  CanSILFunctionType SubstType = ReInfo.getSubstitutedType();
   auto specConv = SpecializedFunc->getConventions();
   (void)specConv;
-  SILFunctionConventions substConv(SubstType, M);
+  SILFunctionConventions substConv(thunkType, M);
 
   assert(specConv.useLoweredAddresses());
 
   // ReInfo.NumIndirectResults corresponds to SubstTy's formal indirect
   // results. SpecTy may have fewer formal indirect results.
-  assert(SubstType->getNumIndirectFormalResults()
+  assert(thunkType->getNumIndirectFormalResults()
          >= SpecType->getNumIndirectFormalResults());
 
-  SILBasicBlock *EntryBB = Thunk->getEntryBlock();
   SILArgument *ReturnValueAddr = nullptr;
   auto SpecArgIter = SpecializedFunc->getArguments().begin();
-  auto cloneSpecializedArgument = [&]() {
-    // No change to the argument.
-    SILArgument *SpecArg = *SpecArgIter++;
-    auto *NewArg =
-        EntryBB->createFunctionArgument(SpecArg->getType(), SpecArg->getDecl());
-    NewArg->setNoImplicitCopy(
-        cast<SILFunctionArgument>(SpecArg)->isNoImplicitCopy());
-    NewArg->setLifetimeAnnotation(
-        cast<SILFunctionArgument>(SpecArg)->getLifetimeAnnotation());
-    NewArg->setClosureCapture(
-        cast<SILFunctionArgument>(SpecArg)->isClosureCapture());
-    Arguments.push_back(NewArg);
-  };
+
   // ReInfo.NumIndirectResults corresponds to SubstTy's formal indirect
   // results. SpecTy may have fewer formal indirect results.
-  assert(SubstType->getNumIndirectFormalResults()
+  assert(thunkType->getNumIndirectFormalResults()
          >= SpecType->getNumIndirectFormalResults());
   unsigned resultIdx = 0;
-  for (auto substRI : SubstType->getIndirectFormalResults()) {
+  for (auto substRI : thunkType->getIndirectFormalResults()) {
     if (ReInfo.isFormalResultConverted(resultIdx++)) {
       // Convert an originally indirect to direct specialized result.
       // Store the result later.
@@ -2703,36 +2736,42 @@ SILArgument *ReabstractionThunkGenerator::convertReabstractionThunkArguments(
           substConv.getSILType(substRI, Builder.getTypeExpansionContext()));
       assert(ResultTy.isAddress());
       assert(!ReturnValueAddr);
-      ReturnValueAddr = EntryBB->createFunctionArgument(ResultTy);
+      ReturnValueAddr = Thunk->getEntryBlock()->createFunctionArgument(ResultTy);
       continue;
     }
     // If the specialized result is already indirect, simply clone the indirect
     // result argument.
-    assert((*SpecArgIter)->getType().isAddress());
-    cloneSpecializedArgument();
+    SILArgument *specArg = *SpecArgIter++;
+    assert(specArg->getType().isAddress());
+    Arguments.push_back(addFunctionArgument(Thunk, specArg->getType(), specArg));
   }
   assert(SpecArgIter
          == SpecializedFunc->getArgumentsWithoutIndirectResults().begin());
   unsigned numParams = SpecType->getNumParameters();
-  assert(numParams == SubstType->getNumParameters());
+  assert(numParams == thunkType->getNumParameters());
   for (unsigned paramIdx = 0; paramIdx < numParams; ++paramIdx) {
+    if (ReInfo.isDroppedMetatypeArg(paramIdx)) {
+      if (paramIdx < ApplySite(OrigPAI).getCalleeArgIndexOfFirstAppliedArg()) {
+        // Only if the dropped metatype argument is not applied by the `partial_apply`,
+        // we need to add it to the thunk.
+        SILType ParamTy = SpecializedFunc->mapTypeIntoContext(
+            substConv.getSILType(thunkType->getParameters()[paramIdx],
+                                 Builder.getTypeExpansionContext()));
+        SILArgument *origArg = OrigF->getArgumentsWithoutIndirectResults()[paramIdx];
+        addFunctionArgument(Thunk, ParamTy, origArg);
+      }
+      continue;
+    }
+    SILArgument *specArg = *SpecArgIter++;
     if (ReInfo.isParamConverted(paramIdx)) {
       // Convert an originally indirect to direct specialized parameter.
       assert(!specConv.isSILIndirect(SpecType->getParameters()[paramIdx]));
       // Instead of passing the address, pass the loaded value.
       SILType ParamTy = SpecializedFunc->mapTypeIntoContext(
-          substConv.getSILType(SubstType->getParameters()[paramIdx],
+          substConv.getSILType(thunkType->getParameters()[paramIdx],
                                Builder.getTypeExpansionContext()));
       assert(ParamTy.isAddress());
-      SILArgument *SpecArg = *SpecArgIter++;
-      SILFunctionArgument *NewArg =
-          EntryBB->createFunctionArgument(ParamTy, SpecArg->getDecl());
-      NewArg->setNoImplicitCopy(
-          cast<SILFunctionArgument>(SpecArg)->isNoImplicitCopy());
-      NewArg->setLifetimeAnnotation(
-          cast<SILFunctionArgument>(SpecArg)->getLifetimeAnnotation());
-      NewArg->setClosureCapture(
-          cast<SILFunctionArgument>(SpecArg)->isClosureCapture());
+      SILFunctionArgument *NewArg = addFunctionArgument(Thunk, ParamTy, specArg);
       if (!NewArg->getArgumentConvention().isGuaranteedConvention()) {
         SILValue argVal = Builder.emitLoadValueOperation(
             Loc, NewArg, LoadOwnershipQualifier::Take);
@@ -2746,7 +2785,7 @@ SILArgument *ReabstractionThunkGenerator::convertReabstractionThunkArguments(
       continue;
     }
     // Simply clone unconverted direct or indirect parameters.
-    cloneSpecializedArgument();
+    Arguments.push_back(addFunctionArgument(Thunk, specArg->getType(), specArg));
   }
   assert(SpecArgIter == SpecializedFunc->getArguments().end());
   return ReturnValueAddr;
@@ -2773,7 +2812,7 @@ static bool createPrespecialized(StringRef UnspecializedName,
   ReabstractionInfo ReInfo(M.getSwiftModule(), M.isWholeModule(), ApplySite(),
                            UnspecFunc, Apply.getSubstitutionMap(),
                            IsNotSerialized,
-                           /*ConvertIndirectToDirect=*/true);
+                           /*ConvertIndirectToDirect= */true, /*dropMetatypeArgs=*/ false);
 
   if (!ReInfo.canBeSpecialized())
     return false;
@@ -2945,7 +2984,7 @@ bool usePrespecialized(
           funcBuilder.getModule().getSwiftModule(),
           funcBuilder.getModule().isWholeModule(), apply, refF, newSubstMap,
           apply.getFunction()->isSerialized() ? IsSerialized : IsNotSerialized,
-          /*ConvertIndirectToDirect=*/true, /*dropMetatypeArgs*/true, nullptr);
+          /*ConvertIndirectToDirect=*/ true, /*dropMetatypeArgs=*/ false, nullptr);
 
       if (layoutReInfo.getSpecializedType() == reInfo.getSpecializedType()) {
         layoutMatches.push_back(
@@ -3005,6 +3044,28 @@ bool usePrespecialized(
   return false;
 }
 
+static bool canDropMetatypeArgs(ApplySite apply, SILFunction *callee) {
+  for (const Operand &argOp : apply.getArgumentOperands()) {
+    SILValue arg = argOp.get();
+    if (auto mt = arg->getType().getAs<MetatypeType>()) {
+      if (mt->hasRepresentation() && mt->getRepresentation() == MetatypeRepresentation::Thin)
+        continue;
+      // If the passed thick metatype value is not a `metatype` instruction
+      // we don't know the real metatype at runtime. It's not necessarily the
+      // same as the declared metatype. It could e.g. be an upcast of a class
+      // metatype.
+      if (isa<MetatypeInst>(arg))
+        continue;
+      // But: if the metatype is not used in the callee we don't have to care
+      // what metatype value is passed. We can just remove it.
+      if (callee->isDefinition() && onlyHaveDebugUses(callee->getArgument(apply.getCalleeArgIndex(argOp))))
+        continue;
+      return false;
+    }
+  }
+  return true;
+}
+
 void swift::trySpecializeApplyOfGeneric(
     SILOptFunctionBuilder &FuncBuilder,
     ApplySite Apply, DeadInstructionSet &DeadApplies,
@@ -3053,7 +3114,7 @@ void swift::trySpecializeApplyOfGeneric(
                            FuncBuilder.getModule().isWholeModule(), Apply, RefF,
                            Apply.getSubstitutionMap(), Serialized,
                            /*ConvertIndirectToDirect=*/ true,
-                           /*dropMetatypeArgs=*/ isMandatory,
+                           /*dropMetatypeArgs=*/ canDropMetatypeArgs(Apply, RefF),
                            &ORE);
   if (!ReInfo.canBeSpecialized())
     return;
@@ -3210,16 +3271,19 @@ void swift::trySpecializeApplyOfGeneric(
     auto *FRI = Builder.createFunctionRef(PAI->getLoc(), Thunk);
     SmallVector<SILValue, 4> Arguments;
     for (auto &Op : PAI->getArgumentOperands()) {
+      unsigned calleeArgIdx = ApplySite(PAI).getCalleeArgIndex(Op);
+      if (ReInfo.isDroppedMetatypeArg(calleeArgIdx))
+        continue;
       Arguments.push_back(Op.get());
     }
     auto Subs = ReInfo.getCallerParamSubstitutionMap();
     auto FnTy = Thunk->getLoweredFunctionType();
     Subs = SubstitutionMap::get(FnTy->getSubstGenericSignature(), Subs);
-    auto *NewPAI = Builder.createPartialApply(
-        PAI->getLoc(), FRI, Subs, Arguments,
-        PAI->getType().getAs<SILFunctionType>()->getCalleeConvention(),
-        PAI->isOnStack());
-    PAI->replaceAllUsesWith(NewPAI);
+    SingleValueInstruction *newPAI = Builder.createPartialApply(
+      PAI->getLoc(), FRI, Subs, Arguments,
+      PAI->getType().getAs<SILFunctionType>()->getCalleeConvention(),
+      PAI->isOnStack());
+    PAI->replaceAllUsesWith(newPAI);
     DeadApplies.insert(PAI);
     return;
   }

test/SIL/Serialization/specializer_can_deserialize.swift

@@ -8,10 +8,10 @@ import Swift
 
 // CHECK-LABEL: sil {{.*}}@main
 // CHECK: bb0({{.*}}):
-// CHECK: function_ref @$ss9ContainerVAByxGycfCBi32__Tg5{{.*}}
+// CHECK: function_ref @$ss9ContainerVAByxGycfCBi32__Tgm5{{.*}}
 // CHECK: function_ref @$ss9ContainerV11doSomethingyyFBi32__Tg5{{.*}} 
 
-// CHECK-LABEL: sil shared [noinline] @$ss9ContainerVAByxGycfCBi32__Tg5Tf4d_n
+// CHECK-LABEL: sil shared [noinline] @$ss9ContainerVAByxGycfCBi32__Tgm5
 
 // CHECK-LABEL: sil shared [noinline] @$ss9ContainerV11doSomethingyyFBi32__Tg5Tf4d_n
 

test/SILOptimizer/existential_metatype.swift

@@ -3,12 +3,12 @@ protocol SomeP {}
 
 public enum SpecialEnum : SomeP {}
 
-// CHECK-LABEL: sil shared [noinline] @$s20existential_metatype17checkProtocolType0aE0Sbxm_tAA5SomePRzlFAA11SpecialEnumO_Tg5Tf4d_n : $@convention(thin) () -> Bool {
+// CHECK-LABEL: sil shared [noinline] @$s20existential_metatype17checkProtocolType0aE0Sbxm_tAA5SomePRzlFAA11SpecialEnumO_Tgm5 : $@convention(thin) () -> Bool {
 // CHECK:       bb0:
 // CHECK-NEXT:   %0 = integer_literal $Builtin.Int1, -1
 // CHECK-NEXT:   %1 = struct $Bool (%0 : $Builtin.Int1)
 // CHECK-NEXT:   return %1 : $Bool
-// CHECK-LABEL: } // end sil function '$s20existential_metatype17checkProtocolType0aE0Sbxm_tAA5SomePRzlFAA11SpecialEnumO_Tg5Tf4d_n'
+// CHECK-LABEL: } // end sil function '$s20existential_metatype17checkProtocolType0aE0Sbxm_tAA5SomePRzlFAA11SpecialEnumO_Tgm5'
 @inline(never)
 func checkProtocolType<P : SomeP>(existentialType: P.Type) -> Bool {
   return existentialType == SpecialEnum.self