basic implementation of performing an init of the id after an assign of the actorSystem.

The resignID call within the initializer moved into DI, because an assignment to
the actorSystem, and thus initialization of the id, is no longer guaranteed to happen.
Thus, while we previously could model the resignation as a clean-up emitted on all
unwind paths in an initializer, now it's conditional, based on whether the id was
initialized. This is exactly what DI is designed to do, so we inject the resignation
call just before we call the identity's deinit.

Author: kavon

include/swift/SILOptimizer/Utils/DistributedActor.h

@@ -32,6 +32,13 @@ class SILLocation;
 class SILType;
 class SILValue;
 
+/// Creates a reference to the distributed actor's \p actorSystem
+/// stored property.
+SILValue refDistributedActorSystem(SILBuilder &B,
+                                   SILLocation loc,
+                                   ClassDecl *actorDecl,
+                                   SILValue actorInstance);
+
 /// Emit a call to a witness of the DistributedActorSystem protocol.
 ///
 /// \param methodName The name of the method on the DistributedActorSystem protocol.
@@ -55,6 +62,11 @@ void emitDistributedActorSystemWitnessCall(
 void emitActorReadyCall(SILBuilder &B, SILLocation loc, SILValue actor,
                         SILValue actorSystem);
 
+/// Emits code to notify the \p actorSystem that the given identity is resigned.
+void emitResignIdentityCall(SILBuilder &B, SILLocation loc, ClassDecl* actDecl,
+                            SILValue actor, SILValue identityRef);
+
+
 } // namespace swift
 
 #endif

lib/SILGen/InitializeDistActorIdentity.h

@@ -0,0 +1,40 @@
+//===--- InitializeDistActorIdentity.h - dist actor ID init for SILGen ----===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#include "Cleanup.h"
+
+namespace swift {
+
+class AllocStackInst;
+
+namespace Lowering {
+
+/// A clean-up designed to emit an initialization of a distributed actor's
+/// identity upon successful initialization of the actor's system.
+struct InitializeDistActorIdentity : Cleanup {
+private:
+  ConstructorDecl *ctor;
+  ManagedValue actorSelf;
+  VarDecl *systemVar;
+public:
+  InitializeDistActorIdentity(ConstructorDecl *ctor, ManagedValue actorSelf);
+
+  void emit(SILGenFunction &SGF, CleanupLocation loc,
+            ForUnwind_t forUnwind) override;
+
+  void dump(SILGenFunction &) const override;
+
+  VarDecl* getSystemVar() const { return systemVar; }
+};
+
+} // end Lowering namespace
+} // end swift namespace

lib/SILGen/SILGenDistributed.cpp

@@ -137,24 +137,24 @@ static SILArgument *findFirstDistributedActorSystemArg(SILFunction &F) {
   auto *module = F.getModule().getSwiftModule();
   auto &C = F.getASTContext();
 
-  auto *transportProto = C.getProtocol(KnownProtocolKind::DistributedActorSystem);
-  Type transportTy = transportProto->getDeclaredInterfaceType();
+  auto *DAS = C.getDistributedActorSystemDecl();
+  Type systemTy = DAS->getDeclaredInterfaceType();
 
   for (auto arg : F.getArguments()) {
-    // TODO(distributed): also be able to locate a generic transport
+    // TODO(distributed): also be able to locate a generic system
     Type argTy = arg->getType().getASTType();
     auto argDecl = arg->getDecl();
 
-    auto conformsToTransport =
-        module->lookupConformance(argDecl->getInterfaceType(), transportProto);
+    auto conformsToSystem =
+        module->lookupConformance(argDecl->getInterfaceType(), DAS);
 
     // Is it a protocol that conforms to DistributedActorSystem?
-    if (argTy->isEqual(transportTy) || conformsToTransport) {
+    if (argTy->isEqual(systemTy) || conformsToSystem) {
       return arg;
     }
 
     // Is it some specific DistributedActorSystem?
-    auto result = module->lookupConformance(argTy, transportProto);
+    auto result = module->lookupConformance(argTy, DAS);
     if (!result.isInvalid()) {
       return arg;
     }
@@ -172,7 +172,8 @@ static SILArgument *findFirstDistributedActorSystemArg(SILFunction &F) {
 static void emitActorSystemInit(SILGenFunction &SGF,
                                         ConstructorDecl *ctor,
                                         SILLocation loc,
-                                        ManagedValue actorSelf) {
+                                        ManagedValue actorSelf,
+                                        SILValue systemValue) {
   assert(ctor->isImplicit() && "unexpected explicit dist actor init");
   assert(ctor->isDesignatedInit());
 
@@ -183,11 +184,10 @@ static void emitActorSystemInit(SILGenFunction &SGF,
   // By construction, automatically generated distributed actor ctors have
   // exactly one ActorSystem-conforming argument to the constructor,
   // so we grab the first one from the params.
-  SILValue systemArg = findFirstDistributedActorSystemArg(SGF.F);
   VarDecl *var = lookupProperty(classDecl, C.Id_actorSystem);
   assert(var);
 
-  initializeProperty(SGF, loc, actorSelf.getValue(), var, systemArg);
+  initializeProperty(SGF, loc, actorSelf.getValue(), var, systemValue);
 }
 
 /// Emits the distributed actor's identity (`id`) initialization.
@@ -196,72 +196,98 @@ static void emitActorSystemInit(SILGenFunction &SGF,
 /// \verbatim
 ///     self.id = system.assignID(Self.self)
 /// \endverbatim
-static void emitIDInit(SILGenFunction &SGF, ConstructorDecl *ctor,
-                             SILLocation loc, ManagedValue borrowedSelfArg) {
-  assert(ctor->isImplicit() && "unexpected explicit dist actor init");
+void SILGenFunction::emitDistActorIdentityInit(ConstructorDecl *ctor,
+                                               SILLocation loc,
+                                               SILValue borrowedSelfArg,
+                                               SILValue actorSystem) {
   assert(ctor->isDesignatedInit());
 
   auto &C = ctor->getASTContext();
-  auto &B = SGF.B;
-  auto &F = SGF.F;
 
   auto *dc = ctor->getDeclContext();
   auto classDecl = dc->getSelfClassDecl();
 
+  assert(classDecl->isDistributedActor());
+
   // --- prepare `Self.self` metatype
   auto *selfTyDecl = ctor->getParent()->getSelfNominalTypeDecl();
   auto selfTy = F.mapTypeIntoContext(selfTyDecl->getDeclaredInterfaceType());
-  auto selfMetatype = SGF.getLoweredType(MetatypeType::get(selfTy));
+  auto selfMetatype = getLoweredType(MetatypeType::get(selfTy));
   SILValue selfMetatypeValue = B.createMetatype(loc, selfMetatype);
 
-  // since we're doing this only for the implicitly defined ctors, grab from arg
-  SILValue actorSystem = findFirstDistributedActorSystemArg(SGF.F);
-
   // --- create a temporary storage for the result of the call
   // it will be deallocated automatically as we exit this scope
   VarDecl *var = lookupProperty(classDecl, C.Id_id);
-  auto resultTy = SGF.getLoweredType(
-      F.mapTypeIntoContext(var->getInterfaceType()));
-  auto temp = SGF.emitTemporaryAllocation(loc, resultTy);
+  auto resultTy = getLoweredType(F.mapTypeIntoContext(var->getInterfaceType()));
+  auto temp = emitTemporaryAllocation(loc, resultTy);
 
   // --- emit the call itself.
   emitDistributedActorSystemWitnessCall(
       B, loc, C.Id_assignID,
-      actorSystem, SGF.getLoweredType(selfTy),
+      actorSystem, getLoweredType(selfTy),
       { temp, selfMetatypeValue });
 
   // --- initialize the property.
-  initializeProperty(SGF, loc, borrowedSelfArg.getValue(), var, temp);
+  initializeProperty(*this, loc, borrowedSelfArg, var, temp);
 }
 
-namespace {
-/// Cleanup to resign the identity of a distributed actor if an abnormal exit happens.
-class ResignIdentity : public Cleanup {
-  ClassDecl *actorDecl;
-  SILValue self;
-public:
-  ResignIdentity(ClassDecl *actorDecl, SILValue self)
-    : actorDecl(actorDecl), self(self) {
-      assert(actorDecl->isDistributedActor());
-    }
+InitializeDistActorIdentity::InitializeDistActorIdentity(ConstructorDecl *ctor,
+                                       ManagedValue actorSelf)
+                                       : ctor(ctor),
+                                         actorSelf(actorSelf) {
+  systemVar = ctor->getDeclContext()
+                  ->getSelfClassDecl()
+                  ->getDistributedActorSystemProperty();
+  assert(systemVar);
+}
 
-  void emit(SILGenFunction &SGF, CleanupLocation l, ForUnwind_t forUnwind) override {
-    if (forUnwind == IsForUnwind) {
-      l.markAutoGenerated();
-      SGF.emitDistributedActorSystemResignIDCall(l, actorDecl,
-                                 ManagedValue::forUnmanaged(self));
-    }
+void InitializeDistActorIdentity::emit(SILGenFunction &SGF, CleanupLocation loc,
+                              ForUnwind_t forUnwind) {
+
+  // If we're unwinding, that must mean we're in the case where the
+  // evaluating the expression being assigned to the actorSystem has
+  // thrown an error. In that case, we cannot initialize the identity,
+  // since there is no actorSystem.
+  if (forUnwind == IsForUnwind)
+    return;
+    
+
+  // Save the current clean-up depth
+  auto baseDepth = SGF.getCleanupsDepth();
+  {
+    loc.markAutoGenerated();
+    auto borrowedSelf = actorSelf.borrow(SGF, loc);
+
+    // load the actorSystem value
+    Type formalType = SGF.F.mapTypeIntoContext(systemVar->getInterfaceType());
+    SILType loweredType = SGF.getLoweredType(formalType).getAddressType();
+    auto ref =
+      SGF.B.createRefElementAddr(loc, borrowedSelf, systemVar, loweredType);
+
+    SGFContext ctx;
+    auto systemVal =
+      SGF.emitLoad(loc, ref.getValue(),
+                   SGF.getTypeLowering(loweredType), ctx, IsNotTake);
+
+    // Important that we mark the location as auto-generated, since the id
+    // is a @_compilerInitialized field.
+    SGF.emitDistActorIdentityInit(ctor, loc,
+                                borrowedSelf.getValue(), systemVal.getValue());
   }
 
-  void dump(SILGenFunction &SGF) const override {
+  // Emit any active clean-ups we just pushed.
+  while (SGF.getTopCleanup() != baseDepth)
+    SGF.Cleanups.popAndEmitCleanup(SGF.getTopCleanup(), loc, forUnwind);
+
+}
+
+void InitializeDistActorIdentity::dump(SILGenFunction &) const {
 #ifndef NDEBUG
-    llvm::errs() << "ResignIdentity "
-                 << "State:" << getState() << " "
-                 << "Self: " << self << "\n";
+  llvm::errs() << "InitializeDistActorIdentity\n"
+               << "State: " << getState()
+               << "\n";
 #endif
-  }
-};
-} // end anonymous namespace
+}
 
 void SILGenFunction::emitDistributedActorImplicitPropertyInits(
     ConstructorDecl *ctor, ManagedValue selfArg) {
@@ -273,18 +299,18 @@ void SILGenFunction::emitDistributedActorImplicitPropertyInits(
 
   selfArg = selfArg.borrow(*this, loc);
 
-  // register a clean-up to resign the identity upon abnormal exit
-  // we do this regardless of initializer kind, since it's easy to do in SILGen.
-  auto *actorDecl = cast<ClassDecl>(ctor->getParent()->getAsDecl());
-  Cleanups.pushCleanup<ResignIdentity>(actorDecl, selfArg.getValue());
-
-  // Users must initialize the actorSystem property explicitly in their ctors
-  if (!ctor->isImplicit())
+  // implicit ctors initialize the system and identity from
+  // its ActorSystem parameter.
+  if (ctor->isImplicit()) {
+    SILValue actorSystem = findFirstDistributedActorSystemArg(F);
+    emitActorSystemInit(*this, ctor, loc, selfArg, actorSystem);
+    emitDistActorIdentityInit(ctor, loc, selfArg.getValue(), actorSystem);
     return;
+  }
 
-  // implicit ctors initialize these from the ActorSystem parameter.
-  emitActorSystemInit(*this, ctor, loc, selfArg);
-  emitIDInit(*this, ctor, loc, selfArg);
+  // for explicit ctors, store (but do not push) a clean-up that will
+  // initialize the identity in whichever scope it's pushed to.
+  DistActorCtorContext = InitializeDistActorIdentity(ctor, selfArg);
 }
 
 void SILGenFunction::emitDistributedActorReady(
@@ -505,7 +531,7 @@ SILGenFunction::emitConditionalResignIdentityCall(SILLocation loc,
                                                   ManagedValue actorSelf,
                                                   SILBasicBlock *continueBB) {
   assert(actorDecl->isDistributedActor() &&
-  "only distributed actors have transport lifecycle hooks in deinit");
+  "only distributed actors have actorSystem lifecycle hooks in deinit");
 
   auto selfTy = actorDecl->getDeclaredInterfaceType();
 

lib/SILGen/SILGenFunction.h

@@ -15,6 +15,7 @@
 
 #include "FormalEvaluation.h"
 #include "Initialization.h"
+#include "InitializeDistActorIdentity.h"
 #include "JumpDest.h"
 #include "RValue.h"
 #include "SGFContext.h"
@@ -399,6 +400,11 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
                       AsyncLetChildTask>
       AsyncLetChildTasks;
 
+  /// Indicates whether this function is a distributed actor's designated
+  /// initializer, providing the needed clean-up to emit an identity
+  /// assignment after initializing the actorSystem property.
+  Optional<InitializeDistActorIdentity> DistActorCtorContext;
+
   /// When rebinding 'self' during an initializer delegation, we have to be
   /// careful to preserve the object at 1 retain count during the delegation
   /// because of assumptions in framework code. This enum tracks the state of
@@ -2071,6 +2077,17 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   void emitDistributedActorImplicitPropertyInits(
       ConstructorDecl *ctor, ManagedValue selfArg);
 
+  /// Initializes just the implicit identity property of a distributed actor.
+  /// \param selfVal a value corresponding to the actor's self
+  /// \param actorSystemVal a value corresponding to the actorSystem, to be used
+  /// to invoke its \p assignIdentity method.
+  void emitDistActorIdentityInit(ConstructorDecl *ctor,
+                                 SILLocation loc,
+                                 SILValue selfVal,
+                                 SILValue actorSystemVal);
+
+  void registerDistActorIdentityInit();
+
   /// Given a function representing a distributed actor factory, emits the
   /// corresponding SIL function for it.
   void emitDistributedActorFactory(

lib/SILGen/SILGenLValue.cpp

@@ -749,6 +749,8 @@ namespace {
 
     virtual bool isLoadingPure() const override { return true; }
 
+    VarDecl *getField() const { return Field; }
+
     ManagedValue project(SILGenFunction &SGF, SILLocation loc,
                          ManagedValue base) && override {
       assert(base.getType().hasReferenceSemantics() &&
@@ -4505,6 +4507,18 @@ void SILGenFunction::emitAssignToLValue(SILLocation loc, RValue &&src,
   emitAssignToLValue(loc, ArgumentSource(loc, std::move(src)), std::move(dest));
 }
 
+namespace {
+
+} // end anonymous namespace
+
+/// Checks if the last component of the LValue refers to the given var decl.
+static bool referencesVar(PathComponent const& comp, VarDecl *var) {
+  if (comp.getKind() != PathComponent::RefElementKind)
+    return false;
+
+  return static_cast<RefElementComponent const&>(comp).getField() == var;
+}
+
 void SILGenFunction::emitAssignToLValue(SILLocation loc,
                                         ArgumentSource &&src,
                                         LValue &&dest) {
@@ -4514,6 +4528,12 @@ void SILGenFunction::emitAssignToLValue(SILLocation loc,
   // LValue. For example: (x[i], x[j]) = a, b
   ArgumentScope argScope(*this, loc);
 
+  // If this is an assignment to a distributed actor's actorSystem, push
+  // a clean-up to also initialize its id.
+  if (LLVM_UNLIKELY(DistActorCtorContext) &&
+      referencesVar(**(dest.end()-1), DistActorCtorContext->getSystemVar()))
+    Cleanups.pushCleanup<InitializeDistActorIdentity>(*DistActorCtorContext);
+
   // If the last component is a getter/setter component, use a special
   // generation pattern that allows us to peephole the emission of the RHS.
   if (trySetterPeephole(*this, loc, std::move(src), std::move(dest))) {
@@ -4605,6 +4625,13 @@ void SILGenFunction::emitAssignLValueToLValue(SILLocation loc, LValue &&src,
     return;
   }
 
+  // If this is an assignment to a distributed actor's actorSystem, push
+  // a clean-up to also initialize its id.
+  // FIXME: in what situations does this lvalue to lvalue assign happen?
+  if (LLVM_UNLIKELY(DistActorCtorContext) &&
+      referencesVar(**(dest.end()-1), DistActorCtorContext->getSystemVar()))
+    Cleanups.pushCleanup<InitializeDistActorIdentity>(*DistActorCtorContext);
+
   auto &rvalueTL = getTypeLowering(src.getTypeOfRValue());
 
   auto srcAddr = emitAddressOfLValue(loc, std::move(src)).getUnmanagedValue();