Merge pull request #67171 from drexin/wip-sp-simp-inject

Support destructiveInjectEnumTag in simple single payload enums with layout strings

Author: drexin

include/swift/Runtime/RuntimeFunctions.def

@@ -2346,6 +2346,17 @@ FUNCTION(EnumSimpleGetEnumTag,
          ATTRS(NoUnwind, WillReturn),
          EFFECT(NoEffect))
 
+// unsigned swift_enumSimple_destructiveInjectEnumTag(swift::OpaqueValue *address,
+//                                                    unsigned tag,
+//                                                    const Metadata *metadata)
+FUNCTION(EnumSimpleDestructiveInjectEnumTag,
+         swift_enumSimple_destructiveInjectEnumTag,
+         C_CC, AlwaysAvailable,
+         RETURNS(VoidTy),
+         ARGS(Int8PtrTy, Int32Ty, TypeMetadataPtrTy),
+         ATTRS(NoUnwind, WillReturn),
+         EFFECT(NoEffect))
+
 // unsigned swift_enumFn_getEnumTag(swift::OpaqueValue *address,
 //                                  const Metadata *metadata);
 FUNCTION(EnumFnGetEnumTag,

lib/IRGen/GenValueWitness.cpp

@@ -960,7 +960,7 @@ getDestructiveInjectEnumTagFunction(IRGenModule &IGM,
          (tzCount % toCount != 0))) {
       return nullptr;
     } else {
-      return nullptr;
+      return IGM.getEnumSimpleDestructiveInjectEnumTagFn();
     }
   }
 }

stdlib/public/runtime/BytecodeLayouts.cpp

@@ -729,6 +729,62 @@ extern "C" unsigned swift_enumSimple_getEnumTag(swift::OpaqueValue *address,
       reader, addr, extraTagBytesHandler, xihandler);
 }
 
+extern "C" void swift_enumSimple_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 extraTagBytesHandler =
+      [addr, tag](size_t payloadNumExtraInhabitants, size_t payloadSize,
+                  uint8_t numExtraTagBytes) -> std::optional<bool> {
+    if (tag <= payloadNumExtraInhabitants) {
+      return std::nullopt;
+    }
+
+    unsigned noPayloadIndex = tag - 1;
+    unsigned caseIndex = noPayloadIndex - payloadNumExtraInhabitants;
+    unsigned payloadIndex, extraTagIndex;
+    if (payloadSize >= 4) {
+      extraTagIndex = 1;
+      payloadIndex = caseIndex;
+    } else {
+      unsigned payloadBits = payloadSize * 8U;
+      extraTagIndex = 1U + (caseIndex >> payloadBits);
+      payloadIndex = caseIndex & ((1U << payloadBits) - 1U);
+    }
+
+    // Store into the value.
+    if (payloadSize)
+      storeEnumElement(addr, payloadIndex, payloadSize);
+    if (numExtraTagBytes)
+      storeEnumElement(addr + payloadSize, extraTagIndex, numExtraTagBytes);
+
+    return true;
+  };
+
+  auto xihandler = [addr, tag](size_t payloadNumExtraInhabitants,
+                               uint64_t zeroTagValue, uint8_t xiTagBytesPattern,
+                               unsigned xiTagBytesOffset, size_t payloadSize,
+                               uint8_t numExtraTagBytes) -> bool {
+    auto xiTagBytes = 1 << (xiTagBytesPattern - 1);
+    if (tag <= payloadNumExtraInhabitants) {
+      if (numExtraTagBytes != 0)
+        storeEnumElement(addr + payloadSize, 0, numExtraTagBytes);
+
+      if (tag == 0)
+        return true;
+
+      storeEnumElement(addr + xiTagBytesOffset, tag - 1 + zeroTagValue,
+                       xiTagBytes);
+    }
+    return true;
+  };
+
+  handleSinglePayloadEnumSimpleTag<bool>(reader, addr, extraTagBytesHandler,
+                                         xihandler);
+}
+
 extern "C"
 unsigned swift_enumFn_getEnumTag(swift::OpaqueValue *address,
                                  const Metadata *metadata) {

stdlib/public/runtime/BytecodeLayouts.h

@@ -124,6 +124,10 @@ SWIFT_RUNTIME_EXPORT
 unsigned swift_enumSimple_getEnumTag(swift::OpaqueValue *address,
                                      const Metadata *metadata);
 SWIFT_RUNTIME_EXPORT
+void swift_enumSimple_destructiveInjectEnumTag(swift::OpaqueValue *address,
+                                               unsigned tag,
+                                               const Metadata *metadata);
+SWIFT_RUNTIME_EXPORT
 unsigned swift_enumFn_getEnumTag(swift::OpaqueValue *address,
                                  const Metadata *metadata);
 SWIFT_RUNTIME_EXPORT

test/Interpreter/layout_string_witnesses_dynamic.swift

@@ -594,6 +594,30 @@ func testResilientMultiPayloadEnumGenericTag() {
 
 testResilientMultiPayloadEnumGenericTag()
 
+@inline(never)
+func matchResilientSinglePayloadEnumSimpleTag(_ x: ResilientSinglePayloadEnumSimple) -> Int {
+    return switch x {
+    case .nonEmpty: 0
+    case .empty0: 1
+    case .empty1: 2
+    }
+}
+
+func testResilientSinglePayloadEnumSimpleInjectTag() {
+    let x = ResilientSinglePayloadEnumSimple.nonEmpty(SimpleClass(x: 23))
+    let y = ResilientSinglePayloadEnumSimple.empty0
+    let z = ResilientSinglePayloadEnumSimple.empty1
+
+    // CHECK: Enum case: 0
+    print("Enum case: \(matchResilientSinglePayloadEnumSimpleTag(x))")
+    // CHECK: Enum case: 1
+    print("Enum case: \(matchResilientSinglePayloadEnumSimpleTag(y))")
+    // CHECK: Enum case: 2
+    print("Enum case: \(matchResilientSinglePayloadEnumSimpleTag(z))")
+}
+
+testResilientSinglePayloadEnumSimpleInjectTag()
+
 @inline(never)
 func matchResilientSinglePayloadEnumGenericTag(_ x: ResilientSinglePayloadEnumGeneric<AnyObject>) -> Int {
     return switch x {