[WIP] SIL: Add devirtualizer support for default witness methods

lib/SILOptimizer/Utils/Devirtualize.cpp

@@ -679,29 +679,42 @@ DevirtualizationResult swift::tryDevirtualizeClassMethod(FullApplySite AI,
 //                        Witness Method Optimization
 //===----------------------------------------------------------------------===//
 
-/// 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 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.
+static void getWitnessMethodSubstitutions(ApplySite AI, SILFunction *F,
+                                          ArrayRef<Substitution> Subs,
+                                          SmallVectorImpl<Substitution> &NewSubs) {
   auto &Module = AI.getModule();
 
-  // Collect all the required substitutions.
-  //
-  // The complete set of substitutions may be different, e.g. because the found
-  // witness thunk F may have been created by a specialization pass and have
-  // additional generic parameters.
-  SmallVector<Substitution, 16> NewSubstList(Subs.begin(), Subs.end());
+  auto CalleeCanType = F->getLoweredFunctionType();
+
+  ProtocolDecl *proto = nullptr;
+  if (CalleeCanType->getRepresentation() ==
+      SILFunctionTypeRepresentation::WitnessMethod) {
+    proto = CalleeCanType->getDefaultWitnessMethodProtocol(
+        *Module.getSwiftModule());
+  }
+
+  ArrayRef<Substitution> origSubs = AI.getSubstitutions();
+
+  if (proto != nullptr) {
+    // If the callee is a default witness method thunk, preserve substitutions
+    // from the call site.
+    NewSubs.append(origSubs.begin(), origSubs.end());
+    return;
+  }
+
+  // If the callee is a concrete witness method thunk, apply substitutions
+  // from the conformance, and drop any substitutions derived from the Self
+  // type.
+  NewSubs.append(Subs.begin(), Subs.end());
+
   if (auto generics = AI.getOrigCalleeType()->getGenericSignature()) {
-    ArrayRef<Substitution> origSubs = AI.getSubstitutions();
     for (auto genericParam : generics->getAllDependentTypes()) {
       auto origSub = origSubs.front();
       origSubs = origSubs.slice(1);
 
-      // Ignore generic parameters derived from 'self', the generic
-      // parameter at depth 0, index 0.
+      // If the callee is a concrete witness method thunk, we ignore
+      // generic parameters derived from 'self', the generic parameter at
+      // depth 0, index 0.
       auto type = genericParam->getCanonicalType();
       while (auto memberType = dyn_cast<DependentMemberType>(type)) {
         type = memberType.getBase();
@@ -714,17 +727,36 @@ static ApplySite devirtualizeWitnessMethod(ApplySite AI, SILFunction *F,
       }
 
       // Okay, remember this substitution.
-      NewSubstList.push_back(origSub);
+      NewSubs.push_back(origSub);
     }
-
-    assert(origSubs.empty() && "subs not parallel to dependent types");
   }
 
+  assert(origSubs.empty() && "subs not parallel to dependent types");
+}
+
+/// 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 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.
+  auto &Module = AI.getModule();
+
+  // Collect all the required substitutions.
+  //
+  // The complete set of substitutions may be different, e.g. because the found
+  // witness thunk F may have been created by a specialization pass and have
+  // additional generic parameters.
+  SmallVector<Substitution, 4> NewSubs;
+
+  getWitnessMethodSubstitutions(AI, F, Subs, NewSubs);
+
   // Figure out the exact bound type of the function to be called by
   // applying all substitutions.
   auto CalleeCanType = F->getLoweredFunctionType();
   auto SubstCalleeCanType = CalleeCanType->substGenericArgs(
-    Module, Module.getSwiftModule(), NewSubstList);
+    Module, Module.getSwiftModule(), NewSubs);
 
   // Collect arguments from the apply instruction.
   auto Arguments = SmallVector<SILValue, 4>();
@@ -754,15 +786,15 @@ static ApplySite devirtualizeWitnessMethod(ApplySite AI, SILFunction *F,
 
   if (auto *A = dyn_cast<ApplyInst>(AI))
     SAI = Builder.createApply(Loc, FRI, SubstCalleeSILType,
-                              ResultSILType, NewSubstList, Arguments,
+                              ResultSILType, NewSubs, Arguments,
                               A->isNonThrowing());
   if (auto *TAI = dyn_cast<TryApplyInst>(AI))
     SAI = Builder.createTryApply(Loc, FRI, SubstCalleeSILType,
-                                 NewSubstList, Arguments,
+                                 NewSubs, Arguments,
                                  TAI->getNormalBB(), TAI->getErrorBB());
   if (auto *PAI = dyn_cast<PartialApplyInst>(AI))
     SAI = Builder.createPartialApply(Loc, FRI, SubstCalleeSILType,
-                                     NewSubstList, Arguments, PAI->getType());
+                                     NewSubs, Arguments, PAI->getType());
 
   NumWitnessDevirt++;
   return SAI;

test/SILOptimizer/devirt_default_witness_method.sil

@@ -0,0 +1,62 @@
+// RUN: %target-sil-opt -enable-sil-verify-all %s -sil-combine -enable-resilience | FileCheck %s
+sil_stage canonical
+
+import Builtin
+import Swift
+import SwiftShims
+
+public protocol ResilientProtocol {
+  func defaultA()
+}
+
+sil @defaultA : $@convention(witness_method) <Self where Self : ResilientProtocol> (@in_guaranteed Self) -> () {
+bb0(%0 : $*Self):
+  %result = tuple ()
+  return %result : $()
+}
+
+sil_default_witness_table ResilientProtocol 1 {
+  method #ResilientProtocol.defaultA!1: @defaultA
+}
+
+struct ConformingStruct : ResilientProtocol {
+  func defaultA()
+}
+
+sil_witness_table ConformingStruct : ResilientProtocol module protocol_resilience {
+  method #ResilientProtocol.defaultA!1: @defaultA
+}
+
+struct ConformingGenericStruct<T> : ResilientProtocol {
+  func defaultA()
+}
+
+sil_witness_table <T> ConformingGenericStruct<T> : ResilientProtocol module protocol_resilience {
+  method #ResilientProtocol.defaultA!1: @defaultA
+}
+
+// CHECK-LABEL: sil hidden @test_devirt_of_default_witness_method : $@convention(thin) (@in_guaranteed ConformingStruct) -> ()
+// CHECK:       bb0(%0 : $*ConformingStruct):
+// CHECK:         [[FN:%.*]] = function_ref @defaultA : $@convention(witness_method) <τ_0_0 where τ_0_0 : ResilientProtocol> (@in_guaranteed τ_0_0) -> ()
+// CHECK-NEXT:    [[RESULT:%.*]] = apply [[FN]]<ConformingStruct>(%0) : $@convention(witness_method) <τ_0_0 where τ_0_0 : ResilientProtocol> (@in_guaranteed τ_0_0)
+// CHECK-NEXT:    return [[RESULT]] : $()
+
+sil hidden @test_devirt_of_default_witness_method : $@convention(thin) (@in_guaranteed ConformingStruct) -> () {
+bb0(%0 : $*ConformingStruct):
+  %fn = witness_method $ConformingStruct, #ResilientProtocol.defaultA!1 : $@convention(witness_method) <T where T : ResilientProtocol> (@in_guaranteed T) -> ()
+  %result = apply %fn<ConformingStruct>(%0) : $@convention(witness_method) <T where T : ResilientProtocol> (@in_guaranteed T) -> ()
+  return %result : $()
+}
+
+// CHECK-LABEL: sil hidden @test_devirt_of_generic_default_witness_method : $@convention(thin) (@in_guaranteed ConformingGenericStruct<Int>) -> ()
+// CHECK:       bb0(%0 : $*ConformingGenericStruct<Int>):
+// CHECK:         [[FN:%.*]] = function_ref @defaultA : $@convention(witness_method) <τ_0_0 where τ_0_0 : ResilientProtocol> (@in_guaranteed τ_0_0) -> ()
+// CHECK-NEXT:    [[RESULT:%.*]] = apply [[FN]]<ConformingGenericStruct<Int>>(%0) : $@convention(witness_method) <τ_0_0 where τ_0_0 : ResilientProtocol> (@in_guaranteed τ_0_0)
+// CHECK-NEXT:    return [[RESULT]] : $()
+
+sil hidden @test_devirt_of_generic_default_witness_method : $@convention(thin) (@in_guaranteed ConformingGenericStruct<Int>) -> () {
+bb0(%0 : $*ConformingGenericStruct<Int>):
+  %fn = witness_method $ConformingGenericStruct<Int>, #ResilientProtocol.defaultA!1 : $@convention(witness_method) <T where T : ResilientProtocol> (@in_guaranteed T) -> ()
+  %result = apply %fn<ConformingGenericStruct<Int>>(%0) : $@convention(witness_method) <T where T : ResilientProtocol> (@in_guaranteed T) -> ()
+  return %result : $()
+}