Merge pull request #74225 from DougGregor/flow-sensitive-actor-init-isolation

Teach `#isolation` to respect the flow-sensitive nature of actor initializers

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,6 +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) -> (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,6 +2089,20 @@ static ValueDecl *getHopToActor(ASTContext &ctx, Identifier id) {
   return builder.build(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),
+                            isDistributed ? _distributedActor : _actor)),
+      _parameters(_typeparam(0)),
+      _optional(_existential(_actor)));
+}
+
 static ValueDecl *getDistributedActorAsAnyActor(ASTContext &ctx, Identifier id) {
   BuiltinFunctionBuilder builder(ctx);
   auto *distributedActorProto = ctx.getProtocol(KnownProtocolKind::DistributedActor);
@@ -3191,6 +3205,12 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
   case BuiltinValueKind::HopToActor:
     return getHopToActor(Context, Id);
 
+  case BuiltinValueKind::FlowSensitiveSelfIsolation:
+    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,8 @@ BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, EndAsyncLetLifetime)
 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

@@ -630,6 +630,8 @@ CONSTANT_OWNERSHIP_BUILTIN(None, CreateTaskGroup)
 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"
@@ -6743,6 +6744,56 @@ RValue RValueEmitter::visitMacroExpansionExpr(MacroExpansionExpr *E,
 
 RValue RValueEmitter::visitCurrentContextIsolationExpr(
     CurrentContextIsolationExpr *E, SGFContext C) {
+  // If we are in an actor initializer that is isolated to, the current context
+  // isolation flow-sensitive: before 'self' has been initialized, it will be
+  // nil. After 'self' has been initialized, it will be 'self'. Introduce a
+  // custom builtin that Definite Initialization will rewrite appropriately.
+  if (auto ctor = dyn_cast_or_null<ConstructorDecl>(
+          SGF.F.getDeclRef().getDecl())) {
+    auto isolation = getActorIsolation(ctor);
+    if (ctor->isDesignatedInit() &&
+        isolation == ActorIsolation::ActorInstance &&
+        isolation.getActorInstance() == ctor->getImplicitSelfDecl()) {
+      ASTContext &ctx = SGF.getASTContext();
+      auto builtinName = ctx.getIdentifier(
+          isolation.isDistributedActor()
+              ? getBuiltinName(BuiltinValueKind::FlowSensitiveDistributedSelfIsolation)
+              : getBuiltinName(BuiltinValueKind::FlowSensitiveSelfIsolation));
+      SILType resultTy = SGF.getLoweredType(E->getType());
+
+      auto injection = cast<InjectIntoOptionalExpr>(E->getActor());
+      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);
+
+      SubstitutionMap subs = SubstitutionMap::getProtocolSubstitutions(
+          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 visit(E->getActor(), C);
 }
 

lib/SILOptimizer/Mandatory/DIMemoryUseCollector.cpp

@@ -1709,6 +1709,19 @@ void ElementUseCollector::collectClassSelfUses(
       Kind = DIUseKind::Escape;
     }
 
+    // Track flow-sensitive 'self' isolation builtins separately, because they
+    // aren't really uses of 'self' until after DI, once we've decided whether
+    // they have a fully-formed 'self' to use.
+    if (auto builtin = dyn_cast<BuiltinInst>(User)) {
+      if (auto builtinKind = builtin->getBuiltinKind()) {
+        if (*builtinKind == BuiltinValueKind::FlowSensitiveSelfIsolation ||
+            *builtinKind ==
+              BuiltinValueKind::FlowSensitiveDistributedSelfIsolation) {
+          Kind = DIUseKind::FlowSensitiveSelfIsolation;
+        }
+      }
+    }
+
     trackUse(DIMemoryUse(User, Kind, 0, TheMemory.getNumElements()));
   }
 }

lib/SILOptimizer/Mandatory/DIMemoryUseCollector.h

@@ -298,7 +298,12 @@ enum DIUseKind {
   LoadForTypeOfSelf,
 
   /// This instruction is a value_metatype on the address of 'self'.
-  TypeOfSelf
+  TypeOfSelf,
+
+  /// This instruction is the builtin for flow-sensitive current isolation
+  /// within an actor initializer. It will be replaced with either a copy of
+  /// its argument (injected into an (any Actor)?) or nil.
+  FlowSensitiveSelfIsolation,
 };
 
 /// This struct represents a single classified access to the memory object

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"
@@ -491,6 +492,7 @@ namespace {
     void handleTypeOfSelfUse(DIMemoryUse &Use);
     void handleInOutUse(const DIMemoryUse &Use);
     void handleEscapeUse(const DIMemoryUse &Use);
+    void handleFlowSensitiveActorIsolationUse(const DIMemoryUse &Use);
 
     bool diagnoseReturnWithoutInitializingStoredProperties(
         const SILInstruction *Inst, SILLocation loc, const DIMemoryUse &Use);
@@ -565,6 +567,7 @@ LifetimeChecker::LifetimeChecker(const DIMemoryObjectInfo &TheMemory,
     case DIUseKind::LoadForTypeOfSelf:
     case DIUseKind::TypeOfSelf:
     case DIUseKind::Escape:
+    case DIUseKind::FlowSensitiveSelfIsolation:
       continue;
     case DIUseKind::Assign:
     case DIUseKind::Set:
@@ -1160,6 +1163,10 @@ void LifetimeChecker::doIt() {
     case DIUseKind::BadExplicitStore:
       diagnoseBadExplicitStore(Inst);
       break;
+
+    case DIUseKind::FlowSensitiveSelfIsolation:
+      handleFlowSensitiveActorIsolationUse(Use);
+      break;
     }
   }
 
@@ -1344,6 +1351,73 @@ void LifetimeChecker::handleTypeOfSelfUse(DIMemoryUse &Use) {
   }
 }
 
+void LifetimeChecker::handleFlowSensitiveActorIsolationUse(
+    const DIMemoryUse &Use) {
+  bool IsSuperInitComplete, FailedSelfUse;
+
+  ASTContext &ctx = F.getASTContext();
+  auto builtinInst = cast<BuiltinInst>(Use.Inst);
+  SILBuilderWithScope B(builtinInst);
+  SILValue replacement;
+  SILType optExistentialType = builtinInst->getType();
+  SILLocation loc = builtinInst->getLoc();
+  if (isInitializedAtUse(Use, &IsSuperInitComplete, &FailedSelfUse)) {
+    // '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, anyActorValue, ctx.getOptionalSomeDecl(), optExistentialType);
+  } else {
+    // 'self' is not initialized yet, so use 'nil'.
+    replacement = B.createEnum(
+        loc, SILValue(), ctx.getOptionalNoneDecl(), optExistentialType);
+  }
+
+  // Introduce the replacement.
+  InstModCallbacks callbacks;
+  replaceAllUsesAndErase(builtinInst, replacement, callbacks);
+}
+
 void LifetimeChecker::emitSelfConsumedDiagnostic(SILInstruction *Inst) {
   if (!shouldEmitError(Inst))
     return;

lib/SILOptimizer/Transforms/AccessEnforcementReleaseSinking.cpp

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