Merge pull request #5571 from DougGregor/devirtualize-default-generic-witness

[Devirtualizer] Handle default witnesses for generic requirements.

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 : $()
+}