Handle flow-sensitive #isolation in distributed actor initializers.

Distributed actors can be treated as actors by accessing the `asLocalActor`
property. When lowering `#isolation` in a distributed actor initializer,
use a separate builtin `flowSensitiveDistributedSelfIsolation` to
capture the conformance to `DistributedActor`, and have Definite
Initialization introduce the call to the `asLocalActor` getter when
needed.

Author: DougGregor

include/swift/AST/ASTSynthesis.h

@@ -56,6 +56,7 @@ enum SingletonTypeSynthesizer {
   _serialExecutor, // the '_Concurrency.SerialExecutor' protocol
   _taskExecutor,   // the '_Concurrency.TaskExecutor' protocol
   _actor,          // the '_Concurrency.Actor' protocol
+  _distributedActor,  // the 'Distributed.DistributedActor' protocol
 };
 inline Type synthesizeType(SynthesisContext &SC,
                            SingletonTypeSynthesizer kind) {
@@ -82,6 +83,9 @@ inline Type synthesizeType(SynthesisContext &SC,
   case _actor:
     return SC.Context.getProtocol(KnownProtocolKind::Actor)
       ->getDeclaredInterfaceType();
+  case _distributedActor:
+    return SC.Context.getProtocol(KnownProtocolKind::DistributedActor)
+      ->getDeclaredInterfaceType();
   case _copyable:
     return SC.Context.getProtocol(KnownProtocolKind::Copyable)
         ->getDeclaredInterfaceType();

include/swift/AST/Builtins.def

@@ -899,14 +899,25 @@ BUILTIN_MISC_OPERATION(StartAsyncLetWithLocalBuffer, "startAsyncLetWithLocalBuff
 /// This is only supported under the task-to-thread concurrency model.
 BUILTIN_MISC_OPERATION(TaskRunInline, "taskRunInline", "", Special)
 
-/// flowSensitiveSelfIsolation<T: Actor>(_ actor: T) -> T?
+/// flowSensitiveSelfIsolation<T: Actor>(_ actor: T) -> (any Actor)?
 ///
 /// Used only in actor initializers, this builtin lowers to either 'actor'
 /// (wrapped in an optional) or 'nil' depending on whether 'self' has been
 /// initialized at this point. 'actor' is always an alias for the 'self'
 /// being initialized.
 BUILTIN_MISC_OPERATION(FlowSensitiveSelfIsolation, "flowSensitiveSelfIsolation", "", Special)
 
+/// flowSensitiveDistributedSelfIsolation<T: DistributedActor>(
+///   _ actor: T
+/// ) -> (any Actor)?
+///
+/// Used only in distributed actor initializers, this builtin lowers to either
+/// 'actor.asLocalActor' or 'nil' depending on whether 'self' has been
+/// initialized at this point. 'actor' is always an alias for the 'self'
+/// being initialized.
+BUILTIN_MISC_OPERATION(FlowSensitiveDistributedSelfIsolation,
+                       "flowSensitiveDistributedSelfIsolation", "", Special)
+
 /// endAsyncLet(): (Builtin.RawPointer) -> Void
 ///
 /// DEPRECATED. The swift_asyncLet_finish intrinsic and endAsyncLetLifetime

lib/AST/Builtins.cpp

@@ -2089,14 +2089,18 @@ static ValueDecl *getHopToActor(ASTContext &ctx, Identifier id) {
   return builder.build(id);
 }
 
-static ValueDecl *getFlowSensitiveSelfIsolation(ASTContext &ctx, Identifier id) {
+static ValueDecl *getFlowSensitiveSelfIsolation(
+  ASTContext &ctx, Identifier id, bool isDistributed
+) {
   BuiltinFunctionBuilder builder(ctx);
-  return getBuiltinFunction(ctx, id, _thin,
-                            _generics(_unrestricted,
-                                      _conformsToDefaults(0),
-                                      _conformsTo(_typeparam(0), _actor)),
-                            _parameters(_typeparam(0)),
-                            _optional(_existential(_actor)));
+  return getBuiltinFunction(
+      ctx, id, _thin,
+      _generics(_unrestricted,
+                _conformsToDefaults(0),
+                _conformsTo(_typeparam(0),
+                            isDistributed ? _distributedActor : _actor)),
+      _parameters(_typeparam(0)),
+      _optional(_existential(_actor)));
 }
 
 static ValueDecl *getDistributedActorAsAnyActor(ASTContext &ctx, Identifier id) {
@@ -3202,7 +3206,10 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
     return getHopToActor(Context, Id);
 
   case BuiltinValueKind::FlowSensitiveSelfIsolation:
-    return getFlowSensitiveSelfIsolation(Context, Id);
+    return getFlowSensitiveSelfIsolation(Context, Id, false);
+
+  case BuiltinValueKind::FlowSensitiveDistributedSelfIsolation:
+    return getFlowSensitiveSelfIsolation(Context, Id, true);
 
   case BuiltinValueKind::AutoDiffCreateLinearMapContextWithType:
     return getAutoDiffCreateLinearMapContext(Context, Id);

lib/SIL/IR/OperandOwnership.cpp

@@ -907,6 +907,7 @@ BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, CreateTaskGroup)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, CreateTaskGroupWithFlags)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, DestroyTaskGroup)
 BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, FlowSensitiveSelfIsolation)
+BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, FlowSensitiveDistributedSelfIsolation)
 
 BUILTIN_OPERAND_OWNERSHIP(ForwardingConsume, COWBufferForReading)
 

lib/SIL/IR/ValueOwnership.cpp

@@ -631,6 +631,7 @@ CONSTANT_OWNERSHIP_BUILTIN(None, CreateTaskGroupWithFlags)
 CONSTANT_OWNERSHIP_BUILTIN(None, DestroyTaskGroup)
 CONSTANT_OWNERSHIP_BUILTIN(None, TaskRunInline)
 CONSTANT_OWNERSHIP_BUILTIN(Owned, FlowSensitiveSelfIsolation)
+CONSTANT_OWNERSHIP_BUILTIN(Owned, FlowSensitiveDistributedSelfIsolation)
 CONSTANT_OWNERSHIP_BUILTIN(None, GetEnumTag)
 CONSTANT_OWNERSHIP_BUILTIN(None, InjectEnumTag)
 CONSTANT_OWNERSHIP_BUILTIN(Owned, DistributedActorAsAnyActor)

lib/SILGen/SILGenExpr.cpp

@@ -31,6 +31,7 @@
 #include "swift/AST/CanTypeVisitor.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/DiagnosticsCommon.h"
+#include "swift/AST/DistributedDecl.h"
 #include "swift/AST/ExistentialLayout.h"
 #include "swift/AST/Expr.h"
 #include "swift/AST/ForeignErrorConvention.h"
@@ -6755,23 +6756,41 @@ RValue RValueEmitter::visitCurrentContextIsolationExpr(
         isolation.getActorInstance() == ctor->getImplicitSelfDecl()) {
       ASTContext &ctx = SGF.getASTContext();
       auto builtinName = ctx.getIdentifier(
-          getBuiltinName(BuiltinValueKind::FlowSensitiveSelfIsolation));
+          isolation.isDistributedActor()
+              ? getBuiltinName(BuiltinValueKind::FlowSensitiveDistributedSelfIsolation)
+              : getBuiltinName(BuiltinValueKind::FlowSensitiveSelfIsolation));
       SILType resultTy = SGF.getLoweredType(E->getType());
 
       auto injection = cast<InjectIntoOptionalExpr>(E->getActor());
-      auto erasure = cast<ErasureExpr>(injection->getSubExpr());
-      auto conformance = erasure->getConformances()[0];
+      ProtocolConformanceRef conformance;
+      Expr *origActorExpr;
+      if (isolation.isDistributedActor()) {
+        // Create a reference to the asLocalActor getter.
+        auto asLocalActorDecl = getDistributedActorAsLocalActorComputedProperty(
+            SGF.F.getDeclContext()->getParentModule());
+        auto asLocalActorGetter = asLocalActorDecl->getAccessor(AccessorKind::Get);
+        SILDeclRef asLocalActorRef = SILDeclRef(
+            asLocalActorGetter, SILDeclRef::Kind::Func);
+        SGF.emitGlobalFunctionRef(E, asLocalActorRef);
+
+        // Extract the base ('self') and the DistributedActor conformance.
+        auto memberRef = cast<MemberRefExpr>(injection->getSubExpr());
+        conformance = memberRef->getDecl().getSubstitutions()
+            .getConformances()[0];
+        origActorExpr = memberRef->getBase();
+      } else {
+        auto erasure = cast<ErasureExpr>(injection->getSubExpr());
+        conformance = erasure->getConformances()[0];
+        origActorExpr = erasure->getSubExpr();
+      }
       SGF.SGM.useConformance(conformance);
 
-      auto origActorExpr = erasure->getSubExpr();
-      auto actorProto = ctx.getProtocol(KnownProtocolKind::Actor);
       SubstitutionMap subs = SubstitutionMap::getProtocolSubstitutions(
-          actorProto, origActorExpr->getType(), conformance);
+          conformance.getRequirement(), origActorExpr->getType(), conformance);
       auto origActor = SGF.maybeEmitValueOfLocalVarDecl(
           ctor->getImplicitSelfDecl(), AccessKind::Read).getValue();
       auto call = SGF.B.createBuiltin(E, builtinName, resultTy, subs, origActor);
-      return RValue(SGF, E,
-                    ManagedValue::forForwardedRValue(SGF, call));
+      return RValue(SGF, E, ManagedValue::forForwardedRValue(SGF, call));
     }
   }
 

lib/SILOptimizer/Mandatory/DIMemoryUseCollector.cpp

@@ -1714,7 +1714,9 @@ void ElementUseCollector::collectClassSelfUses(
     // they have a fully-formed 'self' to use.
     if (auto builtin = dyn_cast<BuiltinInst>(User)) {
       if (auto builtinKind = builtin->getBuiltinKind()) {
-        if (*builtinKind == BuiltinValueKind::FlowSensitiveSelfIsolation) {
+        if (*builtinKind == BuiltinValueKind::FlowSensitiveSelfIsolation ||
+            *builtinKind ==
+              BuiltinValueKind::FlowSensitiveDistributedSelfIsolation) {
           Kind = DIUseKind::FlowSensitiveSelfIsolation;
         }
       }

lib/SILOptimizer/Mandatory/DefiniteInitialization.cpp

@@ -15,6 +15,7 @@
 #include "DIMemoryUseCollector.h"
 #include "swift/AST/DiagnosticEngine.h"
 #include "swift/AST/DiagnosticsSIL.h"
+#include "swift/AST/DistributedDecl.h"
 #include "swift/AST/Expr.h"
 #include "swift/AST/Stmt.h"
 #include "swift/ClangImporter/ClangModule.h"
@@ -1361,25 +1362,51 @@ void LifetimeChecker::handleFlowSensitiveActorIsolationUse(
   SILType optExistentialType = builtinInst->getType();
   SILLocation loc = builtinInst->getLoc();
   if (isInitializedAtUse(Use, &IsSuperInitComplete, &FailedSelfUse)) {
-    // 'self' is initialized, so replace this builtin with an injection of the
-    // argument into (any Actor)?.
-
-    // Create a copy of the actor argument, which we intentionally did not
-    // copy in SILGen.
-    SILValue actor = B.createCopyValue(loc, builtinInst->getArguments()[0]);
-
-    // Inject 'self' into 'any Actor'.
-    ProtocolConformanceRef conformances[1] = {
-      builtinInst->getSubstitutions().getConformances()[0]
-    };
-    SILType existentialType = optExistentialType.getOptionalObjectType();
-    SILValue existentialBox = B.createInitExistentialRef(
-        loc, existentialType, actor->getType().getASTType(), actor,
-        ctx.AllocateCopy(conformances));
+    // 'self' is initialized, so replace this builtin with the appropriate
+    // operation to produce `any Actor.
+
+    SILValue anyActorValue;
+    auto conformance = builtinInst->getSubstitutions().getConformances()[0];
+    if (builtinInst->getBuiltinKind() == BuiltinValueKind::FlowSensitiveSelfIsolation) {
+      // Create a copy of the actor argument, which we intentionally did not
+      // copy in SILGen.
+      SILValue actor = B.createCopyValue(loc, builtinInst->getArguments()[0]);
+
+      // Inject 'self' into 'any Actor'.
+      ProtocolConformanceRef conformances[1] = { conformance };
+      SILType existentialType = optExistentialType.getOptionalObjectType();
+      anyActorValue = B.createInitExistentialRef(
+          loc, existentialType, actor->getType().getASTType(), actor,
+          ctx.AllocateCopy(conformances));
+    } else {
+      // Borrow the actor argument, which we need to form the appropriate
+      // call to the asLocalActor getter.
+      SILValue actor = B.createBeginBorrow(loc, builtinInst->getArguments()[0]);
+
+      // Dig out the getter for asLocalActor.
+      auto asLocalActorDecl = getDistributedActorAsLocalActorComputedProperty(
+          F.getDeclContext()->getParentModule());
+      auto asLocalActorGetter = asLocalActorDecl->getAccessor(AccessorKind::Get);
+      SILDeclRef asLocalActorRef = SILDeclRef(
+          asLocalActorGetter, SILDeclRef::Kind::Func);
+      SILFunction *asLocalActorFunc = F.getModule()
+          .lookUpFunction(asLocalActorRef);
+      SILValue asLocalActorValue = B.createFunctionRef(loc, asLocalActorFunc);
+
+      // Call asLocalActor. It produces an 'any Actor'.
+      anyActorValue = B.createApply(
+          loc,
+          asLocalActorValue,
+          SubstitutionMap::get(asLocalActorGetter->getGenericSignature(),
+                               { actor->getType().getASTType() },
+                               { conformance }),
+                               { actor });
+      B.createEndBorrow(loc, actor);
+    }
 
     // Then, wrap it in an optional.
     replacement = B.createEnum(
-        loc, existentialBox, ctx.getOptionalSomeDecl(), optExistentialType);
+        loc, anyActorValue, ctx.getOptionalSomeDecl(), optExistentialType);
   } else {
     // 'self' is not initialized yet, so use 'nil'.
     replacement = B.createEnum(

lib/SILOptimizer/Transforms/AccessEnforcementReleaseSinking.cpp

@@ -160,6 +160,7 @@ static bool isBarrier(SILInstruction *inst) {
     case BuiltinValueKind::InjectEnumTag:
     case BuiltinValueKind::ExtractFunctionIsolation:
     case BuiltinValueKind::FlowSensitiveSelfIsolation:
+    case BuiltinValueKind::FlowSensitiveDistributedSelfIsolation:
     case BuiltinValueKind::AddressOfRawLayout:
       return false;
 

test/SILOptimizer/definite_init_flow_sensitive_distributed_actor_self.swift

@@ -0,0 +1,37 @@
+// RUN: %target-swift-frontend -emit-sil %s -module-name test -swift-version 5 -sil-verify-all | %FileCheck %s
+// REQUIRES: concurrency, distributed
+
+import Distributed
+
+@available(SwiftStdlib 5.1, *)
+func f(isolatedTo actor: isolated (any Actor)?) async -> Int { 0 }
+
+@available(SwiftStdlib 5.7, *)
+distributed actor NotCodableDA<ActorSystem>
+  where ActorSystem: DistributedActorSystem<any Codable> {
+  let number: Int
+
+  // CHECK-LABEL: sil hidden{{.*}}@$s4test12NotCodableDAC11actorSystemACyxGx_tYacfc : $@convention(method) @async <ActorSystem where ActorSystem : DistributedActorSystem, ActorSystem.SerializationRequirement == any Decodable & Encodable> (@in ActorSystem, @sil_isolated @owned NotCodableDA<ActorSystem>) -> @owned NotCodableDA<ActorSystem> {
+  init(actorSystem: ActorSystem) async {
+    self.actorSystem = actorSystem
+
+    // First use of #isolation, which is replaced by 'nil'.
+    // CHECK: [[ISOLATION_1:%.*]] = enum $Optional<any Actor>, #Optional.none!enumelt
+    // CHECK: [[F_1:%.*]] = function_ref @$s4test1f10isolatedToSiScA_pSgYi_tYaF
+    // CHECK-NEXT: [[F_RESULT:%.*]] = apply [[F_1]]([[ISOLATION_1]])
+
+    // Assignment to "number" of the result.
+    // CHECK: [[NUMBER:%.*]] = ref_element_addr {{%.*}} : $NotCodableDA<ActorSystem>, #NotCodableDA.number
+    // CHECK: store [[F_RESULT]] to [[NUMBER]]
+    self.number = await f(isolatedTo: #isolation)
+
+    // Second use of #isolation, which uses 'self.asLocalActor''
+    // CHECK: [[AS_LOCAL_ACTOR_FN:%.*]] = function_ref @$s11Distributed0A5ActorPAAE07asLocalB0ScA_pvg : $@convention(method) <τ_0_0 where τ_0_0 : DistributedActor> (@sil_isolated @guaranteed τ_0_0) -> @owned any Actor
+    // CHECK-NEXT: [[ACTOR_EXISTENTIAL:%.*]] =  apply [[AS_LOCAL_ACTOR_FN]]<NotCodableDA<ActorSystem>>(%1) : $@convention(method) <τ_0_0 where τ_0_0 : DistributedActor> (@sil_isolated @guaranteed τ_0_0) -> @owned any Actor
+    // CHECK: [[ISOLATION_2:%.*]] = enum $Optional<any Actor>, #Optional.some!enumelt, [[ACTOR_EXISTENTIAL]]
+    // CHECK: [[F_2:%.*]] = function_ref @$s4test1f10isolatedToSiScA_pSgYi_tYaF
+    // CHECK-NEXT: apply [[F_2]]([[ISOLATION_2]])
+
+    _ = await f(isolatedTo: #isolation)
+  }
+}