[IRGen+Runtime] Add getEnumTag for layout strings on simple single payload enums (#66941)

Author: drexin

include/swift/Runtime/RuntimeFunctions.def

@@ -2315,6 +2315,16 @@ FUNCTION(GenericInitWithTake,
          ATTRS(NoUnwind),
          EFFECT(Refcounting))
 
+// unsigned swift_enumSimple_getEnumTag(swift::OpaqueValue *address,
+//                                      const Metadata *metadata);
+FUNCTION(EnumSimpleGetEnumTag,
+         swift_enumSimple_getEnumTag,
+         C_CC, AlwaysAvailable,
+         RETURNS(Int32Ty),
+         ARGS(Int8PtrTy, TypeMetadataPtrTy),
+         ATTRS(NoUnwind),
+         EFFECT(NoEffect))
+
 // unsigned swift_enumFn_getEnumTag(swift::OpaqueValue *address,
 //                                  const Metadata *metadata);
 FUNCTION(EnumFnGetEnumTag,

lib/IRGen/GenValueWitness.cpp

@@ -900,8 +900,11 @@ static bool isRuntimeInstatiatedLayoutString(IRGenModule &IGM,
 }
 
 static llvm::Constant *getEnumTagFunction(IRGenModule &IGM,
-                                   const EnumTypeLayoutEntry *typeLayoutEntry) {
-  if (typeLayoutEntry->isSingleton()) {
+                                     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()) {
@@ -923,7 +926,7 @@ static llvm::Constant *getEnumTagFunction(IRGenModule &IGM,
          (tzCount % toCount != 0))) {
       return IGM.getEnumFnGetEnumTagFn();
     } else {
-      return nullptr;
+      return IGM.getEnumSimpleGetEnumTagFn();
     }
   }
 }
@@ -1156,7 +1159,9 @@ static void addValueWitness(IRGenModule &IGM, ConstantStructBuilder &B,
       if (auto *typeLayoutEntry = typeInfo.buildTypeLayoutEntry(
               IGM, ty, /*useStructLayouts*/ true)) {
         if (auto *enumLayoutEntry = typeLayoutEntry->getAsEnum()) {
-          if (auto *fn = getEnumTagFunction(IGM, enumLayoutEntry)) {
+          auto genericSig = concreteType.getNominalOrBoundGenericNominal()
+                              ->getGenericSignature();
+          if (auto *fn = getEnumTagFunction(IGM, enumLayoutEntry, genericSig)) {
             return addFunction(fn);
           }
         }

stdlib/public/runtime/BytecodeLayouts.cpp

@@ -652,6 +652,53 @@ swift_generic_assignWithTake(swift::OpaqueValue *dest, swift::OpaqueValue *src,
   return swift_generic_initWithTake(dest, src, metadata);
 }
 
+extern "C"
+unsigned swift_enumSimple_getEnumTag(swift::OpaqueValue *address,
+                                     const Metadata *metadata) {
+  auto addr = reinterpret_cast<uint8_t *>(address);
+  LayoutStringReader reader{metadata->getLayoutString(),
+                            layoutStringHeaderSize + sizeof(uint64_t)};
+  auto byteCountsAndOffset = reader.readBytes<uint64_t>();
+  auto extraTagBytesPattern = (uint8_t)(byteCountsAndOffset >> 62);
+  auto xiTagBytesPattern = ((uint8_t)(byteCountsAndOffset >> 59)) & 0x7;
+  auto xiTagBytesOffset =
+      byteCountsAndOffset & std::numeric_limits<uint32_t>::max();
+
+  if (extraTagBytesPattern) {
+    auto extraTagBytes = 1 << (extraTagBytesPattern - 1);
+    auto payloadSize = reader.readBytes<size_t>();
+    auto tagBytes = readTagBytes(addr + payloadSize, extraTagBytes);
+    if (tagBytes) {
+      reader.skip(sizeof(uint64_t));
+      auto payloadNumExtraInhabitants = reader.readBytes<size_t>();
+      unsigned caseIndexFromExtraTagBits =
+          payloadSize >= 4 ? 0 : (tagBytes - 1U) << (payloadSize * 8U);
+      unsigned caseIndexFromValue = loadEnumElement(addr, payloadSize);
+      unsigned noPayloadIndex =
+          (caseIndexFromExtraTagBits | caseIndexFromValue) +
+          payloadNumExtraInhabitants;
+      return noPayloadIndex + 1;
+    }
+  } else {
+    reader.skip(sizeof(size_t));
+  }
+
+  if (xiTagBytesPattern) {
+    auto zeroTagValue = reader.readBytes<uint64_t>();
+    auto xiTagValues = reader.readBytes<size_t>();
+
+    auto xiTagBytes = 1 << (xiTagBytesPattern - 1);
+    uint64_t tagBytes =
+        readTagBytes(addr + xiTagBytesOffset, xiTagBytes) -
+        zeroTagValue;
+    if (tagBytes < xiTagValues) {
+      return tagBytes + 1;
+    }
+  }
+
+  return 0;
+}
+
 extern "C"
 unsigned swift_enumFn_getEnumTag(swift::OpaqueValue *address,
                                  const Metadata *metadata) {

stdlib/public/runtime/BytecodeLayouts.h

@@ -113,6 +113,9 @@ swift::OpaqueValue *swift_generic_initWithTake(swift::OpaqueValue *dest,
                                                swift::OpaqueValue *src,
                                                const Metadata *metadata);
 SWIFT_RUNTIME_EXPORT
+unsigned swift_enumSimple_getEnumTag(swift::OpaqueValue *address,
+                                     const Metadata *metadata);
+SWIFT_RUNTIME_EXPORT
 unsigned swift_enumFn_getEnumTag(swift::OpaqueValue *address,
                                  const Metadata *metadata);
 SWIFT_RUNTIME_EXPORT

test/Interpreter/Inputs/layout_string_witnesses_types_resilient.swift

@@ -45,6 +45,12 @@ public enum ResilientSinglePayloadEnumComplex {
     case nonEmpty(ResilientMultiPayloadEnum)
 }
 
+public enum ResilientSinglePayloadEnumSimple {
+    case empty0
+    case empty1
+    case nonEmpty(AnyObject)
+}
+
 public func getResilientSinglePayloadEnumGenericEmpty0<T>(_ t: T.Type) -> ResilientSinglePayloadEnumGeneric<T> {
     return .empty0
 }
@@ -60,3 +66,7 @@ public func getResilientMultiPayloadEnumEmpty0() -> ResilientMultiPayloadEnum {
 public func getResilientSinglePayloadEnumComplexEmpty0() -> ResilientSinglePayloadEnumComplex {
     return .empty0
 }
+
+public func getResilientSinglePayloadEnumSimpleEmpty0() -> ResilientSinglePayloadEnumSimple {
+    return .empty0
+}

test/Interpreter/layout_string_witnesses_dynamic.swift

@@ -516,6 +516,19 @@ func testGenericSinglePayloadEnumManyXI() {
 
 testGenericSinglePayloadEnumManyXI()
 
+func testResilientSinglePayloadEnumSimpleTag() {
+    let x = switch getResilientSinglePayloadEnumSimpleEmpty0() {
+    case .nonEmpty: 0
+    case .empty0: 1
+    case .empty1: 2
+    }
+
+    // CHECK: Enum case: 1
+    print("Enum case: \(x)")
+}
+
+testResilientSinglePayloadEnumSimpleTag()
+
 func testResilientSinglePayloadEnumComplexTag() {
     let x = switch getResilientSinglePayloadEnumComplexEmpty0() {
     case .nonEmpty: 0