inject actorReady calls for async dist actor ctors

Immediately after the hop_to_executor in an async, distributed
actor init, we need to notify the transport that the actor is ready.

This patch does not yet account for more complex cases. In particular,
we will need a mechanism to prevent multiple calls to actorReady,
which can happen if a loop appears in the init:

distributed actor Dactor {
  var x: Int
  init(tport: ActorTransport, count: Int) async {
    var i = count
    repeat {
      self.x = count
      // hop is injected here
      i -= 1
    } while i > 0
  }
}

Author: kavon

include/swift/SILOptimizer/Utils/DistributedActor.h

@@ -24,19 +24,17 @@ class SILFunction;
 class SILLocation;
 class SILValue;
 
-/// Perform a load of the given distributed actor's transport.
-/// \param actorSelf the value representing `self` for the distributed actor
-/// instance. \returns the transport value
-SILValue loadActorTransport(SILBuilder &B, SILLocation loc,
-                            ClassDecl *actorDecl, SILValue actorSelf);
+/// Finds the first `ActorTransport`-compatible parameter of the given function.
+/// \returns nullptr if the function does not have such a parameter.
+SILArgument *findFirstActorTransportArg(SILFunction &F);
 
 /// Emits code that notifies the distributed actor's transport that the
 /// actor is ready for execution.
 /// \param B the builder to use when emitting the code.
 /// \param actor the distributed actor instance to pass to the transport as
 /// being "ready" \param transport a value representing the ActorTransport
-void emitActorReadyCall(SILBuilder &B, SILLocation loc, ClassDecl *actorDecl,
-                        SILValue actor, SILValue transport);
+void emitActorReadyCall(SILBuilder &B, SILLocation loc, SILValue actor,
+                        SILValue transport);
 
 } // namespace swift
 

lib/SILGen/SILGenDistributed.cpp

@@ -100,39 +100,6 @@ static void emitDistributedIfRemoteBranch(SILGenFunction &SGF,
 
 // MARK: local instance initialization
 
-/// Finds the first `ActorTransport`-compatible parameter of the given function.
-/// Crashes if the given function does not have such a parameter.
-static SILArgument *findFirstActorTransportArg(SILFunction &F) {
-  auto *module = F.getModule().getSwiftModule();
-  auto &C = F.getASTContext();
-
-  auto *transportProto = C.getProtocol(KnownProtocolKind::ActorTransport);
-  Type transportTy = transportProto->getDeclaredInterfaceType();
-
-  for (auto arg : F.getArguments()) {
-    // TODO(distributed): also be able to locate a generic transport
-    Type argTy = arg->getType().getASTType();
-    auto argDecl = arg->getDecl();
-
-    auto conformsToTransport =
-        module->lookupConformance(argDecl->getInterfaceType(), transportProto);
-
-    // Is it a protocol that conforms to ActorTransport?
-    if (argTy->isEqual(transportTy) || conformsToTransport) {
-      return arg;
-    }
-
-    // Is it some specific ActorTransport?
-    auto result = module->lookupConformance(argTy, transportProto);
-    if (!result.isInvalid()) {
-      return arg;
-    }
-  }
-
-  // did not find argument of ActorTransport type!
-  llvm_unreachable("Missing required ActorTransport argument!");
-}
-
 /// For the initialization of a local distributed actor instance, emits code to initialize the instance's
 /// stored property corresponding to the transport.
 static void emitTransportInit(SILGenFunction &SGF,
@@ -200,8 +167,7 @@ static void emitIdentityInit(SILGenFunction &SGF, ConstructorDecl *ctor,
   assert(distributedActorProto);
   assert(transportProto);
 
-  SILValue transportValue =
-      loadActorTransport(B, loc, classDecl, borrowedSelfArg.getValue());
+  SILValue transportValue = findFirstActorTransportArg(F);
 
   // --- Open the transport existential, if needed.
   auto transportASTType = transportValue->getType().getASTType();
@@ -296,15 +262,12 @@ void SILGenFunction::emitDistributedActorReady(
   // Only designated initializers get the lifecycle handling injected
   assert(ctor->isDesignatedInit());
 
-  auto *dc = ctor->getDeclContext();
-  auto classDecl = dc->getSelfClassDecl();
-
   SILValue transport = findFirstActorTransportArg(F);
 
   FullExpr scope(Cleanups, CleanupLocation(loc));
   auto borrowedSelf = actorSelf.borrow(*this, loc);
 
-  emitActorReadyCall(B, loc, classDecl, borrowedSelf.getValue(), transport);
+  emitActorReadyCall(B, loc, borrowedSelf.getValue(), transport);
 }
 
 // MARK: remote instance initialization

lib/SILOptimizer/Mandatory/DefiniteInitialization.cpp

@@ -27,6 +27,7 @@
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "swift/SILOptimizer/Utils/CFGOptUtils.h"
+#include "swift/SILOptimizer/Utils/DistributedActor.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallBitVector.h"
@@ -797,6 +798,16 @@ void LifetimeChecker::diagnoseInitError(const DIMemoryUse &Use,
     diagnose(Module, TheMemory.getLoc(), diag::variable_defined_here, isLet);
 }
 
+/// Determines whether the given function is a constructor that belogs to a
+/// distributed actor declaration.
+static bool isDistributedActorCtor(SILFunction &F) {
+  auto *context = F.getDeclContext();
+  if (auto *ctor = dyn_cast_or_null<ConstructorDecl>(context->getAsDecl()))
+    if (auto *cls = dyn_cast<ClassDecl>(ctor->getDeclContext()->getAsDecl()))
+      return cls->isDistributedActor();
+  return false;
+}
+
 /// Injects a hop_to_executor instruction after the specified insertion point.
 static void injectHopToExecutorAfter(SILLocation loc,
                                      SILBasicBlock::iterator insertPt,
@@ -820,6 +831,13 @@ static void injectHopToExecutorAfter(SILLocation loc,
 
       b.createHopToExecutor(loc.asAutoGenerated(), actor, /*mandatory=*/false);
 
+      // Distributed actors also need to notify their transport immediately
+      // after performing the hop.
+      if (isDistributedActorCtor(b.getFunction())) {
+        auto transport = findFirstActorTransportArg(b.getFunction());
+        emitActorReadyCall(b, loc.asAutoGenerated(), actor, transport);
+      }
+
       if (needsBorrow)
         b.createEndBorrow(loc, actor);
     });

lib/SILOptimizer/Utils/DistributedActor.cpp

@@ -18,33 +18,46 @@
 
 namespace swift {
 
-// MARK: utilities
-
-SILValue loadActorTransport(SILBuilder &B, SILLocation loc,
-                            ClassDecl *actorDecl, SILValue actorSelf) {
-  assert(actorDecl->isDistributedActor());
+SILArgument *findFirstActorTransportArg(SILFunction &F) {
+  auto *module = F.getModule().getSwiftModule();
+  auto &C = F.getASTContext();
+
+  auto *transportProto = C.getProtocol(KnownProtocolKind::ActorTransport);
+  Type transportTy = transportProto->getDeclaredInterfaceType();
+
+  for (auto arg : F.getArguments()) {
+    // TODO(distributed): also be able to locate a generic transport
+    Type argTy = arg->getType().getASTType();
+    auto argDecl = arg->getDecl();
+
+    auto conformsToTransport =
+        module->lookupConformance(argDecl->getInterfaceType(), transportProto);
+
+    // Is it a protocol that conforms to ActorTransport?
+    if (argTy->isEqual(transportTy) || conformsToTransport) {
+      return arg;
+    }
+
+    // Is it some specific ActorTransport?
+    auto result = module->lookupConformance(argTy, transportProto);
+    if (!result.isInvalid()) {
+      return arg;
+    }
+  }
 
-  // get the VarDecl corresponding to the transport
-  auto &C = actorDecl->getASTContext();
-  auto refs = actorDecl->lookupDirect(C.Id_actorTransport);
-  assert(refs.size() == 1);
-  VarDecl *prop = dyn_cast<VarDecl>(refs.front());
+#ifndef NDEBUG
+  llvm_unreachable("Missing required ActorTransport argument!");
+#endif
 
-  // form a reference and load it
-  auto &F = B.getFunction();
-  auto fieldAddr = B.createRefElementAddr(
-      loc, actorSelf, prop, F.getLoweredType(prop->getInterfaceType()));
-  return B.createLoad(loc, fieldAddr, LoadOwnershipQualifier::Copy);
+  return nullptr;
 }
 
-// MARK: exposed interface
-
-void emitActorReadyCall(SILBuilder &B, SILLocation loc, ClassDecl *actorDecl,
-                        SILValue actor, SILValue transport) {
+void emitActorReadyCall(SILBuilder &B, SILLocation loc, SILValue actor,
+                        SILValue transport) {
 
   auto &F = B.getFunction();
   auto &M = B.getModule();
-  auto &C = actorDecl->getASTContext();
+  auto &C = F.getASTContext();
 
   ProtocolDecl *actorProto = C.getProtocol(KnownProtocolKind::DistributedActor);
   ProtocolDecl *transProto = C.getProtocol(KnownProtocolKind::ActorTransport);
@@ -69,10 +82,7 @@ void emitActorReadyCall(SILBuilder &B, SILLocation loc, ClassDecl *actorDecl,
   assert(!transportConfRef.isInvalid() &&
          "Missing conformance to `ActorTransport`");
 
-  auto *selfTyDecl = actorDecl->getSelfNominalTypeDecl();
-  auto selfTy = F.mapTypeIntoContext(
-      selfTyDecl->getDeclaredInterfaceType()); // TODO: thats just self var devl
-                                               // getType
+  Type selfTy = F.mapTypeIntoContext(actor->getType().getASTType());
 
   // Note: it does not matter on what module we perform the lookup,
   // it is currently ignored. So the Stdlib module is good enough.