[Distributed] make remote proxy allocation minimally small

Author: ktoso

include/swift/AST/KnownSDKTypes.def

@@ -57,6 +57,7 @@ KNOWN_SDK_TYPE_DECL(Distributed, DistributedTargetInvocationEncoder, ProtocolDec
 KNOWN_SDK_TYPE_DECL(Distributed, DistributedTargetInvocationDecoder, ProtocolDecl, 0)
 KNOWN_SDK_TYPE_DECL(Distributed, RemoteCallTarget, StructDecl, 0)
 KNOWN_SDK_TYPE_DECL(Distributed, RemoteCallArgument, StructDecl, 1)
+KNOWN_SDK_TYPE_DECL(Distributed, _DistributedRemoteReferenceLayout, ClassDecl, 2)
 
 // String processing
 KNOWN_SDK_TYPE_DECL(StringProcessing, Regex, StructDecl, 1)

include/swift/Runtime/RuntimeFunctions.def

@@ -2447,7 +2447,7 @@ FUNCTION(NonDefaultDistributedActorInitialize,
 // );
 FUNCTION(DistributedActorInitializeRemote,
          swift_distributedActor_remote_initialize, SwiftCC,
-         ConcurrencyAvailability, // TODO(distributed): Introduce DistributedAvailability once shipping somewhere
+         ConcurrencyAvailability,
          RETURNS(OpaquePtrTy),
          ARGS(TypeMetadataPtrTy),
          ATTRS(NoUnwind),

lib/SILGen/SILGenDistributed.cpp

@@ -459,16 +459,30 @@ void SILGenFunction::emitDistributedActorFactory(FuncDecl *fd) { // TODO(distrib
   // ==== Case 'none') Create the remote instance
   {
     B.emitBlock(makeProxyBB);
+
     // ==== Create 'remote' distributed actor instance
 
-    // --- Call: _distributedActorRemoteInitialize(Self.self)
+    // --- Prepare metadata for remote "proxy" object initialization
+    auto remoteReferenceStubDecl = C.get_DistributedRemoteReferenceLayoutDecl();
+    auto systemTy = getDistributedActorSystemType(selfTyDecl);
+
+    auto remoteReferenceStubType = BoundGenericClassType::get(
+        remoteReferenceStubDecl, nullptr, {systemTy});
+    auto remoteReferenceStubMetatypeType =
+        getLoweredType(MetatypeType::get(remoteReferenceStubType, C));
+
+    SILValue remoteReferenceStubMetaTypeValue =
+        B.createMetatype(loc, remoteReferenceStubMetatypeType);
+
+    // --- Call: _distributedActorRemoteInitialize(_DistributedRemoteReferenceLayout.self)
     auto builtinName = C.getIdentifier(
         getBuiltinName(BuiltinValueKind::InitializeDistributedRemoteActor));
     auto *remote = B.createBuiltin(
         loc, builtinName,
         /*returnTy*/returnTy,
         /*subs*/ {},
-        {selfMetatypeValue});
+//        {selfMetatypeValue});
+        {remoteReferenceStubMetaTypeValue});
 
     // ==== Initialize distributed actor properties
     loc.markAutoGenerated();

lib/Sema/CodeSynthesisDistributedActor.cpp

@@ -831,6 +831,8 @@ addDistributedActorCodableConformance(
 /******************************************************************************/
 /******************************************************************************/
 
+// IMPORTANT: If adding new fields, make sure that the remote stub type has the
+// same fields and order: _DistributedRemoteReferenceLayout
 void swift::assertRequiredSynthesizedPropertyOrder(ASTContext &Context,
                                                    NominalTypeDecl *nominal) {
 #ifndef NDEBUG

stdlib/public/Concurrency/Actor.cpp

@@ -2311,7 +2311,11 @@ OpaqueValue*
 swift::swift_distributedActor_remote_initialize(const Metadata *actorType) {
   const ClassMetadata *metadata = actorType->getClassObject();
 
-  // TODO(distributed): make this allocation smaller
+  fprintf(stderr, "[%s:%d](%s) initialize remote, size: %d\n", __FILE_NAME__, __LINE__, __FUNCTION__,
+          metadata->getInstanceSize());
+  fprintf(stderr, "[%s:%d](%s) initialize remote, align mask: %d\n", __FILE_NAME__, __LINE__, __FUNCTION__,
+          metadata->getInstanceAlignMask());
+
   // ==== Allocate the memory for the remote instance
   HeapObject *alloc = swift_allocObject(metadata,
                                         metadata->getInstanceSize(),

stdlib/public/Distributed/DistributedActor.swift

@@ -449,3 +449,16 @@ public protocol _DistributedActorStub where Self: DistributedActor {}
 public func _distributedStubFatalError(function: String = #function) -> Never {
   fatalError("Unexpected invocation of distributed method '\(function)' stub!")
 }
+
+/// Type that is used to replicate the expected memory layout of a remote
+// distributed actor reference object.
+@available(SwiftStdlib 6.0, *)
+public class _DistributedRemoteReferenceLayout<ActorSystem>
+  where ActorSystem: DistributedActorSystem {
+  // Fields must be in this specific order.
+  let id: ActorSystem.ActorID
+  let actorSystem: ActorSystem
+  let unownedExecutor: UnownedSerialExecutor
+
+  private init() { fatalError("Not intended to have instances") }
+}

test/Concurrency/Runtime/async_task_locals_groups.swift

@@ -17,6 +17,7 @@ enum TL {
 @available(SwiftStdlib 5.1, *)
 @discardableResult
 func printTaskLocal<V>(
+
     _ key: TaskLocal<V>,
     _ expected: V? = nil,
     file: String = #file, line: UInt = #line