[sil-devirtualizer] Improve devirtualization of witness_method instructions.

Handle such cases like partial applications of witness methods and applications of witness methods with substitutions.

Some of these uses-cases occur when there is a protocol defining an operator, a generic struct conforming to this protocol, and the operator conformance of this struct is expressed as a global function.

Author: swiftix

include/swift/SIL/SILInstruction.h

@@ -652,6 +652,16 @@ class ApplyInstBase<Impl, Base, false> : public Base {
     return {getSubstitutionsStorage(), NumSubstitutions};
   }
 
+  ArrayRef<Substitution> getSubstitutionsWithoutSelfSubstitution() const {
+    assert(getNumArguments() && "Should only be called when Callee has "
+           "at least a self parameter.");
+    assert(hasSubstitutions() && "Should only be called when Callee has "
+           "substitutions.");
+    if (getSubstCalleeType()->hasSelfParam())
+      return getSubstitutions().slice(1);
+    return getSubstitutions();
+  }
+
   /// The arguments passed to this instruction.
   MutableArrayRef<Operand> getArgumentOperands() {
     return Operands.getDynamicAsArray();
@@ -4266,6 +4276,8 @@ class ApplySite {
       return cast<ApplyInst>(Inst)->getSubstitutionsWithoutSelfSubstitution();
     case ValueKind::TryApplyInst:
       return cast<TryApplyInst>(Inst)->getSubstitutionsWithoutSelfSubstitution();
+    case ValueKind::PartialApplyInst:
+      return cast<PartialApplyInst>(Inst)->getSubstitutionsWithoutSelfSubstitution();
     default:
       llvm_unreachable("not implemented for this instruction!");
     }

include/swift/SILPasses/Utils/Devirtualize.h

@@ -40,7 +40,7 @@ namespace swift {
 /// Two elements are required, because a result of the new devirtualized
 /// apply/try_apply instruction (second element) eventually needs to be
 /// casted to produce a properly typed value (first element).
-typedef std::pair<ValueBase *, FullApplySite> DevirtualizationResult;
+typedef std::pair<ValueBase *, ApplySite> DevirtualizationResult;
 
 DevirtualizationResult tryDevirtualizeApply(FullApplySite AI);
 bool isClassWithUnboundGenericParameters(SILType C, SILModule &M);
@@ -49,6 +49,7 @@ DevirtualizationResult devirtualizeClassMethod(FullApplySite AI,
                                                SILValue ClassInstance);
 DevirtualizationResult tryDevirtualizeClassMethod(FullApplySite AI,
                                                   SILValue ClassInstance);
+DevirtualizationResult tryDevirtualizeWitnessMethod(ApplySite AI); 
 }
 
 #endif

lib/IRGen/GenFunc.cpp

@@ -3591,7 +3591,9 @@ static llvm::Function *emitPartialApplicationForwarder(IRGenModule &IGM,
   }
 
   // Emit the polymorphic arguments.
-  assert(subs.empty() != hasPolymorphicParameters(origType)
+  assert((subs.empty() != hasPolymorphicParameters(origType) ||
+         (subs.empty() && origType->getRepresentation() ==
+             SILFunctionTypeRepresentation::WitnessMethod))
          && "should have substitutions iff original function is generic");
   WitnessMetadata witnessMetadata;
   if (hasPolymorphicParameters(origType)) {

lib/IRGen/IRGenSIL.cpp

@@ -2104,7 +2104,6 @@ getPartialApplicationFunction(IRGenSILFunction &IGF,
       break;
 
     case SILFunctionTypeRepresentation::WitnessMethod:
-      assert(false && "partial_apply of witness functions not implemented");
       break;
 
     case SILFunctionTypeRepresentation::Thick:

lib/SILPasses/EarlySIL/MandatoryInlining.cpp

@@ -320,14 +320,14 @@ runOnFunctionRecursively(SILFunction *F, FullApplySite AI,
           I = II->getIterator();
         else
           I = NewInst->getParentBB()->begin();
-        auto NewAI = NewInstPair.second;
+        auto NewAI = FullApplySite::isa(NewInstPair.second.getInstruction());
         if (!NewAI)
           continue;
 
         InnerAI = NewAI;
       }
 
-      SILLocation Loc = InnerAI.getLoc();
+	  SILLocation Loc = InnerAI.getLoc();
       SILValue CalleeValue = InnerAI.getCallee();
       bool IsThick;
       PartialApplyInst *PAI;

lib/SILPasses/IPO/PerformanceInliner.cpp

@@ -826,7 +826,7 @@ FullApplySite SILPerformanceInliner::devirtualizeUpdatingCallGraph(
   if (!NewInstPair.second)
     return FullApplySite();
 
-  auto NewAI = NewInstPair.second;
+  auto NewAI = FullApplySite::isa(NewInstPair.second.getInstruction());
   CallGraphEditor(&CG).replaceApplyWithNew(Apply, NewAI);
 
   replaceDeadApply(Apply, NewInstPair.first);

lib/SILPasses/SILCombiner/SILCombiner.h

@@ -243,7 +243,7 @@ class SILCombiner :
   SILInstruction *visitAllocRefDynamicInst(AllocRefDynamicInst *ARDI);
   SILInstruction *visitEnumInst(EnumInst *EI);
   SILInstruction *visitConvertFunctionInst(ConvertFunctionInst *CFI);
-
+  SILInstruction *visitWitnessMethodInst(WitnessMethodInst *WMI);
 
   /// Instruction visitor helpers.
   SILInstruction *optimizeBuiltinCanBeObjCClass(BuiltinInst *AI);

lib/SILPasses/SILCombiner/SILCombinerMiscVisitors.cpp

@@ -23,6 +23,7 @@
 #include "swift/SILAnalysis/CFG.h"
 #include "swift/SILAnalysis/ValueTracking.h"
 #include "swift/SILPasses/Utils/Local.h"
+#include "swift/SILPasses/Utils/Devirtualize.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/DenseMap.h"
@@ -1070,3 +1071,33 @@ visitAllocRefDynamicInst(AllocRefDynamicInst *ARDI) {
 SILInstruction *SILCombiner::visitEnumInst(EnumInst *EI) {
   return nullptr;
 }
+
+SILInstruction *SILCombiner::visitWitnessMethodInst(WitnessMethodInst *WMI) {
+  // Try to devirtualize a witness_method if it is statically possible.
+  // Many cases are handled by the inliner/devirtualizer, but certain
+  // special cases are not covered there, e.g. partial_apply(witness_method)
+  SILFunction *F;
+  ArrayRef<Substitution> Subs;
+  SILWitnessTable *WT;
+
+  std::tie(F, WT, Subs) =
+    WMI->getModule().lookUpFunctionInWitnessTable(WMI->getConformance(),
+						                          WMI->getMember());
+
+  if (!F)
+    return nullptr;
+
+  for (auto U = WMI->use_begin(), E = WMI->use_end(); U != E;) {
+    auto User = U->getUser();
+    ++U;
+    if (auto AI = ApplySite::isa(User)) {
+      auto Result = tryDevirtualizeWitnessMethod(AI);
+      if (Result.first) {
+        User->replaceAllUsesWith(Result.first);
+        eraseInstFromFunction(*User);
+      }
+    }
+  }
+  return nullptr;
+}
+

lib/SILPasses/Utils/Devirtualize.cpp

@@ -503,7 +503,7 @@ DevirtualizationResult swift::tryDevirtualizeClassMethod(FullApplySite AI,
 /// Generate a new apply of a function_ref to replace an apply of a
 /// witness_method when we've determined the actual function we'll end
 /// up calling.
-static FullApplySite devirtualizeWitnessMethod(FullApplySite AI, SILFunction *F,
+static ApplySite devirtualizeWitnessMethod(ApplySite AI, SILFunction *F,
                                             ArrayRef<Substitution> Subs) {
   // We know the witness thunk and the corresponding set of substitutions
   // required to invoke the protocol method at this point.
@@ -517,7 +517,10 @@ static FullApplySite devirtualizeWitnessMethod(FullApplySite AI, SILFunction *F,
   SmallVector<Substitution, 16> NewSubstList(Subs.begin(), Subs.end());
 
   // Add the non-self-derived substitutions from the original application.
-  for (auto &origSub : AI.getSubstitutionsWithoutSelfSubstitution())
+  ArrayRef<Substitution>  SubstList;
+  SubstList = AI.getSubstitutionsWithoutSelfSubstitution();
+
+  for (auto &origSub : SubstList)
     if (!origSub.getArchetype()->isSelfDerived())
       NewSubstList.push_back(origSub);
 
@@ -531,18 +534,17 @@ static FullApplySite devirtualizeWitnessMethod(FullApplySite AI, SILFunction *F,
   auto Arguments = SmallVector<SILValue, 4>();
 
   auto ParamTypes = SubstCalleeCanType->getParameterSILTypes();
-  // Type of the current parameter being processed
-  auto ParamType = ParamTypes.begin();
 
   // Iterate over the non self arguments and add them to the
   // new argument list, upcasting when required.
   SILBuilderWithScope B(AI.getInstruction());
-  for (SILValue A : AI.getArguments()) {
-    if (A.getType() != *ParamType)
-      A = B.createUpcast(AI.getLoc(), A, *ParamType);
+  for (unsigned ArgN = 0, ArgE = AI.getNumArguments(); ArgN != ArgE; ++ArgN) {
+	SILValue A = AI.getArgument(ArgN);
+	auto ParamType = ParamTypes[ParamTypes.size() - AI.getNumArguments() + ArgN];
+    if (A.getType() != ParamType)
+      A = B.createUpcast(AI.getLoc(), A, ParamType);
 
     Arguments.push_back(A);
-    ++ParamType;
   }
 
   // Replace old apply instruction by a new apply instruction that invokes
@@ -553,7 +555,7 @@ static FullApplySite devirtualizeWitnessMethod(FullApplySite AI, SILFunction *F,
 
   auto SubstCalleeSILType = SILType::getPrimitiveObjectType(SubstCalleeCanType);
   auto ResultSILType = SubstCalleeCanType->getSILResult();
-  FullApplySite SAI;
+  ApplySite SAI;
 
   if (auto *A = dyn_cast<ApplyInst>(AI))
     SAI = Builder.createApply(Loc, FRI, SubstCalleeSILType,
@@ -563,6 +565,9 @@ static FullApplySite devirtualizeWitnessMethod(FullApplySite AI, SILFunction *F,
     SAI = Builder.createTryApply(Loc, FRI, SubstCalleeSILType,
                                  NewSubstList, Arguments,
                                  TAI->getNormalBB(), TAI->getErrorBB());
+  if (auto *PAI = dyn_cast<PartialApplyInst>(AI))
+    SAI = Builder.createPartialApply(Loc, FRI, SubstCalleeSILType,
+                                     NewSubstList, Arguments, PAI->getType());
 
   NumWitnessDevirt++;
   return SAI;
@@ -571,7 +576,7 @@ static FullApplySite devirtualizeWitnessMethod(FullApplySite AI, SILFunction *F,
 /// In the cases where we can statically determine the function that
 /// we'll call to, replace an apply of a witness_method with an apply
 /// of a function_ref, returning the new apply.
-static DevirtualizationResult tryDevirtualizeWitnessMethod(FullApplySite AI) {
+DevirtualizationResult swift::tryDevirtualizeWitnessMethod(ApplySite AI) {
   SILFunction *F;
   ArrayRef<Substitution> Subs;
   SILWitnessTable *WT;

test/SILPasses/devirt_static_witness_method.sil

@@ -0,0 +1,93 @@
+// RUN: %target-sil-opt -enable-sil-verify-all %s -sil-combine | FileCheck %s
+sil_stage canonical
+
+import Builtin
+import Swift
+import SwiftShims
+
+protocol CanAdd  {
+  func +(lhs: Self, rhs: Self) -> Self
+}
+
+extension Int64 : CanAdd {
+}
+
+struct S<T> : CanAdd {
+}
+
+func +<T>(lhs: S<T>, rhs: S<T>) -> S<T>
+
+
+sil hidden [transparent] [thunk] @operator_plus_static_non_generic_witness_for_S : $@convention(thin) <T where T : CanAdd> (@out S<T>, @in S<T>, @in S<T>, @thick S<T>.Type) -> () {
+bb0(%0 : $*S<T>, %1 : $*S<T>, %2 : $*S<T>, %3 : $@thick S<T>.Type) :
+  %17 = tuple ()
+  return %17 : $()
+}
+
+sil hidden [transparent] [thunk] @operator_plus_static_non_generic_witness : $@convention(witness_method) (@out Int64, @in Int64, @in Int64, @thick Int64.Type) -> () {
+bb0(%0 : $*Int64, %1 : $*Int64, %2 : $*Int64, %3 : $@thick Int64.Type):
+  %4 = struct_element_addr %1 : $*Int64, #Int64._value
+  %5 = load %4 : $*Builtin.Int64
+  %6 = struct_element_addr %2 : $*Int64, #Int64._value
+  %7 = load %6 : $*Builtin.Int64
+  %8 = integer_literal $Builtin.Int1, -1
+  %9 = builtin "sadd_with_overflow_Int64"(%5 : $Builtin.Int64, %7 : $Builtin.Int64, %8 : $Builtin.Int1) : $(Builtin.Int64, Builtin.Int1)
+  %10 = tuple_extract %9 : $(Builtin.Int64, Builtin.Int1), 0
+  %11 = tuple_extract %9 : $(Builtin.Int64, Builtin.Int1), 1
+  cond_fail %11 : $Builtin.Int1
+  %15 = struct $Int64 (%10 : $Builtin.Int64)
+  store %15 to %0 : $*Int64
+  %17 = tuple ()
+  return %17 : $()
+}
+
+// Test that this partial application of a function reference referring to static non generic witness does not crash the IRGen.
+// Such code may be produced e.g. when users refer to global operators defined on builtin types.
+// CHECK-LABEL: sil hidden @test_partial_apply_of_static_witness
+sil hidden @test_partial_apply_of_static_witness : $@convention(thin) () -> @callee_owned (@out Int64, @in Int64, @in Int64) -> () {
+bb0:
+  %1 = metatype $@thick Int64.Type
+  %2 = function_ref @operator_plus_static_non_generic_witness : $@convention(witness_method) (@out Int64, @in Int64, @in Int64, @thick Int64.Type) -> ()
+  %3 = partial_apply %2(%1) : $@convention(witness_method) (@out Int64, @in Int64, @in Int64, @thick Int64.Type) -> ()
+  return %3 : $@callee_owned (@out Int64, @in Int64, @in Int64) -> ()
+}
+
+// Test that this partial application of witness_method can be devirtualized.
+// CHECK-LABEL: sil hidden @test_devirt_of_partial_apply_of_witness_method
+// CHECK-NOT: witness_method $S<Int64>, #CanAdd."+"!1 
+// CHECK: function_ref @operator_plus_static_non_generic_witness_for_S
+// CHECK: partial_apply
+// CHECK: return
+sil hidden @test_devirt_of_partial_apply_of_witness_method : $@convention(thin) () -> @callee_owned (@out S<Int64>, @in S<Int64>, @in S<Int64>) -> () {
+bb0:
+  %1 = metatype $@thick S<Int64>.Type
+  %2 = witness_method $S<Int64>, #CanAdd."+"!1 : $@convention(witness_method) <τ_0_0 where τ_0_0 : CanAdd> (@out τ_0_0, @in τ_0_0, @in τ_0_0, @thick τ_0_0.Type) -> () 
+  %3 = partial_apply %2<S<Int64>>(%1) : $@convention(witness_method) <τ_0_0 where τ_0_0 : CanAdd> (@out τ_0_0, @in τ_0_0, @in τ_0_0, @thick τ_0_0.Type) -> () 
+  return %3 : $@callee_owned (@out S<Int64>, @in S<Int64>, @in S<Int64>) -> ()
+}
+
+// Test that this application of witness_method can be devirtualized.
+// CHECK-LABEL: sil hidden @test_devirt_of_apply_of_witness_method
+// CHECK-NOT: witness_method $S<Int64>, #CanAdd."+"!1 
+// CHECK: function_ref @operator_plus_static_non_generic_witness_for_S
+// CHECK: apply
+// CHECK: return
+sil hidden @test_devirt_of_apply_of_witness_method : $@convention(thin) (@in S<Int64>) -> S<Int64> {
+bb0(%0 : $*S<Int64>):
+  %1 = alloc_stack $S<Int64>
+  %5 = metatype $@thick S<Int64>.Type
+  %6 = witness_method $S<Int64>, #CanAdd."+"!1 : $@convention(witness_method) <τ_0_0 where τ_0_0 : CanAdd> (@out τ_0_0, @in τ_0_0, @in τ_0_0, @thick τ_0_0.Type) -> () 
+  %7 = apply %6<S<Int64>>(%1#1, %0, %0, %5) : $@convention(witness_method) <τ_0_0 where τ_0_0 : CanAdd> (@out τ_0_0, @in τ_0_0, @in τ_0_0, @thick τ_0_0.Type) -> () 
+  %8 = load %1#1: $*S<Int64>
+  dealloc_stack %1#0 : $*@local_storage S<Int64>
+  return %8 : $S<Int64>
+}
+
+sil_witness_table hidden Int64: CanAdd module Test {
+  method #CanAdd."+"!1: @operator_plus_static_non_generic_witness
+}
+
+sil_witness_table hidden <E where E : CanAdd> S<E>: CanAdd module Test {
+  method #CanAdd."+"!1: @operator_plus_static_non_generic_witness_for_S
+}
+