[Distributed] Distributed remote ref should be as small as possible (swiftlang#87449)

Because we guaranteed that one cannot access a distribted "remote"
instance's state via type system, we can allocate the instance much
smaller because we never access those user initialized fields.

This way a remote distributed actor reference is always 128 bytes, this
is because actors have special storage 12*8 bytes of PrivateData in
addition to the 16 bytes object header.

With this optimization, regardless how many fields an actor has, the
remote ref always is the same size.

resolves rdar://81825648

Used claude to generate a bunch of tests to check the sizes of actors
etc, I think this is correct (and getting it wrong totally just crashes
immediately anyway). Could use a review though, thank you!

cc @mikeash @xedin

Author: ktoso

include/swift/Runtime/Metadata.h

@@ -217,6 +217,7 @@ bool swift_compareTypeContextDescriptors(const TypeContextDescriptor *lhs,
                                          const TypeContextDescriptor *rhs);
 
 /// Compute the bounds of class metadata with a resilient superclass.
+SWIFT_EXPORT_FROM_ATTRIBUTE(swiftCore) // Cannot use SWIFT_RUNTIME_EXPORT because it is not compatible with C linkage
 ClassMetadataBounds getResilientMetadataBounds(
                                            const ClassDescriptor *descriptor);
 int32_t getResilientImmediateMembersOffset(const ClassDescriptor *descriptor);

stdlib/public/Concurrency/Actor.cpp

@@ -1962,12 +1962,18 @@ void DefaultActorImpl::deallocate() {
 #endif
 }
 
+static size_t
+getDistributedRemoteActorAllocSize(const ClassMetadata *metadata);
+
 void DefaultActorImpl::deallocateUnconditional() {
   concurrency::trace::actor_deallocate(this);
 
 #if !SWIFT_CONCURRENCY_EMBEDDED
   auto metadata = cast<ClassMetadata>(this->metadata);
-  swift_deallocClassInstance(this, metadata->getInstanceSize(),
+  size_t deallocSize = isDistributedRemote() ?
+      getDistributedRemoteActorAllocSize(metadata)
+      : metadata->getInstanceSize();
+  swift_deallocClassInstance(this, deallocSize,
                              metadata->getInstanceAlignMask());
 #else
   // Embedded Swift's runtime doesn't actually use the size/mask values.
@@ -2299,12 +2305,17 @@ static bool isDefaultActorClass(const ClassMetadata *metadata) {
 void swift::swift_defaultActor_deallocateResilient(HeapObject *actor) {
 #if !SWIFT_CONCURRENCY_EMBEDDED
   auto metadata = cast<ClassMetadata>(actor->metadata);
+
   if (isDefaultActorClass(metadata))
     return swift_defaultActor_deallocate(static_cast<DefaultActor*>(actor));
 
-  swift_deallocObject(actor, metadata->getInstanceSize(),
+  size_t deallocSize = swift_distributed_actor_is_remote(actor) ?
+        getDistributedRemoteActorAllocSize(metadata)
+        : metadata->getInstanceSize();
+  swift_deallocObject(actor, deallocSize,
                       metadata->getInstanceAlignMask());
 #else
+  // TODO(distributed): Embedded should be able to handle remote actor references here
   return swift_defaultActor_deallocate(static_cast<DefaultActor*>(actor));
 #endif
 }
@@ -2774,25 +2785,50 @@ void swift::swift_nonDefaultDistributedActor_initialize(NonDefaultDistributedAct
   asImpl(_actor)->initialize();
 }
 
+/// Compute the minimal allocation size for a 'remote' distributed actor reference.
+///
+/// Remote distributed actors are proxy objects that only need storage for the
+/// `id`, `actorSystem`, and `unownedExecutor` fields. Any user-defined stored
+/// properties beyond those three are never initialized or accessed on a remote
+/// instance, so we can trim the allocation at the offset where the first
+/// user-defined field would begin.
+static size_t
+getDistributedRemoteActorAllocSize(const ClassMetadata *metadata) {
+  auto description = metadata->getDescription();
+  uint32_t numFields = description->NumFields;
+
+  assert(numFields >= 3 &&
+         "distributed actor must have at least id, actorSystem, and "
+         "unownedExecutor fields");
+
+  if (numFields >= 4) {
+    // The 4th field is the first user-defined stored property.
+    // Its offset marks the end of the synthesized fields,
+    // so it is exactly  how much storage we need.
+    const auto *fieldOffsets = metadata->getFieldOffsets();
+    return fieldOffsets[3];
+  }
+
+  // Only the three required fields exist, remote-ref and local instances have the same size.
+  return metadata->getInstanceSize();
+}
+
 OpaqueValue*
 swift::swift_distributedActor_remote_initialize(const Metadata *actorType) {
   const ClassMetadata *metadata = actorType->getClassObject();
 
-  // TODO(distributed): make this allocation smaller
-  // ==== Allocate the memory for the remote instance
+  // Use a smaller allocation that only covers the id and actorSystem fields,
+  // since remote proxies never access user-defined stored properties.
+  size_t allocSize = getDistributedRemoteActorAllocSize(metadata);
+  assert(allocSize <= metadata->getInstanceSize() && "Remote reference size was larger than a local instance size!?");
   HeapObject *alloc = swift_allocObject(metadata,
-                                        metadata->getInstanceSize(),
+                                        allocSize,
                                         metadata->getInstanceAlignMask());
 
-  // TODO: remove this memset eventually, today we only do this to not have
-  //       to modify the destructor logic, as releasing zeroes is no-op
-  memset((void *)(alloc + 1), 0,
-         metadata->getInstanceSize() - sizeof(HeapObject));
+  // Zero the body so that the destructor can safely release fields without
+  // encountering uninitialized memory.
+  memset((void *)(alloc + 1), 0, allocSize - sizeof(HeapObject));
 
-  // TODO(distributed): a remote one does not have to have the "real"
-  //  default actor body, e.g. we don't need an executor at all; so
-  //  we can allocate more efficiently and only share the flags/status field
-  //  between the both memory representations
   // If it is a default actor, we reuse the same layout as DefaultActorImpl,
   // and store flags in the allocation directly as we initialize it.
   if (isDefaultActorClass(metadata)) {

test/abi/macOS/arm64/stdlib.swift

@@ -1209,6 +1209,9 @@ Added: _$ss4SpanVsRi_zrlE22_makeBorrowingIteratorAByxGyF
 Added: _$ss4SpanVyxGs18_BorrowingSequencesRi_zrlMc
 Added: _$ss4SpanVyxGs26_BorrowingIteratorProtocolsRi_zrlMc
 
+// Necessary for Distributed to compute distributed remote reference object size
+Added: __ZN5swift26getResilientMetadataBoundsEPKNS_21TargetClassDescriptorINS_9InProcessEEE
+
 // Borrow
 Added: _$ss6BorrowVMa
 Added: _$ss6BorrowVMn

test/abi/macOS/x86_64/stdlib.swift

@@ -1204,6 +1204,9 @@ Added: _$ss4SpanVsRi_zrlE22_makeBorrowingIteratorAByxGyF
 Added: _$ss4SpanVyxGs18_BorrowingSequencesRi_zrlMc
 Added: _$ss4SpanVyxGs26_BorrowingIteratorProtocolsRi_zrlMc
 
+// Necessary for Distributed to compute distributed remote reference object size
+Added: __ZN5swift26getResilientMetadataBoundsEPKNS_21TargetClassDescriptorINS_9InProcessEEE
+
 // Borrow
 Added: _$ss6BorrowVMa
 Added: _$ss6BorrowVMn