Support destructiveInjectEnumTag in generic single payload enums with layout strings

include/swift/Runtime/RuntimeFunctions.def

@@ -2333,7 +2333,7 @@ FUNCTION(SingletonEnumGetEnumTag,
          C_CC, AlwaysAvailable,
          RETURNS(Int32Ty),
          ARGS(Int8PtrTy, TypeMetadataPtrTy),
-         ATTRS(NoUnwind),
+         ATTRS(NoUnwind, WillReturn),
          EFFECT(NoEffect))
 
 // unsigned swift_enumSimple_getEnumTag(swift::OpaqueValue *address,
@@ -2343,7 +2343,7 @@ FUNCTION(EnumSimpleGetEnumTag,
          C_CC, AlwaysAvailable,
          RETURNS(Int32Ty),
          ARGS(Int8PtrTy, TypeMetadataPtrTy),
-         ATTRS(NoUnwind),
+         ATTRS(NoUnwind, WillReturn),
          EFFECT(NoEffect))
 
 // unsigned swift_enumFn_getEnumTag(swift::OpaqueValue *address,
@@ -2353,7 +2353,7 @@ FUNCTION(EnumFnGetEnumTag,
          C_CC, AlwaysAvailable,
          RETURNS(Int32Ty),
          ARGS(Int8PtrTy, TypeMetadataPtrTy),
-         ATTRS(NoUnwind),
+         ATTRS(NoUnwind, WillReturn),
          EFFECT(NoEffect))
 
 // unsigned swift_multiPayloadEnumGeneric_getEnumTag(opaque* address,
@@ -2363,7 +2363,7 @@ FUNCTION(MultiPayloadEnumGenericGetEnumTag,
          C_CC, AlwaysAvailable,
          RETURNS(Int32Ty),
          ARGS(Int8PtrTy, TypeMetadataPtrTy),
-         ATTRS(NoUnwind),
+         ATTRS(NoUnwind, WillReturn),
          EFFECT(NoEffect))
 
 // unsigned swift_singlePayloadEnumGeneric_getEnumTag(swift::OpaqueValue *address,
@@ -2373,7 +2373,18 @@ FUNCTION(SinglePayloadEnumGenericGetEnumTag,
          C_CC, AlwaysAvailable,
          RETURNS(Int32Ty),
          ARGS(Int8PtrTy, TypeMetadataPtrTy),
-         ATTRS(NoUnwind),
+         ATTRS(NoUnwind, WillReturn),
+         EFFECT(NoEffect))
+
+// void swift_singlePayloadEnumGeneric_destructiveInjectEnumTag(swift::OpaqueValue *address,
+//                                                              unsigned tag,
+//                                                              const Metadata *metadata)
+FUNCTION(SinglePayloadEnumGenericDestructiveInjectEnumTag,
+         swift_singlePayloadEnumGeneric_destructiveInjectEnumTag,
+         C_CC, AlwaysAvailable,
+         RETURNS(VoidTy),
+         ARGS(Int8PtrTy, Int32Ty, TypeMetadataPtrTy),
+         ATTRS(NoUnwind, WillReturn),
          EFFECT(NoEffect))
 
 // void swift_generic_instantiateLayoutString(const uint8_t* layoutStr,

lib/IRGen/GenValueWitness.cpp

@@ -931,6 +931,38 @@ static llvm::Constant *getEnumTagFunction(IRGenModule &IGM,
   }
 }
 
+static llvm::Constant *
+getDestructiveInjectEnumTagFunction(IRGenModule &IGM,
+                                    const EnumTypeLayoutEntry *typeLayoutEntry,
+                                    GenericSignature genericSig) {
+  if ((!typeLayoutEntry->layoutString(IGM, genericSig) &&
+       !isRuntimeInstatiatedLayoutString(IGM, typeLayoutEntry)) ||
+      typeLayoutEntry->isSingleton()) {
+    return nullptr;
+  } else if (!typeLayoutEntry->isFixedSize(IGM)) {
+    if (typeLayoutEntry->isMultiPayloadEnum()) {
+      return nullptr;
+    } else {
+      return IGM.getSinglePayloadEnumGenericDestructiveInjectEnumTagFn();
+    }
+  } else if (typeLayoutEntry->isMultiPayloadEnum()) {
+    return nullptr;
+  } else {
+    auto &payloadTI = **(typeLayoutEntry->cases[0]->getFixedTypeInfo());
+    auto mask = payloadTI.getFixedExtraInhabitantMask(IGM);
+    auto tzCount = mask.countTrailingZeros();
+    auto shiftedMask = mask.lshr(tzCount);
+    auto toCount = shiftedMask.countTrailingOnes();
+    if (payloadTI.mayHaveExtraInhabitants(IGM) &&
+        (mask.countPopulation() > 64 || toCount != mask.countPopulation() ||
+         (tzCount % toCount != 0))) {
+      return nullptr;
+    } else {
+      return nullptr;
+    }
+  }
+}
+
 static bool
 valueWitnessRequiresCopyability(ValueWitness index) {
   switch (index) {
@@ -1189,8 +1221,31 @@ addValueWitness(IRGenModule &IGM, ConstantStructBuilder &B, ValueWitness index,
     }
     goto standard;
   }
+  case ValueWitness::DestructiveInjectEnumTag: {
+    assert(concreteType.getEnumOrBoundGenericEnum());
+    if (IGM.Context.LangOpts.hasFeature(Feature::LayoutStringValueWitnesses) &&
+        IGM.getOptions().EnableLayoutStringValueWitnesses) {
+      auto ty = boundGenericCharacteristics
+                    ? boundGenericCharacteristics->concreteType
+                    : concreteType;
+      auto &typeInfo = boundGenericCharacteristics
+                           ? *boundGenericCharacteristics->TI
+                           : concreteTI;
+      if (auto *typeLayoutEntry = typeInfo.buildTypeLayoutEntry(
+              IGM, ty, /*useStructLayouts*/ true)) {
+        if (auto *enumLayoutEntry = typeLayoutEntry->getAsEnum()) {
+          auto genericSig = concreteType.getNominalOrBoundGenericNominal()
+                                ->getGenericSignature();
+          if (auto *fn = getDestructiveInjectEnumTagFunction(
+                  IGM, enumLayoutEntry, genericSig)) {
+            return addFunction(fn);
+          }
+        }
+      }
+    }
+    goto standard;
+  }
   case ValueWitness::DestructiveProjectEnumData:
-  case ValueWitness::DestructiveInjectEnumTag:
     assert(concreteType.getEnumOrBoundGenericEnum());
     goto standard;
   }

stdlib/public/runtime/BytecodeLayouts.cpp

@@ -751,29 +751,29 @@ static inline T handleSinglePayloadEnumGenericTag(
     LayoutStringReader &reader, uint8_t *addr,
     std::function<std::optional<T>(const Metadata *, size_t, uint8_t)>
         extraTagBytesHandler,
-    std::function<T(const Metadata *, unsigned, unsigned)> xiHandler) {
+    std::function<T(const Metadata *, unsigned, unsigned, size_t, uint8_t)>
+        xiHandler) {
   auto tagBytesAndOffset = reader.readBytes<uint64_t>();
   auto extraTagBytesPattern = (uint8_t)(tagBytesAndOffset >> 62);
   auto xiTagBytesOffset =
       tagBytesAndOffset & std::numeric_limits<uint32_t>::max();
-  const Metadata *xiType = nullptr;
+  auto numExtraTagBytes = 1 << (extraTagBytesPattern - 1);
+  auto payloadSize = reader.readBytes<size_t>();
+  auto xiType = reader.readBytes<const Metadata *>();
 
   if (extraTagBytesPattern) {
-    auto numExtraTagBytes = 1 << (extraTagBytesPattern - 1);
-    auto payloadSize = reader.readBytes<size_t>();
-    xiType = reader.readBytes<const Metadata *>();
     if (auto result =
             extraTagBytesHandler(xiType, payloadSize, numExtraTagBytes)) {
       return *result;
     }
   } else {
     reader.skip(sizeof(size_t));
-    xiType = reader.readBytes<const Metadata *>();
   }
 
   auto numEmptyCases = reader.readBytes<unsigned>();
 
-  return xiHandler(xiType, xiTagBytesOffset, numEmptyCases);
+  return xiHandler(xiType, xiTagBytesOffset, numEmptyCases, payloadSize,
+                   numExtraTagBytes);
 }
 
 extern "C" unsigned
@@ -803,7 +803,8 @@ swift_singlePayloadEnumGeneric_getEnumTag(swift::OpaqueValue *address,
   };
 
   auto xihandler = [addr](const Metadata *xiType, unsigned xiTagBytesOffset,
-                          unsigned numEmptyCases) -> unsigned {
+                          unsigned numEmptyCases, size_t payloadSize,
+                          uint8_t numExtraTagBytes) -> unsigned {
     if (xiType) {
       return xiType->vw_getEnumTagSinglePayload(
           (const OpaqueValue *)(addr + xiTagBytesOffset), numEmptyCases);
@@ -816,6 +817,64 @@ swift_singlePayloadEnumGeneric_getEnumTag(swift::OpaqueValue *address,
       reader, addr, extraTagBytesHandler, xihandler);
 }
 
+extern "C" void swift_singlePayloadEnumGeneric_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 =
+      [=](const Metadata *xiType, size_t payloadSize,
+          uint8_t numExtraTagBytes) -> std::optional<bool> {
+    unsigned payloadNumExtraInhabitants =
+        xiType ? xiType->vw_getNumExtraInhabitants() : 0;
+    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 = [=](const Metadata *xiType, unsigned xiTagBytesOffset,
+                       unsigned numEmptyCases, size_t payloadSize,
+                       uint8_t numExtraTagBytes) -> bool {
+    unsigned payloadNumExtraInhabitants =
+        xiType ? xiType->vw_getNumExtraInhabitants() : 0;
+    if (tag <= payloadNumExtraInhabitants) {
+      if (numExtraTagBytes != 0)
+        storeEnumElement(addr + payloadSize, 0, numExtraTagBytes);
+
+      if (tag == 0)
+        return true;
+
+      xiType->vw_storeEnumTagSinglePayload(
+          (swift::OpaqueValue *)(addr + xiTagBytesOffset), tag, numEmptyCases);
+    }
+    return true;
+  };
+
+  handleSinglePayloadEnumGenericTag<bool>(reader, addr, extraTagBytesHandler,
+                                          xihandler);
+}
+
 extern "C" swift::OpaqueValue *
 swift_generic_initializeBufferWithCopyOfBuffer(swift::ValueBuffer *dest,
                                                swift::ValueBuffer *src,

stdlib/public/runtime/BytecodeLayouts.h

@@ -133,6 +133,9 @@ SWIFT_RUNTIME_EXPORT
 unsigned swift_singlePayloadEnumGeneric_getEnumTag(swift::OpaqueValue *address,
                                                    const Metadata *metadata);
 SWIFT_RUNTIME_EXPORT
+void swift_singlePayloadEnumGeneric_destructiveInjectEnumTag(
+    swift::OpaqueValue *address, unsigned tag, const Metadata *metadata);
+SWIFT_RUNTIME_EXPORT
 void swift_generic_instantiateLayoutString(const uint8_t *layoutStr,
                                            Metadata *type);
 

test/Interpreter/Inputs/layout_string_witnesses_types_resilient.swift

@@ -22,7 +22,7 @@ public struct GenericResilient<C, T> {
 public enum ResilientSinglePayloadEnumGeneric<T> {
     case empty0
     case empty1
-    case nonEmpty0(T)
+    case nonEmpty(T)
 }
 
 public enum ResilientMultiPayloadEnumGeneric<T> {

test/Interpreter/layout_string_witnesses_dynamic.swift

@@ -569,7 +569,7 @@ testResilientMultiPayloadEnumTag()
 
 func testResilientSinglePayloadEnumGenericTag() {
     let x = switch getResilientSinglePayloadEnumGenericEmpty0(AnyObject.self) {
-    case .nonEmpty0: 0
+    case .nonEmpty: 0
     case .empty0: 1
     case .empty1: 2
     }
@@ -594,6 +594,30 @@ func testResilientMultiPayloadEnumGenericTag() {
 
 testResilientMultiPayloadEnumGenericTag()
 
+@inline(never)
+func matchResilientSinglePayloadEnumGenericTag(_ x: ResilientSinglePayloadEnumGeneric<AnyObject>) -> Int {
+    return switch x {
+    case .nonEmpty: 0
+    case .empty0: 1
+    case .empty1: 2
+    }
+}
+
+func testResilientSinglePayloadEnumGenericInjectTag() {
+    let x = ResilientSinglePayloadEnumGeneric<AnyObject>.nonEmpty(SimpleClass(x: 23))
+    let y = ResilientSinglePayloadEnumGeneric<AnyObject>.empty0
+    let z = ResilientSinglePayloadEnumGeneric<AnyObject>.empty1
+
+    // CHECK: Enum case: 0
+    print("Enum case: \(matchResilientSinglePayloadEnumGenericTag(x))")
+    // CHECK: Enum case: 1
+    print("Enum case: \(matchResilientSinglePayloadEnumGenericTag(y))")
+    // CHECK: Enum case: 2
+    print("Enum case: \(matchResilientSinglePayloadEnumGenericTag(z))")
+}
+
+testResilientSinglePayloadEnumGenericInjectTag()
+
 #if os(macOS)
 
 import Foundation