Revert "Revert "Merge pull request #38938 from drexin/wip-dist-resolve" (#38994)" (#39011)

This reverts commit f6ae9f3.

Author: drexin

lib/IRGen/GenBuiltin.cpp

@@ -390,7 +390,7 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
 
   if (Builtin.ID == BuiltinValueKind::InitializeDistributedRemoteActor) {
     auto actorMetatype = args.claimNext();
-    emitDistributedActorInitializeRemote(IGF, actorMetatype, out);
+    emitDistributedActorInitializeRemote(IGF, resultType, actorMetatype, out);
     return;
   }
 

lib/IRGen/GenDistributed.cpp

@@ -17,6 +17,7 @@
 #include "GenDistributed.h"
 
 #include "BitPatternBuilder.h"
+#include "ClassTypeInfo.h"
 #include "ExtraInhabitants.h"
 #include "GenProto.h"
 #include "GenType.h"
@@ -32,7 +33,12 @@ using namespace swift;
 using namespace irgen;
 
 llvm::Value *irgen::emitDistributedActorInitializeRemote(
-    IRGenFunction &IGF, llvm::Value *actorMetatype, Explosion &out) {
+    IRGenFunction &IGF, SILType selfType, llvm::Value *actorMetatype, Explosion &out) {
+  auto &classTI = IGF.getTypeInfo(selfType).as<ClassTypeInfo>();
+  auto &classLayout = classTI.getClassLayout(IGF.IGM, selfType,
+                                             /*forBackwardDeployment=*/false);
+  llvm::Type *destType = classLayout.getType()->getPointerTo();
+
   auto fn = IGF.IGM.getDistributedActorInitializeRemoteFn();
   actorMetatype =
       IGF.Builder.CreateBitCast(actorMetatype, IGF.IGM.TypeMetadataPtrTy);
@@ -41,9 +47,11 @@ llvm::Value *irgen::emitDistributedActorInitializeRemote(
   call->setCallingConv(IGF.IGM.SwiftCC);
   call->setDoesNotThrow();
 
-  out.add(call);
+  auto result = IGF.Builder.CreateBitCast(call, destType);
+
+  out.add(result);
 
-  return call;
+  return result;
 }
 
 void irgen::emitDistributedActorDestroy(IRGenFunction &IGF,

lib/IRGen/GenDistributed.h

@@ -43,6 +43,7 @@ class IRGenFunction;
 /// Emit the '_distributedActorRemoteInitialize' call.
 llvm::Value *emitDistributedActorInitializeRemote(
     IRGenFunction &IGF,
+    SILType selfType,
     llvm::Value *actorMetatype,
     Explosion &out);
 

lib/SILGen/SILGenDistributed.cpp

@@ -70,6 +70,34 @@ static void emitDistributedIfRemoteBranch(SILGenFunction &SGF,
   B.createCondBranch(Loc, isRemoteResultUnwrapped, isRemoteBB, isLocalBB);
 }
 
+static AbstractFunctionDecl *lookupActorTransportResolveFunc(ASTContext &C) {
+  auto transportDecl = C.getActorTransportDecl();
+
+  for (auto decl : transportDecl->lookupDirect(DeclName(C.Id_resolve)))
+    if (auto funcDecl = dyn_cast<AbstractFunctionDecl>(decl))
+      return funcDecl;
+
+  llvm_unreachable("Missing ActorTransport.resolve function");
+}
+
+static VarDecl *lookupActorTransportProperty(ASTContext &C, ClassDecl *cd,
+                                             SILValue selfValue) {
+  auto transportVarDeclRefs = cd->lookupDirect(C.Id_actorTransport);
+  assert(transportVarDeclRefs.size() == 1);
+  return dyn_cast<VarDecl>(transportVarDeclRefs.front());
+}
+
+static EnumElementDecl *lookupEnumCase(ASTContext &C, EnumDecl *target,
+                                       Identifier identifier) {
+  auto elementDecls = target->lookupDirect(DeclName(identifier));
+  if (elementDecls.empty())
+    return nullptr;
+
+  auto *elementDecl = elementDecls.front();
+
+  return dyn_cast<EnumElementDecl>(elementDecl);
+}
+
 /******************************************************************************/
 /****************** DISTRIBUTED ACTOR STORAGE INITIALIZATION ******************/
 /******************************************************************************/
@@ -120,6 +148,7 @@ static void emitDistributedActorStore_transport(
     SILValue actorSelf, AbstractFunctionDecl *func,
     SILArgument *transportArg) {
   auto &B = SGF.B;
+
   auto &SGM = SGF.SGM;
   SILGenFunctionBuilder builder(SGM);
 
@@ -135,6 +164,7 @@ static void emitDistributedActorStore_transport(
   auto *var = dyn_cast<VarDecl>(vars.front());
 
   // ----
+
   auto fieldAddr = B.createRefElementAddr(
       loc, actorSelf, var,
       SGF.getLoweredType(var->getInterfaceType()));
@@ -166,6 +196,7 @@ emitDistributedActor_init_transportStore(
   SILValue transportArgValue = getActorTransportArgument(C, F, ctor);
 
   // ----
+
   auto transportFieldAddr = B.createRefElementAddr(
       loc, borrowedSelfArg.getValue(), var,
       SGF.getLoweredType(var->getInterfaceType()));
@@ -187,6 +218,7 @@ static void emitDistributedActorStore_id(
     SILValue actorSelf, AbstractFunctionDecl *func,
     SILArgument *identityArg) {
   auto &B = SGF.B;
+
   auto &SGM = SGF.SGM;
   SILGenFunctionBuilder builder(SGM);
 
@@ -237,7 +269,9 @@ static void emitDistributedActorStore_init_assignIdentity(
 
   // --- Prepare the arguments
   SILValue transportArgValue = getActorTransportArgument(C, F, ctor);
+
   ProtocolDecl *distributedActorProto = C.getProtocol(KnownProtocolKind::DistributedActor);
+
   assert(distributedActorProto);
   assert(transportProto);
 
@@ -382,6 +416,7 @@ void SILGenFunction::initializeDistributedActorImplicitStorageInit(
 
 void SILGenFunction::emitDistributedActorReady(
     ConstructorDecl *ctor, ManagedValue selfArg) {
+
   auto *dc = ctor->getDeclContext();
   auto classDecl = dc->getSelfClassDecl();
   auto &C = classDecl->getASTContext();
@@ -462,6 +497,86 @@ void SILGenFunction::emitDistributedActorReady(
 /******************* DISTRIBUTED ACTOR RESOLVE FUNCTION ***********************/
 /******************************************************************************/
 
+/// Synthesize the distributed actor's identity (`id`) initialization:
+///
+/// \verbatim
+///     transport.resolve(_, as:)
+/// \endverbatim
+static void createDistributedActorFactory_resolve(
+    SILGenFunction &SGF, ASTContext &C, FuncDecl *fd, SILValue identityValue,
+    SILValue transportValue, Type selfTy, SILValue selfMetatypeValue,
+    SILType resultTy, SILBasicBlock *normalBB, SILBasicBlock *errorBB) {
+  auto &B = SGF.B;
+  auto &SGM = SGF.SGM;
+  auto &F = SGF.F;
+  SILGenFunctionBuilder builder(SGM);
+
+  auto loc = SILLocation(fd);
+  loc.markAutoGenerated();
+
+  ProtocolDecl *distributedActorProto =
+      C.getProtocol(KnownProtocolKind::DistributedActor);
+  ProtocolDecl *transportProto =
+      C.getProtocol(KnownProtocolKind::ActorTransport);
+  assert(distributedActorProto);
+  assert(transportProto);
+
+  // // --- Open the transport existential
+  OpenedArchetypeType *Opened;
+  auto transportASTType = transportValue->getType().getASTType();
+  auto openedTransportType =
+      transportASTType->openAnyExistentialType(Opened)->getCanonicalType();
+  auto openedTransportSILType = F.getLoweredType(openedTransportType);
+  auto transportArchetypeValue =
+      B.createOpenExistentialAddr(loc, transportValue, openedTransportSILType,
+                                  OpenedExistentialAccess::Immutable);
+
+  // --- prepare the witness_method
+  // Note: it does not matter on what module we perform the lookup,
+  // it is currently ignored. So the Stdlib module is good enough.
+  auto *module = SGF.getModule().getSwiftModule();
+
+  // the conformance here is just an abstract thing so we can simplify
+  auto transportConfRef = ProtocolConformanceRef(transportProto);
+  assert(!transportConfRef.isInvalid() &&
+         "Missing conformance to `ActorTransport`");
+
+  auto distributedActorConfRef =
+      module->lookupConformance(selfTy, distributedActorProto);
+  assert(!distributedActorConfRef.isInvalid() &&
+         "Missing conformance to `DistributedActor`");
+
+  auto resolveMethod =
+      cast<FuncDecl>(transportProto->getSingleRequirement(C.Id_resolve));
+  auto resolveRef = SILDeclRef(resolveMethod, SILDeclRef::Kind::Func);
+  auto constantInfo =
+      SGF.getConstantInfo(SGF.getTypeExpansionContext(), resolveRef);
+  auto resolveSILTy = constantInfo.getSILType();
+
+  auto resolveWitnessMethod =
+      B.createWitnessMethod(loc,
+                            /*lookupTy*/ openedTransportType,
+                            /*Conformance*/ transportConfRef,
+                            /*member*/ resolveRef,
+                            /*methodTy*/ resolveSILTy);
+
+  // // --- prepare conformance subs
+  auto genericSig = resolveMethod->getGenericSignature();
+
+  SubstitutionMap subs =
+      SubstitutionMap::get(genericSig, {openedTransportType, selfTy},
+                           {transportConfRef, distributedActorConfRef});
+
+  // // ---- actually call transport.resolve(id, as: Self.self)
+
+  SmallVector<SILValue, 3> params;
+  params.push_back(identityValue);
+  params.push_back(selfMetatypeValue);
+  params.push_back(transportArchetypeValue); // self for the call, as last param
+
+  B.createTryApply(loc, resolveWitnessMethod, subs, params, normalBB, errorBB);
+}
+
 /// Function body of:
 /// \verbatim
 /// DistributedActor.resolve(
@@ -487,28 +602,66 @@ void SILGenFunction::emitDistributedActorFactory(FuncDecl *fd) {
   assert(
       transportArg->getType().getASTType()->isEqual(C.getActorTransportType()));
 
-  // --- Parameter: self
+  SILValue selfArgValue = F.getSelfArgument();
+  ManagedValue selfArg = ManagedValue::forUnmanaged(selfArgValue);
+
+  // type: SpecificDistributedActor.Type
+  auto selfArgType = F.mapTypeIntoContext(selfArg.getType().getASTType());
+  auto selfMetatype = getLoweredType(selfArgType);
+  SILValue selfMetatypeValue = B.createMetatype(loc, selfMetatype);
+
+  // type: SpecificDistributedActor
   auto *selfTyDecl = fd->getParent()->getSelfNominalTypeDecl();
   assert(selfTyDecl->isDistributedActor());
+  auto selfTy = F.mapTypeIntoContext(selfTyDecl->getDeclaredInterfaceType());
+  auto returnTy = getLoweredType(selfTy);
 
-  SILValue selfArgValue = F.getSelfArgument();
-  ManagedValue selfArg = ManagedValue::forUnmanaged(selfArgValue);
+  // ==== Prepare all the basic blocks
+  auto returnBB = createBasicBlock();
+  auto resolvedBB = createBasicBlock();
+  auto makeProxyBB = createBasicBlock();
+  auto switchBB = createBasicBlock();
+  auto errorBB = createBasicBlock();
+
+  SILFunctionConventions fnConv = F.getConventions(); // TODO: no idea?
+
+  // --- get the uninitialized allocation from the runtime system.
+  FullExpr scope(Cleanups, CleanupLocation(fd));
+
+  auto optionalReturnTy = SILType::getOptionalType(returnTy);
 
-  // ==== Case 'remote') Create the remote instance
+  // ==== Call `try transport.resolve(id, as: Self.self)`
   {
-    // ==== Create 'remote' distributed actor instance
-    // --- Prepare param: Self.self
-    // type: SpecificDistributedActor
-    auto returnTy = getLoweredType(
-        F.mapTypeIntoContext(selfTyDecl->getDeclaredInterfaceType()));
+    createDistributedActorFactory_resolve(
+        *this, C, fd, identityArg, transportArg, selfTy, selfMetatypeValue,
+        optionalReturnTy, switchBB, errorBB);
+  }
+
+  // ==== switch resolved { ... }
+  {
+    B.emitBlock(switchBB);
+    auto resolve =
+        switchBB->createPhiArgument(optionalReturnTy, OwnershipKind::Owned);
+
+    B.createSwitchEnum(
+        loc, resolve, nullptr,
+        {{C.getOptionalSomeDecl(), resolvedBB},
+         {std::make_pair(C.getOptionalNoneDecl(), makeProxyBB)}});
+  }
 
-    // type: SpecificDistributedActor.Type
-    auto selfMetatype =
-        getLoweredType(F.mapTypeIntoContext(selfArg.getType().getASTType()));
-    SILValue selfMetatypeValue = B.createMetatype(loc, selfMetatype);
+  // ==== Case 'some') return the resolved instance
+  {
+    B.emitBlock(resolvedBB);
+
+    auto local = resolvedBB->createPhiArgument(returnTy, OwnershipKind::Owned);
+
+    B.createBranch(loc, returnBB, {local});
+  }
 
-    // --- get the uninitialized allocation from the runtime system.
-    FullExpr scope(Cleanups, CleanupLocation(fd));
+  // ==== Case 'none') Create the remote instance
+  {
+    B.emitBlock(makeProxyBB);
+    // ==== Create 'remote' distributed actor instance
 
     // --- Call: _distributedActorRemoteInitialize(Self.self)
     auto builtinName = C.getIdentifier(
@@ -528,8 +681,32 @@ void SILGenFunction::emitDistributedActorFactory(FuncDecl *fd) {
     emitDistributedActorStore_transport(
         C, *this, /*actorSelf*/remote, fd, transportArg);
 
-    // ==== Return the fully initialized remote instance
-    B.createReturn(loc, remote);
+    // ==== Branch to return the fully initialized remote instance
+    B.createBranch(loc, returnBB, {remote});
+  }
+
+  // --- Emit return logic
+  // return <remote>
+  {
+    B.emitBlock(returnBB);
+
+    auto local = returnBB->createPhiArgument(returnTy, OwnershipKind::Owned);
+
+    Cleanups.emitCleanupsForReturn(CleanupLocation(loc), NotForUnwind);
+    B.createReturn(loc, local);
+  }
+
+  // --- Emit rethrow logic
+  // throw error
+  {
+    B.emitBlock(errorBB);
+
+    auto error = errorBB->createPhiArgument(
+        fnConv.getSILErrorType(F.getTypeExpansionContext()),
+        OwnershipKind::Owned);
+
+    Cleanups.emitCleanupsForReturn(CleanupLocation(loc), IsForUnwind);
+    B.createThrow(loc, error);
   }
 }
 
@@ -563,9 +740,8 @@ void SILGenFunction::emitDistributedActor_resignAddress(
                              getLoweredType(idVarDeclRef->getType()));
 
   // ==== locate: self.actorTransport
-  auto transportVarDeclRefs = cd->lookupDirect(ctx.Id_actorTransport);
-  assert(transportVarDeclRefs.size() == 1);
-  auto *transportVarDeclRef = dyn_cast<VarDecl>(transportVarDeclRefs.front());
+  auto transportVarDeclRef = lookupActorTransportProperty(ctx, cd, selfValue);
+
   auto transportRef =
       B.createRefElementAddr(Loc, selfValue, transportVarDeclRef,
                              getLoweredType(transportVarDeclRef->getType()));

lib/SILGen/SILGenFunction.cpp

@@ -514,17 +514,18 @@ void SILGenFunction::emitFunction(FuncDecl *fd) {
   emitProfilerIncrement(fd->getTypecheckedBody());
   emitProlog(captureInfo, fd->getParameters(), fd->getImplicitSelfDecl(), fd,
              fd->getResultInterfaceType(), fd->hasThrows(), fd->getThrowsLoc());
-  prepareEpilog(true, fd->hasThrows(), CleanupLocation(fd));
 
   if (fd->isDistributedActorFactory()) {
     // Synthesize the factory function body
     emitDistributedActorFactory(fd);
   } else {
+    prepareEpilog(true, fd->hasThrows(), CleanupLocation(fd));
+    
     // Emit the actual function body as usual
     emitStmt(fd->getTypecheckedBody());
-  }
 
-  emitEpilog(fd);
+    emitEpilog(fd);
+  }
 
   mergeCleanupBlocks();
 }

stdlib/public/Distributed/ActorTransport.swift

@@ -48,7 +48,7 @@ public protocol ActorTransport: Sendable {
   ///
   /// Detecting liveness of such remote actors shall be offered / by transport libraries
   /// by other means, such as "watching an actor for termination" or similar.
-  func resolve<Act>(_ identity: AnyActorIdentity, as actorType: Act.Type) throws -> ActorResolved<Act> // TODO(distributed): make just optional
+  func resolve<Act>(_ identity: AnyActorIdentity, as actorType: Act.Type) throws -> Act? // TODO(distributed): make just optional
       where Act: DistributedActor
 
   // ==== ---------------------------------------------------------------------
@@ -73,9 +73,3 @@ public protocol ActorTransport: Sendable {
   func resignIdentity(_ id: AnyActorIdentity)
 
 }
-
-@available(SwiftStdlib 5.5, *)
-public enum ActorResolved<Act: DistributedActor> {
-  case resolved(Act)
-  case makeProxy
-}