Devirtualize: Extend SE-0309 bail-out logic to variadic function parameters and Self-rooted type parameters

Author: AnthonyLatsis

lib/SILOptimizer/Utils/Devirtualize.cpp

@@ -1094,9 +1094,10 @@ static bool canDevirtualizeWitnessMethod(ApplySite applySite) {
     return false;
   }
 
-  // FIXME: devirtualizeWitnessMethod below does not support cases with
-  // covariant 'Self' nested inside a collection type,
-  // like '[Self]' or '[* : Self]'.
+  // FIXME: devirtualizeWitnessMethod does not support cases with covariant
+  // 'Self'-rooted type parameters nested inside a collection type, like
+  // '[Self]' or '[* : Self.A]', because it doesn't know how to deal with
+  // associated collection upcasts.
   const Type interfaceTy = wmi->getMember()
                                .getDecl()
                                ->getInterfaceType()
@@ -1107,35 +1108,35 @@ static bool canDevirtualizeWitnessMethod(ApplySite applySite) {
   if (!interfaceTy->hasTypeParameter())
     return true;
 
-  class HasSelfNestedInsideCollection final : public TypeWalker {
-    unsigned CollectionDepth;
+  auto *const selfGP = wmi->getLookupProtocol()->getProtocolSelfType();
+  auto isSelfRootedTypeParameter = [selfGP](Type T) -> bool {
+    if (!T->hasTypeParameter())
+      return false;
 
-  public:
-    Action walkToTypePre(Type T) override {
-      if (!T->hasTypeParameter())
-        return Action::SkipChildren;
+    if (T->isTypeParameter()) {
+      return T->getRootGenericParam()->isEqual(selfGP);
+    }
 
-      if (auto *GP = T->getAs<GenericTypeParamType>()) {
-        // Only 'Self' will have zero depth in the type of a requirement.
-        if (GP->getDepth() == 0 && CollectionDepth)
-          return Action::Stop;
-      }
+    return false;
+  };
 
-      if (T->isArray() || T->isDictionary())
-        ++CollectionDepth;
+  return !interfaceTy.findIf([&](Type T) -> bool {
+    if (!T->hasTypeParameter())
+      return false;
 
-      return Action::Continue;
+    if (T->isArray() || T->isDictionary()) {
+      return T.findIf(isSelfRootedTypeParameter);
     }
 
-    Action walkToTypePost(Type T) override {
-      if (T->isArray() || T->isDictionary())
-        --CollectionDepth;
-
-      return Action::Continue;
+    if (auto *FT = T->getAs<FunctionType>()) {
+      for (const auto &Param : FT->getParams()) {
+        if (Param.isVariadic() && T.findIf(isSelfRootedTypeParameter))
+          return true;
+      }
     }
-  };
 
-  return !interfaceTy.walk(HasSelfNestedInsideCollection());
+    return false;
+  });
 }
 
 /// In the cases where we can statically determine the function that

test/SILOptimizer/devirt_protocol_method_invocations.swift

@@ -290,28 +290,56 @@ public func testReabstracted(f: Optional<()->()>) {
 }
 
 
-// Test that we don't devirtualize calls to protocol requirements with covariant `Self` nested
-// inside a collection type – the devirtualizer does not support handling these yet.
-protocol CovariantSelfInsideCollection {
-  func array() -> Array<Self>
-  func dictionary() -> Dictionary<String, Self>
-  func mixed(_: (Array<(Dictionary<String, String>, Self)>) -> Void)
+// Test that we don't devirtualize calls to protocol requirements with
+// covariant `Self`-rooted type parameters nested inside a collection type;
+// the devirtualizer doesn't know how to handle these yet.
+protocol CovariantSelfInCollection {
+  associatedtype Assoc
+
+  func self1() -> Array<Self>
+  func self2() -> Dictionary<String, Self>
+  func self3(_: (Self...) -> Void)
+  func self4(_: (Array<(Dictionary<String, String>, Self)>) -> Void)
+
+  func assoc1() -> Array<Assoc>
+  func assoc2() -> Dictionary<String, Assoc>
+  func assoc3(_: (Assoc...) -> Void)
+  func assoc4(_: (Array<(Dictionary<String, String>, Assoc)>) -> Void)
+}
+struct CovariantSelfInCollectionImpl: CovariantSelfInCollection {
+  typealias Assoc = Bool
+
+  func self1() -> Array<Self> { [self] }
+  func self2() -> Dictionary<String, Self> { [#file : self] }
+  func self3(_: (Self...) -> Void) {}
+  func self4(_: (Array<(Dictionary<String, String>, Self)>) -> Void) {}
+
+  func assoc1() -> Array<Assoc> { [true] }
+  func assoc2() -> Dictionary<String, Assoc> { [#file : true] }
+  func assoc3(_: (Assoc...) -> Void) {}
+  func assoc4(_: (Array<(Dictionary<String, String>, Assoc)>) -> Void) {}
 }
 // CHECK-LABEL: sil @$s34devirt_protocol_method_invocations12testNoDevirtyyF
 //
-// CHECK: witness_method $S, #CovariantSelfInsideCollection.array
-// CHECK: witness_method $S, #CovariantSelfInsideCollection.dictionary
-// CHECK: witness_method $S, #CovariantSelfInsideCollection.mixed
+// CHECK: witness_method $CovariantSelfInCollectionImpl, #CovariantSelfInCollection.self1
+// CHECK: witness_method $CovariantSelfInCollectionImpl, #CovariantSelfInCollection.self2
+// CHECK: witness_method $CovariantSelfInCollectionImpl, #CovariantSelfInCollection.self3
+// CHECK: witness_method $CovariantSelfInCollectionImpl, #CovariantSelfInCollection.self4
+// CHECK: witness_method $CovariantSelfInCollectionImpl, #CovariantSelfInCollection.assoc1
+// CHECK: witness_method $CovariantSelfInCollectionImpl, #CovariantSelfInCollection.assoc2
+// CHECK: witness_method $CovariantSelfInCollectionImpl, #CovariantSelfInCollection.assoc3
+// CHECK: witness_method $CovariantSelfInCollectionImpl, #CovariantSelfInCollection.assoc4
 // CHECK: end sil function '$s34devirt_protocol_method_invocations12testNoDevirtyyF'
 public func testNoDevirt() {
-  struct S: CovariantSelfInsideCollection {
-    func array() -> Array<Self> { fatalError() }
-    func dictionary() -> Dictionary<String, Self> { fatalError() }
-    func mixed(_: (Array<(Dictionary<String, String>, Self)>) -> Void) {}
-  }
+  let p: any CovariantSelfInCollection = CovariantSelfInCollectionImpl()
+
+  _ = p.self1()
+  _ = p.self2()
+  p.self3 { _ in }
+  p.self4 { _ in }
 
-  let p: CovariantSelfInsideCollection = S()
-  _ = p.array()
-  _ = p.dictionary()
-  p.mixed { _ in }
+  _ = p.assoc1()
+  _ = p.assoc2()
+  p.assoc3 { _ in }
+  p.assoc4 { _ in }
 }