Merge pull request #67039 from drexin/wip-mp-gen-destr-inject-tag

Support destructiveInjectEnumTag in generic multi payload enums with layout strings

Author: drexin

include/swift/Runtime/RuntimeFunctions.def

@@ -2366,6 +2366,17 @@ FUNCTION(MultiPayloadEnumGenericGetEnumTag,
          ATTRS(NoUnwind, WillReturn),
          EFFECT(NoEffect))
 
+// void swift_multiPayloadEnumGeneric_destructiveInjectEnumTag(swift::OpaqueValue *address,
+//                                                             unsigned tag,
+//                                                             const Metadata *metadata)
+FUNCTION(MultiPayloadEnumGenericDestructiveInjectEnumTag,
+         swift_multiPayloadEnumGeneric_destructiveInjectEnumTag,
+         C_CC, AlwaysAvailable,
+         RETURNS(VoidTy),
+         ARGS(Int8PtrTy, Int32Ty, TypeMetadataPtrTy),
+         ATTRS(NoUnwind, WillReturn),
+         EFFECT(NoEffect))
+
 // unsigned swift_singlePayloadEnumGeneric_getEnumTag(swift::OpaqueValue *address,
 //                                                    const Metadata *metadata);
 FUNCTION(SinglePayloadEnumGenericGetEnumTag,

lib/IRGen/GenValueWitness.cpp

@@ -937,12 +937,13 @@ getDestructiveInjectEnumTagFunction(IRGenModule &IGM,
                                     const EnumTypeLayoutEntry *typeLayoutEntry,
                                     GenericSignature genericSig) {
   if ((!typeLayoutEntry->layoutString(IGM, genericSig) &&
-       !isRuntimeInstatiatedLayoutString(IGM, typeLayoutEntry)) ||
-      typeLayoutEntry->isSingleton()) {
+       !isRuntimeInstatiatedLayoutString(IGM, typeLayoutEntry))) {
+    return nullptr;
+  } else if (typeLayoutEntry->isSingleton()) {
     return nullptr;
   } else if (!typeLayoutEntry->isFixedSize(IGM)) {
     if (typeLayoutEntry->isMultiPayloadEnum()) {
-      return nullptr;
+      return IGM.getMultiPayloadEnumGenericDestructiveInjectEnumTagFn();
     } else {
       return IGM.getSinglePayloadEnumGenericDestructiveInjectEnumTagFn();
     }

stdlib/public/runtime/BytecodeLayouts.cpp

@@ -746,6 +746,41 @@ swift_multiPayloadEnumGeneric_getEnumTag(swift::OpaqueValue *address,
   }
 }
 
+extern "C" void swift_multiPayloadEnumGeneric_destructiveInjectEnumTag(
+    swift::OpaqueValue *address, unsigned tag, const Metadata *metadata) {
+  auto addr = reinterpret_cast<uint8_t *>(address);
+  LayoutStringReader reader{metadata->getLayoutString(),
+                            layoutStringHeaderSize + sizeof(uint64_t)};
+
+  auto numTagBytes = reader.readBytes<size_t>();
+  auto numPayloads = reader.readBytes<size_t>();
+  reader.skip(sizeof(size_t));
+  auto enumSize = reader.readBytes<size_t>();
+  auto payloadSize = enumSize - numTagBytes;
+
+  if (tag < numPayloads) {
+    // For a payload case, store the tag after the payload area.
+    auto tagBytes = addr + payloadSize;
+    storeEnumElement(tagBytes, tag, numTagBytes);
+  } else {
+    // For an empty case, factor out the parts that go in the payload and
+    // tag areas.
+    unsigned whichEmptyCase = tag - numPayloads;
+    unsigned whichTag, whichPayloadValue;
+    if (payloadSize >= 4) {
+      whichTag = numPayloads;
+      whichPayloadValue = whichEmptyCase;
+    } else {
+      unsigned numPayloadBits = payloadSize * CHAR_BIT;
+      whichTag = numPayloads + (whichEmptyCase >> numPayloadBits);
+      whichPayloadValue = whichEmptyCase & ((1U << numPayloadBits) - 1U);
+    }
+    auto tagBytes = addr + payloadSize;
+    storeEnumElement(tagBytes, whichTag, numTagBytes);
+    storeEnumElement(addr, whichPayloadValue, payloadSize);
+  }
+}
+
 template <typename T>
 static inline T handleSinglePayloadEnumGenericTag(
     LayoutStringReader &reader, uint8_t *addr,

stdlib/public/runtime/BytecodeLayouts.h

@@ -130,6 +130,9 @@ SWIFT_RUNTIME_EXPORT
 unsigned swift_multiPayloadEnumGeneric_getEnumTag(swift::OpaqueValue *address,
                                                   const Metadata *metadata);
 SWIFT_RUNTIME_EXPORT
+void swift_multiPayloadEnumGeneric_destructiveInjectEnumTag(
+    swift::OpaqueValue *address, unsigned tag, const Metadata *metadata);
+SWIFT_RUNTIME_EXPORT
 unsigned swift_singlePayloadEnumGeneric_getEnumTag(swift::OpaqueValue *address,
                                                    const Metadata *metadata);
 SWIFT_RUNTIME_EXPORT

test/Interpreter/layout_string_witnesses_dynamic.swift

@@ -618,6 +618,34 @@ func testResilientSinglePayloadEnumGenericInjectTag() {
 
 testResilientSinglePayloadEnumGenericInjectTag()
 
+@inline(never)
+func matchResilientMultiPayloadEnumGenericTag(_ x: ResilientMultiPayloadEnumGeneric<AnyObject>) -> Int {
+    return switch x {
+    case .nonEmpty0: 0
+    case .nonEmpty1: 1
+    case .empty0: 2
+    case .empty1: 3
+    }
+}
+
+func testResilientMultiPayloadEnumGenericInjectTag() {
+    let x = ResilientMultiPayloadEnumGeneric<AnyObject>.nonEmpty0(SimpleClass(x: 23))
+    let y = ResilientMultiPayloadEnumGeneric<AnyObject>.nonEmpty1(SimpleClass(x: 32))
+    let z = ResilientMultiPayloadEnumGeneric<AnyObject>.empty0
+    let w = ResilientMultiPayloadEnumGeneric<AnyObject>.empty1
+
+    // CHECK: Enum case: 0
+    print("Enum case: \(matchResilientMultiPayloadEnumGenericTag(x))")
+    // CHECK: Enum case: 1
+    print("Enum case: \(matchResilientMultiPayloadEnumGenericTag(y))")
+    // CHECK: Enum case: 2
+    print("Enum case: \(matchResilientMultiPayloadEnumGenericTag(z))")
+    // CHECK: Enum case: 3
+    print("Enum case: \(matchResilientMultiPayloadEnumGenericTag(w))")
+}
+
+testResilientMultiPayloadEnumGenericInjectTag()
+
 #if os(macOS)
 
 import Foundation