Devirtualize: Skip requirements with covariant 'Self' nested inside a collection

The devirtualizer does not support handling these yet, which wasn't anticipated in #34140

Author: AnthonyLatsis

lib/SILOptimizer/Utils/Devirtualize.cpp

@@ -1094,7 +1094,48 @@ static bool canDevirtualizeWitnessMethod(ApplySite applySite) {
     return false;
   }
 
-  return true;
+  // FIXME: devirtualizeWitnessMethod below does not support cases with
+  // covariant 'Self' nested inside a collection type,
+  // like '[Self]' or '[* : Self]'.
+  const Type interfaceTy = wmi->getMember()
+                               .getDecl()
+                               ->getInterfaceType()
+                               // Skip the 'self' parameter.
+                               ->castTo<AnyFunctionType>()
+                               ->getResult();
+
+  if (!interfaceTy->hasTypeParameter())
+    return true;
+
+  class HasSelfNestedInsideCollection final : public TypeWalker {
+    unsigned CollectionDepth;
+
+  public:
+    Action walkToTypePre(Type T) override {
+      if (!T->hasTypeParameter())
+        return Action::SkipChildren;
+
+      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;
+      }
+
+      if (T->isArray() || T->isDictionary())
+        ++CollectionDepth;
+
+      return Action::Continue;
+    }
+
+    Action walkToTypePost(Type T) override {
+      if (T->isArray() || T->isDictionary())
+        --CollectionDepth;
+
+      return Action::Continue;
+    }
+  };
+
+  return !interfaceTy.walk(HasSelfNestedInsideCollection());
 }
 
 /// In the cases where we can statically determine the function that

test/SILOptimizer/devirt_protocol_method_invocations.swift

@@ -288,3 +288,30 @@ func testReabstractedWitness(_ f: ReabstractedP) {
 public func testReabstracted(f: Optional<()->()>) {
   testReabstractedWitness(f)
 }
+
+
+// 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)
+}
+// 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: 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: CovariantSelfInsideCollection = S()
+  _ = p.array()
+  _ = p.dictionary()
+  p.mixed { _ in }
+}

test/decl/protocol/protocols_with_self_or_assoc_reqs_executable.swift

@@ -48,10 +48,7 @@ Tests.test("Basic") {
   expectEqual(3, collection.count)
 }
 
-// FIXME: Teach the devirtualizer how to handle calls to requirements with covariant `Self` nested
-// inside known-covariant stdlib types such as an array or dictionary.
-@_optimize(none)
-func convariantSelfErasureTest() {
+Tests.test("Covariant 'Self' erasure") {
   struct S: P {
     static let str = "Success"
     func getString() -> String { Self.str }
@@ -107,7 +104,4 @@ func convariantSelfErasureTest() {
   expectEqual(true, S() is P)
 }
 
-
-Tests.test("Covariant 'Self' erasure", convariantSelfErasureTest)
-
 runAllTests()