Make sure we can dig out the (distributed) actor from an archetype.

Fixes #59356 / rdar://94976378.

(cherry picked from commit 774a94f)

Author: DougGregor

include/swift/AST/Types.h

@@ -786,6 +786,10 @@ class alignas(1 << TypeAlignInBits) TypeBase
   /// Break an existential down into a set of constraints.
   ExistentialLayout getExistentialLayout();
 
+  /// If this is an actor or distributed type, get the nominal type declaration
+  /// for the actor.
+  NominalTypeDecl *getAnyActor();
+
   /// Determines whether this type is an actor type.
   bool isActorType();
 

lib/AST/Decl.cpp

@@ -9152,12 +9152,11 @@ ActorIsolation swift::getActorIsolationOfContext(DeclContext *dc) {
 
     case ClosureActorIsolation::ActorInstance: {
       auto selfDecl = isolation.getActorInstance();
-      auto actorClass = selfDecl->getType()->getReferenceStorageReferent()
-          ->getClassOrBoundGenericClass();
-      // FIXME: Doesn't work properly with generics #59356
-      assert(actorClass && "Bad closure actor isolation?");
-      return ActorIsolation::forActorInstance(actorClass)
-                .withPreconcurrency(isolation.preconcurrency());
+      auto actor = selfDecl->getType()->getReferenceStorageReferent()
+          ->getAnyActor();
+      assert(actor && "Bad closure actor isolation?");
+      return ActorIsolation::forActorInstance(actor)
+        .withPreconcurrency(isolation.preconcurrency());
     }
     }
   }

lib/AST/Type.cpp

@@ -391,84 +391,51 @@ bool CanType::isTypeErasedGenericClassTypeImpl(CanType type) {
   return false;
 }
 
-static bool archetypeConformsTo(ArchetypeType *archetype, KnownProtocolKind protocol) {
-  auto &ctx = archetype->getASTContext();
-  auto expectedProto = ctx.getProtocol(protocol);
-  if (!expectedProto)
-    return false;
-
-  for (auto proto : archetype->getConformsTo()) {
-    if (proto == expectedProto || proto->inheritsFrom(expectedProto))
-      return true;
-  }
-
-  return false;
-}
-
-bool TypeBase::isActorType() {
+NominalTypeDecl *TypeBase::getAnyActor() {
   // Nominal types: check whether the declaration is an actor.
-  if (auto nominal = getAnyNominal())
-    return nominal->isActor();
+  if (auto nominal = getAnyNominal()) {
+    if (nominal->isAnyActor())
+      return nominal;
+  }
 
   // Archetypes check for conformance to Actor.
   if (auto archetype = getAs<ArchetypeType>()) {
-    return archetypeConformsTo(archetype, KnownProtocolKind::Actor);
+    for (auto proto : archetype->getConformsTo()) {
+      if (proto->isAnyActor())
+        return proto;
+    }
+
+    return nullptr;
   }
 
   // Existential types: check for Actor protocol.
   if (isExistentialType()) {
-    auto actorProto = getASTContext().getProtocol(KnownProtocolKind::Actor);
-    if (!actorProto)
-      return false;
-
     auto layout = getExistentialLayout();
     if (auto superclass = layout.getSuperclass()) {
-      if (superclass->isActorType())
-        return true;
+      if (auto actor = superclass->getAnyActor())
+        return actor;
     }
 
     for (auto proto : layout.getProtocols()) {
-      if (proto->isActor())
-        return true;
+      if (proto->isAnyActor())
+        return proto;
     }
 
-    return false;
+    return nullptr;
   }
 
+  return nullptr;
+}
+
+bool TypeBase::isActorType() {
+  if (auto actor = getAnyActor())
+    return actor->isActor();
   return false;
 }
 
 bool TypeBase::isDistributedActor() {
-  // Nominal types: check whether the declaration is an actor.
-  if (auto *nominal = getAnyNominal()) {
-    if (auto *classDecl = dyn_cast<ClassDecl>(nominal))
-      return classDecl->isDistributedActor();
-
-    if (isa<StructDecl>(nominal) || isa<EnumDecl>(nominal))
-      return false;
-  }
-
-  // Archetypes check for conformance to DistributedActor.
-  if (auto archetype = getAs<ArchetypeType>()) {
-    return archetypeConformsTo(archetype, KnownProtocolKind::DistributedActor);
-  }
-
-  // Existential types: check for DistributedActor protocol conformance.
-  if (isExistentialType()) {
-    auto actorProto = getASTContext().getDistributedActorDecl();
-    if (!actorProto)
-      return false;
-
-    // TODO(distributed): Inheritance is not yet supported.
-
-    auto layout = getExistentialLayout();
-    return llvm::any_of(layout.getProtocols(),
-                        [&actorProto](ProtocolDecl *protocol) {
-                          return protocol == actorProto ||
-                                 protocol->inheritsFrom(actorProto);
-                        });
-  }
-
+  if (auto actor = getAnyActor())
+    return actor->isDistributedActor();
   return false;
 }
 

test/Distributed/distributed_protocol_isolation.swift

@@ -224,11 +224,10 @@ func test_watchingDA<WDA: TerminationWatchingDA>(da: WDA) async throws {
   // expected-warning@-2{{no calls to throwing functions occur within 'try' expression}}
 
   let __secretlyKnownToBeLocal = da
-  await __secretlyKnownToBeLocal.terminated(da: "local calls are okey!") // OK // FIXME(#59356): (the __secretlyKnown is a hack, but the whenLocal crashes now on pending isolation getting with generic actors for closures)
-  // FIXME: pending fix of closure isolation checking with actors #59356
-  // await da.whenLocal { __secretlyKnownToBeLocal in
-  //   await __secretlyKnownToBeLocal.terminated(da: "local calls are okey!") // OK
-  // }
+  await __secretlyKnownToBeLocal.terminated(da: "local calls are okey!") // OK
+  await da.whenLocal { __secretlyKnownToBeLocal in
+    await __secretlyKnownToBeLocal.terminated(da: "local calls are okey!") // OK
+  }
 }
 
 func test_watchingDA_erased(da: DA_TerminationWatchingDA) async throws {
@@ -238,11 +237,10 @@ func test_watchingDA_erased(da: DA_TerminationWatchingDA) async throws {
   // expected-warning@-2{{no calls to throwing functions occur within 'try' expression}}
 
   let __secretlyKnownToBeLocal = wda
-  await __secretlyKnownToBeLocal.terminated(da: "local calls are okey!") // OK // FIXME(#59356): (the __secretlyKnown is a hack, but the whenLocal crashes now on pending isolation getting with generic actors for closures)
-  // FIXME: pending fix of closure isolation checking with actors #59356
-  // await wda.whenLocal { __secretlyKnownToBeLocal in
-  //   await __secretlyKnownToBeLocal.terminated(da: "local calls are okey!") // OK
-  // }
+  await __secretlyKnownToBeLocal.terminated(da: "local calls are okey!") // OK
+  await wda.whenLocal { __secretlyKnownToBeLocal in
+    await __secretlyKnownToBeLocal.terminated(da: "local calls are okey!") // OK
+  }
 }
 
 func test_watchingDA_any(da: any TerminationWatchingDA) async throws {

validation-test/compiler_crashers_2_fixed/rdar94976378.swift

@@ -0,0 +1,14 @@
+// RUN: %target-swift-frontend -typecheck %s
+
+// REQUIRES: concurrency
+
+@available(SwiftStdlib 5.1, *)
+extension Actor {
+  func f() { }
+
+  func g(a: [Int]) {
+    a.forEach { i in
+      f()
+    }
+  }
+}