[IRGen+Runtime] Add getEnumTag for layout strings on generic single payload enum (swiftlang#66911)

Author: drexin

include/swift/Runtime/RuntimeFunctions.def

@@ -2335,6 +2335,16 @@ FUNCTION(MultiPayloadEnumGenericGetEnumTag,
          ATTRS(NoUnwind),
          EFFECT(NoEffect))
 
+// unsigned swift_singlePayloadEnumGeneric_getEnumTag(swift::OpaqueValue *address,
+//                                                    const Metadata *metadata);
+FUNCTION(SinglePayloadEnumGenericGetEnumTag,
+         swift_singlePayloadEnumGeneric_getEnumTag,
+         C_CC, AlwaysAvailable,
+         RETURNS(Int32Ty),
+         ARGS(Int8PtrTy, TypeMetadataPtrTy),
+         ATTRS(NoUnwind),
+         EFFECT(NoEffect))
+
 // void swift_generic_instantiateLayoutString(const uint8_t* layoutStr,
 //                                            Metadata* type);
 FUNCTION(GenericInstantiateLayoutString,

lib/IRGen/GenValueWitness.cpp

@@ -908,7 +908,7 @@ static llvm::Constant *getEnumTagFunction(IRGenModule &IGM,
       return IGM.getMultiPayloadEnumGenericGetEnumTagFn();
     } else {
       // TODO: implement support for single payload generic enums
-      return nullptr;
+      return IGM.getSinglePayloadEnumGenericGetEnumTagFn();
     }
   } else if (typeLayoutEntry->isMultiPayloadEnum()) {
     return IGM.getEnumFnGetEnumTagFn();

stdlib/public/runtime/BytecodeLayouts.cpp

@@ -692,6 +692,51 @@ swift_multiPayloadEnumGeneric_getEnumTag(swift::OpaqueValue *address,
   }
 }
 
+extern "C" unsigned
+swift_singlePayloadEnumGeneric_getEnumTag(swift::OpaqueValue *address,
+                                          const Metadata *metadata) {
+  auto addr = reinterpret_cast<const uint8_t *>(address);
+  LayoutStringReader reader{metadata->getLayoutString(),
+                            layoutStringHeaderSize + sizeof(uint64_t)};
+
+  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;
+
+  if (extraTagBytesPattern) {
+    auto extraTagBytes = 1 << (extraTagBytesPattern - 1);
+    auto payloadSize = reader.readBytes<size_t>();
+    auto tagBytes = readTagBytes(addr + payloadSize, extraTagBytes);
+    if (tagBytes) {
+      xiType = reader.readBytes<const Metadata *>();
+      unsigned payloadNumExtraInhabitants =
+          xiType ? xiType->vw_getNumExtraInhabitants() : 0;
+      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));
+  }
+
+  xiType = reader.readBytes<const Metadata *>();
+
+  if (xiType) {
+    auto numEmptyCases = reader.readBytes<unsigned>();
+
+    return xiType->vw_getEnumTagSinglePayload(
+        (const OpaqueValue *)(addr + xiTagBytesOffset), numEmptyCases);
+  }
+
+  return 0;
+}
+
 void swift::swift_resolve_resilientAccessors(uint8_t *layoutStr,
                                              size_t layoutStrOffset,
                                              const uint8_t *fieldLayoutStr,

stdlib/public/runtime/BytecodeLayouts.h

@@ -119,6 +119,9 @@ SWIFT_RUNTIME_EXPORT
 unsigned swift_multiPayloadEnumGeneric_getEnumTag(swift::OpaqueValue *address,
                                                   const Metadata *metadata);
 SWIFT_RUNTIME_EXPORT
+unsigned swift_singlePayloadEnumGeneric_getEnumTag(swift::OpaqueValue *address,
+                                                   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

@@ -19,6 +19,12 @@ public struct GenericResilient<C, T> {
     }
 }
 
+public enum ResilientSinglePayloadEnumGeneric<T> {
+    case empty0
+    case empty1
+    case nonEmpty0(T)
+}
+
 public enum ResilientMultiPayloadEnumGeneric<T> {
     case empty0
     case empty1
@@ -39,6 +45,10 @@ public enum ResilientSinglePayloadEnumComplex {
     case nonEmpty(ResilientMultiPayloadEnum)
 }
 
+public func getResilientSinglePayloadEnumGenericEmpty0<T>(_ t: T.Type) -> ResilientSinglePayloadEnumGeneric<T> {
+    return .empty0
+}
+
 public func getResilientMultiPayloadEnumGenericEmpty0<T>(_ t: T.Type) -> ResilientMultiPayloadEnumGeneric<T> {
     return .empty0
 }

test/Interpreter/layout_string_witnesses_dynamic.swift

@@ -543,6 +543,19 @@ func testResilientMultiPayloadEnumTag() {
 
 testResilientMultiPayloadEnumTag()
 
+func testResilientSinglePayloadEnumGenericTag() {
+    let x = switch getResilientSinglePayloadEnumGenericEmpty0(AnyObject.self) {
+    case .nonEmpty0: 0
+    case .empty0: 1
+    case .empty1: 2
+    }
+
+    // CHECK: Enum case: 1
+    print("Enum case: \(x)")
+}
+
+testResilientSinglePayloadEnumGenericTag()
+
 func testResilientMultiPayloadEnumGenericTag() {
     let x = switch getResilientMultiPayloadEnumGenericEmpty0(AnyObject.self) {
     case .nonEmpty0: 0