[Devirtualizer] Handle default witnesses for generic requirements.

The witness thunks for default witnesses are different from the
witness thunks for normal witnesses, because default witnesses take
'Self' (the whole conforming type) rather than having it substituted
away. Cope with this difference while still substituting the innermost
generic parameters for a generic requirement.

Author: DougGregor

lib/SILOptimizer/Utils/Devirtualize.cpp

@@ -791,6 +791,7 @@ static void getWitnessMethodSubstitutions(
     GenericSignature *requirementSig,
     GenericSignature *witnessThunkSig,
     ArrayRef<Substitution> origSubs,
+    bool isDefaultWitness,
     SmallVectorImpl<Substitution> &newSubs) {
 
   if (witnessThunkSig == nullptr)
@@ -806,19 +807,28 @@ static void getWitnessMethodSubstitutions(
   SubstitutionMap subMap;
 
   // Take apart substitutions from the conforming type.
-  //
-  // If `Self` maps to a bound generic type, this gives us the
-  // substitutions for the concrete type's generic parameters.
-  auto witnessSubs = getSubstitutionsForProtocolConformance(conformanceRef);
-
+  ArrayRef<Substitution> witnessSubs;
+  auto *rootConformance = conformance->getRootNormalConformance();
+  auto *witnessSig = rootConformance->getGenericSignature();
   unsigned depth = 0;
-  if (!witnessSubs.empty()) {
-    auto *rootConformance = conformance->getRootNormalConformance();
-    depth = rootConformance->getGenericSignature()->getGenericParams().back()
-              ->getDepth() + 1;
-    auto *witnessSig = rootConformance->getGenericSignature();
+  if (isDefaultWitness) {
+    // For default witnesses, we substitute all of Self.
+    auto gp = witnessThunkSig->getGenericParams().front()->getCanonicalType();
+    subMap.addSubstitution(gp, origSubs.front().getReplacement());
+    subMap.addConformances(gp, origSubs.front().getConformances());
+
+    // For default witnesses, innermost generic parameters are always at
+    // depth 1.
+    depth = 1;
+  } else {
+    // If `Self` maps to a bound generic type, this gives us the
+    // substitutions for the concrete type's generic parameters.
+    witnessSubs = getSubstitutionsForProtocolConformance(conformanceRef);
 
-    witnessSig->getSubstitutionMap(witnessSubs, subMap);
+    if (!witnessSubs.empty()) {
+      witnessSig->getSubstitutionMap(witnessSubs, subMap);
+      depth = witnessSig->getGenericParams().back()->getDepth() + 1;
+    }
   }
 
   // Next, take apart caller-side substitutions.
@@ -890,17 +900,15 @@ static void getWitnessMethodSubstitutions(ApplySite AI, SILFunction *F,
 
   ArrayRef<Substitution> origSubs = AI.getSubstitutions();
 
-  if (F->getLoweredFunctionType()->getRepresentation()
-          == SILFunctionTypeRepresentation::WitnessMethod &&
-      F->getLoweredFunctionType()->getDefaultWitnessMethodProtocol(
-          *Module.getSwiftModule())) {
-    // Default witness thunks use the generic signature of the requirement.
-    NewSubs.append(origSubs.begin(), origSubs.end());
-    return;
-  }
+  bool isDefaultWitness =
+    F->getLoweredFunctionType()->getRepresentation()
+      == SILFunctionTypeRepresentation::WitnessMethod &&
+    F->getLoweredFunctionType()->getDefaultWitnessMethodProtocol(
+                                                     *Module.getSwiftModule())
+      == CRef.getRequirement();
 
   getWitnessMethodSubstitutions(Module, CRef, requirementSig, witnessThunkSig,
-                                origSubs, NewSubs);
+                                origSubs, isDefaultWitness, NewSubs);
 }
 
 /// Check if an upcast is legal.

test/SILOptimizer/devirt_default_witness_method.sil

@@ -5,6 +5,10 @@ import Builtin
 import Swift
 import SwiftShims
 
+public protocol P { }
+public protocol Q : P { }
+extension Int : P, Q { }
+
 public protocol ResilientProtocol {
   func defaultA()
 }
@@ -35,6 +39,28 @@ sil_witness_table <T> ConformingGenericStruct<T> : ResilientProtocol module prot
   method #ResilientProtocol.defaultA!1: @defaultA
 }
 
+public protocol ResilientProtocolWithGeneric {
+  func defaultB<T: Q>(_: T)
+}
+
+sil @defaultB : $@convention(witness_method) <Self where Self : ResilientProtocolWithGeneric><T where T : P> (@in T, @in_guaranteed Self) -> () {
+bb0(%0 : $*T, %1 : $*Self):
+  %result = tuple ()
+  return %result : $()
+}
+
+sil_default_witness_table ResilientProtocolWithGeneric {
+  method #ResilientProtocolWithGeneric.defaultB!1: @defaultB
+}
+
+extension ConformingGenericStruct : ResilientProtocolWithGeneric {
+  func defaultB<T>(_: T)
+}
+
+sil_witness_table <T> ConformingGenericStruct<T> : ResilientProtocolWithGeneric module protocol_resilience {
+  method #ResilientProtocolWithGeneric.defaultB!1: @defaultB
+}
+
 // 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) -> ()
@@ -60,3 +86,16 @@ bb0(%0 : $*ConformingGenericStruct<Int>):
   %result = apply %fn<ConformingGenericStruct<Int>>(%0) : $@convention(witness_method) <T where T : ResilientProtocol> (@in_guaranteed T) -> ()
   return %result : $()
 }
+
+// CHECK-LABEL: test_devirt_of_inner_generic_default_witness_method
+// CHECK:       bb0(%0 : $*ConformingGenericStruct<Int>, %1 : $*Int):
+// CHECK:       [[FN:%[0-9]+]] = function_ref @defaultB : $@convention(witness_method) <τ_0_0 where τ_0_0 : ResilientProtocolWithGeneric><τ_1_0 where τ_1_0 : P> (@in τ_1_0, @in_guaranteed τ_0_0) -> ()
+// CHECK:       [[RESULT:%[0-9]+]] = apply [[FN]]<ConformingGenericStruct<Int>, Int>(%1, %0) : $@convention(witness_method) <τ_0_0 where τ_0_0 : ResilientProtocolWithGeneric><τ_1_0 where τ_1_0 : P> (@in τ_1_0, @in_guaranteed τ_0_0) -> ()
+// CHECK:       return [[RESULT]] : $()
+// CHECK:     }
+sil hidden @test_devirt_of_inner_generic_default_witness_method : $@convention(thin) (@in_guaranteed ConformingGenericStruct<Int>, @in Int) -> () {
+bb0(%0 : $*ConformingGenericStruct<Int>, %1 : $*Int):
+  %fn = witness_method $ConformingGenericStruct<Int>, #ResilientProtocolWithGeneric.defaultB!1 : $@convention(witness_method) <T where T : ResilientProtocolWithGeneric><U where U : Q> (@in U, @in_guaranteed T) -> ()
+  %result = apply %fn<ConformingGenericStruct<Int>, Int>(%1, %0) : $@convention(witness_method) <T where T : ResilientProtocolWithGeneric><U where U : Q> (@in U, @in_guaranteed T) -> ()
+  return %result : $()
+}