MandatoryPerformanceOptimizations: specialize witness_method instructions

In Embedded Swift, witness method lookup is done from specialized witness tables.
For this to work, the type of witness_method must be specialized as well.
Otherwise the method call would be done with wrong parameter conventions (indirect instead of direct).

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/ModulePasses/MandatoryPerformanceOptimizations.swift

@@ -116,6 +116,10 @@ private func optimize(function: Function, _ context: FunctionPassContext, _ modu
         if context.options.enableEmbeddedSwift {
           _ = context.specializeClassMethodInst(classMethod)
         }
+      case let witnessMethod as WitnessMethodInst:
+        if context.options.enableEmbeddedSwift {
+          _ = context.specializeWitnessMethodInst(witnessMethod)
+        }
 
       case let initExRef as InitExistentialRefInst:
         if context.options.enableEmbeddedSwift {

SwiftCompilerSources/Sources/Optimizer/PassManager/Context.swift

@@ -357,6 +357,15 @@ struct FunctionPassContext : MutatingContext {
     return false
   }
 
+  func specializeWitnessMethodInst(_ wm: WitnessMethodInst) -> Bool {
+    if _bridged.specializeWitnessMethodInst(wm.bridged) {
+      notifyInstructionsChanged()
+      notifyCallsChanged()
+      return true
+    }
+    return false
+  }
+
   func specializeApplies(in function: Function, isMandatory: Bool) -> Bool {
     if _bridged.specializeAppliesInFunction(function.bridged, isMandatory) {
       notifyInstructionsChanged()

include/swift/SILOptimizer/OptimizerBridging.h

@@ -256,6 +256,7 @@ struct BridgedPassContext {
                                                                  BridgedSubstitutionMap substitutions) const;
   void deserializeAllCallees(BridgedFunction function, bool deserializeAll) const;
   bool specializeClassMethodInst(BridgedInstruction cm) const;
+  bool specializeWitnessMethodInst(BridgedInstruction wm) const;
   bool specializeAppliesInFunction(BridgedFunction function, bool isMandatory) const;
   BridgedOwnedString mangleOutlinedVariable(BridgedFunction function) const;
   BridgedOwnedString mangleAsyncRemoved(BridgedFunction function) const;

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -607,6 +607,7 @@ bool optimizeMemoryAccesses(SILFunction *fn);
 bool eliminateDeadAllocations(SILFunction *fn, DominanceInfo *domInfo);
 
 bool specializeClassMethodInst(ClassMethodInst *cm);
+bool specializeWitnessMethodInst(WitnessMethodInst *wm);
 
 bool specializeAppliesInFunction(SILFunction &F,
                                  SILTransform *transform,

lib/SILOptimizer/PassManager/PassManager.cpp

@@ -1771,6 +1771,10 @@ bool BridgedPassContext::specializeClassMethodInst(BridgedInstruction cm) const
   return ::specializeClassMethodInst(cm.getAs<ClassMethodInst>());
 }
 
+bool BridgedPassContext::specializeWitnessMethodInst(BridgedInstruction wm) const {
+  return ::specializeWitnessMethodInst(wm.getAs<WitnessMethodInst>());
+}
+
 bool BridgedPassContext::specializeAppliesInFunction(BridgedFunction function, bool isMandatory) const {
   return ::specializeAppliesInFunction(*function.getFunction(), invocation->getTransform(), isMandatory);
 }

lib/SILOptimizer/Utils/Generics.cpp

@@ -725,8 +725,6 @@ void ReabstractionInfo::createSubstitutedAndSpecializedTypes() {
     SubstitutedType = createSubstitutedType(Callee, CallerInterfaceSubs,
                                             HasUnboundGenericParams);
   }
-  assert(!SubstitutedType->hasArchetype() &&
-         "Substituted function type should not contain archetypes");
 
   // Check which parameters and results can be converted from
   // indirect to direct ones.
@@ -2363,6 +2361,58 @@ bool swift::specializeClassMethodInst(ClassMethodInst *cm) {
   return true;
 }
 
+bool swift::specializeWitnessMethodInst(WitnessMethodInst *wm) {
+  SILFunction *f = wm->getFunction();
+  SILModule &m = f->getModule();
+
+  CanType astType = wm->getLookupType();
+  if (!isa<OpenedArchetypeType>(astType))
+    return false;
+
+  if (wm->isSpecialized())
+    return false;
+
+  // This should not happen. Just to be on the safe side.
+  if (wm->use_empty())
+    return false;
+
+  Operand *firstUse = *wm->use_begin();
+  ApplySite AI = ApplySite::isa(firstUse->getUser());
+  assert(AI && AI.getCalleeOperand() == firstUse && "wrong use of witness_method instruction");
+
+  SubstitutionMap subs = AI.getSubstitutionMap();
+
+  SILType funcTy = wm->getType();
+  SILType substitutedType =
+      funcTy.substGenericArgs(m, subs, TypeExpansionContext::minimal());
+
+  ReabstractionInfo reInfo(substitutedType.getAs<SILFunctionType>(), wm->getMember(), m);
+  reInfo.createSubstitutedAndSpecializedTypes();
+  CanSILFunctionType finalFuncTy = reInfo.getSpecializedType();
+  SILType finalSILTy = SILType::getPrimitiveObjectType(finalFuncTy);
+
+  SILBuilder builder(wm);
+  auto *newWM = builder.createWitnessMethod(wm->getLoc(), wm->getLookupType(),
+                                            wm->getConformance(), wm->getMember(), finalSILTy);
+
+  while (!wm->use_empty()) {
+    Operand *use = *wm->use_begin();
+    SILInstruction *user = use->getUser();
+    ApplySite AI = ApplySite::isa(user);
+    if (AI && AI.getCalleeOperand() == use) {
+      replaceWithSpecializedCallee(AI, newWM, reInfo);
+      AI.getInstruction()->eraseFromParent();
+      continue;
+    }
+    llvm::errs() << "unsupported use of witness method "
+                 << newWM->getMember().getDecl()->getName() << " in function "
+                 << newWM->getFunction()->getName() << '\n';
+    llvm::report_fatal_error("unsupported witness method");
+  }
+
+  return true;
+}
+
 /// Create a new apply based on an old one, but with a different
 /// function being applied.
 ApplySite