MandatoryGenericSpecializer: 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.

rdar://94388453

Author: eeckstein

include/swift/SILOptimizer/Utils/Generics.h

@@ -62,6 +62,11 @@ class ReabstractionInfo {
   /// argument has a trivial type.
   SmallBitVector TrivialArgs;
 
+  /// A 1-bit means that the argument is a metatype argument. The argument is
+  /// dropped and replaced by a `metatype` instruction in the entry block.
+  /// Only used if `dropMetatypeArgs` is true.
+  SmallBitVector droppedMetatypeArgs;
+
   /// Set to true if the function has a re-abstracted (= converted from
   /// indirect to direct) resilient argument or return type. This can happen if
   /// the function is compiled within the type's resilience domain, i.e. in
@@ -79,6 +84,10 @@ class ReabstractionInfo {
   /// specializer.
   bool ConvertIndirectToDirect;
 
+  /// If true, drop metatype arguments.
+  /// See `droppedMetatypeArgs`.
+  bool dropMetatypeArgs = false;
+  
   /// The first NumResults bits in Conversions refer to formal indirect
   /// out-parameters.
   unsigned NumFormalIndirectResults;
@@ -191,6 +200,7 @@ class ReabstractionInfo {
                     SubstitutionMap ParamSubs,
                     IsSerialized_t Serialized,
                     bool ConvertIndirectToDirect = true,
+                    bool dropMetatypeArgs = false,
                     OptRemark::Emitter *ORE = nullptr);
 
   /// Constructs the ReabstractionInfo for generic function \p Callee with
@@ -243,6 +253,16 @@ class ReabstractionInfo {
   /// Returns true if there are any conversions from indirect to direct values.
   bool hasConversions() const { return Conversions.any(); }
 
+  /// Returns true if the argument at `ArgIdx` is a dropped metatype argument.
+  /// See `droppedMetatypeArgs`.
+  bool isDroppedMetatypeArg(unsigned ArgIdx) const {
+    return droppedMetatypeArgs.test(ArgIdx);
+  }
+
+  /// Returns true if there are any dropped metatype arguments.
+  /// See `droppedMetatypeArgs`.
+  bool hasDroppedMetatypeArgs() const { return droppedMetatypeArgs.any(); }
+
   /// Remove the arguments of a partial apply, leaving the arguments for the
   /// partial apply result function.
   void prunePartialApplyArgs(unsigned numPartialApplyArgs) {

lib/SILOptimizer/Utils/GenericCloner.cpp

@@ -105,6 +105,12 @@ void GenericCloner::populateCloned() {
           entryArgs.push_back(ASI);
           return true;
         }
+      } else if (ReInfo.isDroppedMetatypeArg(ArgIdx)) {
+        // Replace the metatype argument with an `metatype` instruction in the
+        // entry block.
+        auto *mtInst = getBuilder().createMetatype(Loc, mappedType);
+        entryArgs.push_back(mtInst);
+        return true;
       } else {
         // Handle arguments for formal parameters.
         unsigned paramIdx = ArgIdx - origConv.getSILArgIndexOfFirstParam();

lib/SILOptimizer/Utils/Generics.cpp

@@ -550,8 +550,9 @@ bool ReabstractionInfo::canBeSpecialized(ApplySite Apply, SILFunction *Callee,
 ReabstractionInfo::ReabstractionInfo(
     ModuleDecl *targetModule, bool isWholeModule, ApplySite Apply,
     SILFunction *Callee, SubstitutionMap ParamSubs, IsSerialized_t Serialized,
-    bool ConvertIndirectToDirect, OptRemark::Emitter *ORE)
+    bool ConvertIndirectToDirect, bool dropMetatypeArgs, OptRemark::Emitter *ORE)
     : ConvertIndirectToDirect(ConvertIndirectToDirect),
+      dropMetatypeArgs(dropMetatypeArgs),
       TargetModule(targetModule), isWholeModule(isWholeModule),
       Serialized(Serialized) {
   if (!prepareAndCheck(Apply, Callee, ParamSubs, ORE))
@@ -683,6 +684,7 @@ void ReabstractionInfo::createSubstitutedAndSpecializedTypes() {
     SubstitutedType->getParameters().size();
   Conversions.resize(NumArgs);
   TrivialArgs.resize(NumArgs);
+  droppedMetatypeArgs.resize(NumArgs);
 
   SILFunctionConventions substConv(SubstitutedType, M);
   TypeExpansionContext resilienceExp = getResilienceExpansion();
@@ -737,11 +739,17 @@ void ReabstractionInfo::createSubstitutedAndSpecializedTypes() {
     case ParameterConvention::Indirect_In_Constant:
     case ParameterConvention::Indirect_Inout:
     case ParameterConvention::Indirect_InoutAliasable:
+      break;
+      
     case ParameterConvention::Direct_Owned:
     case ParameterConvention::Direct_Unowned:
-    case ParameterConvention::Direct_Guaranteed:
+    case ParameterConvention::Direct_Guaranteed: {
+      CanType ty = PI.getInterfaceType();
+      if (dropMetatypeArgs && isa<MetatypeType>(ty) && !ty->hasArchetype())
+        droppedMetatypeArgs.set(IdxToInsert);
       break;
     }
+    }
   }
 
   // Produce a specialized type, which is the substituted type with
@@ -852,11 +860,16 @@ createSpecializedType(CanSILFunctionType SubstFTy, SILModule &M) const {
     // No conversion: re-use the original, substituted result info.
     SpecializedResults.push_back(RI);
   }
-  unsigned ParamIdx = 0;
+  unsigned idx = 0;
   for (SILParameterInfo PI : SubstFTy->getParameters()) {
+    unsigned paramIdx = idx++;
     PI = PI.getUnsubstituted(M, SubstFTy, context);
-    bool isTrivial = TrivialArgs.test(param2ArgIndex(ParamIdx));
-    if (!isParamConverted(ParamIdx++)) {
+
+    if (isDroppedMetatypeArg(param2ArgIndex(paramIdx)))
+      continue;
+
+    bool isTrivial = TrivialArgs.test(param2ArgIndex(paramIdx));
+    if (!isParamConverted(paramIdx)) {
       // No conversion: re-use the original, substituted parameter info.
       SpecializedParams.push_back(PI);
       continue;
@@ -1854,7 +1867,8 @@ GenericFuncSpecializer::GenericFuncSpecializer(
       ClonedName = Mangler.manglePrespecialized(ParamSubs);
     } else {
       ClonedName = Mangler.mangleReabstracted(ParamSubs,
-                                              ReInfo.needAlternativeMangling());
+                                              ReInfo.needAlternativeMangling(),
+                                              ReInfo.hasDroppedMetatypeArgs());
     }
   }
   LLVM_DEBUG(llvm::dbgs() << "    Specialized function " << ClonedName << '\n');
@@ -1995,6 +2009,9 @@ prepareCallArguments(ApplySite AI, SILBuilder &Builder,
       return true;
     }
 
+    if (ReInfo.isDroppedMetatypeArg(ArgIdx))
+      return true;
+
     // Handle arguments for formal parameters.
     unsigned paramIdx = ArgIdx - substConv.getSILArgIndexOfFirstParam();
     if (!ReInfo.isParamConverted(paramIdx)) {
@@ -2418,7 +2435,7 @@ static bool createPrespecialized(StringRef UnspecializedName,
   ReabstractionInfo ReInfo(M.getSwiftModule(), M.isWholeModule(), ApplySite(),
                            UnspecFunc, Apply.getSubstitutionMap(),
                            IsNotSerialized,
-                           /*ConvertIndirectToDirect=*/true, nullptr);
+                           /*ConvertIndirectToDirect=*/true);
 
   if (!ReInfo.canBeSpecialized())
     return false;
@@ -2599,7 +2616,9 @@ void swift::trySpecializeApplyOfGeneric(
   ReabstractionInfo ReInfo(FuncBuilder.getModule().getSwiftModule(),
                            FuncBuilder.getModule().isWholeModule(), Apply, RefF,
                            Apply.getSubstitutionMap(), Serialized,
-                           /*ConvertIndirectToDirect=*/true, &ORE);
+                           /*ConvertIndirectToDirect=*/ true,
+                           /*dropMetatypeArgs=*/ isMandatory,
+                           &ORE);
   if (!ReInfo.canBeSpecialized())
     return;
 

test/SILOptimizer/performance-annotations.swift

@@ -131,3 +131,11 @@ func testGlobalWithComplexInit() -> Int {
   return Str.s3 // expected-note {{called from here}}
 }
 
+func metatypeArg<T>(_ t: T.Type, _ b: Bool) { // expected-error {{Using type 'Int' can cause metadata allocation or locks}}
+}
+
+@_noAllocation
+func callFuncWithMetatypeArg() {
+  metatypeArg(Int.self, false)                // expected-note {{called from here}}
+}
+