[IRGen+Runtime] Layout string getEnumTag for generic multi payload enums (#66872)

We already store all necessary information to extract the tag in the layout string, so we should utilize it for the getEnumTag witness

Author: drexin

include/swift/Runtime/RuntimeFunctions.def

@@ -2315,6 +2315,16 @@ FUNCTION(GenericInitWithTake,
          ATTRS(NoUnwind),
          EFFECT(Refcounting))
 
+// unsigned swift_multiPayloadEnumGeneric_getEnumTag(opaque* address,
+//                                                   const Metadata *type);
+FUNCTION(MultiPayloadEnumGenericGetEnumTag,
+         swift_multiPayloadEnumGeneric_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,
          swift_generic_instantiateLayoutString,

lib/IRGen/GenValueWitness.cpp

@@ -1113,7 +1113,29 @@ static void addValueWitness(IRGenModule &IGM, ConstantStructBuilder &B,
     goto standard;
   }
 
-  case ValueWitness::GetEnumTag:
+  case ValueWitness::GetEnumTag: {
+    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()) {
+          if (enumLayoutEntry->isMultiPayloadEnum() &&
+              !typeLayoutEntry->isFixedSize(IGM)) {
+            return addFunction(IGM.getMultiPayloadEnumGenericGetEnumTagFn());
+          }
+        }
+      }
+    }
+    goto standard;
+  }
   case ValueWitness::DestructiveProjectEnumData:
   case ValueWitness::DestructiveInjectEnumTag:
     assert(concreteType.getEnumOrBoundGenericEnum());

stdlib/public/runtime/BytecodeLayouts.cpp

@@ -18,6 +18,7 @@
 
 #include "BytecodeLayouts.h"
 #include "../SwiftShims/swift/shims/HeapObject.h"
+#include "EnumImpl.h"
 #include "WeakReference.h"
 #include "swift/ABI/MetadataValues.h"
 #include "swift/ABI/System.h"
@@ -160,7 +161,7 @@ inline static void handleRefCounts(const Metadata *metadata, Params... params) {
                               std::forward<Params>(params)...);
 }
 
-static uint64_t readTagBytes(uint8_t *addr, uint8_t byteCount) {
+static uint64_t readTagBytes(const uint8_t *addr, uint8_t byteCount) {
   switch (byteCount) {
   case 1:
     return addr[0];
@@ -651,6 +652,35 @@ swift_generic_assignWithTake(swift::OpaqueValue *dest, swift::OpaqueValue *src,
   return swift_generic_initWithTake(dest, src, metadata);
 }
 
+extern "C" unsigned
+swift_multiPayloadEnumGeneric_getEnumTag(swift::OpaqueValue *address,
+                                         const Metadata *metadata) {
+  auto addr = reinterpret_cast<const uint8_t *>(address);
+  LayoutStringReader reader{metadata->getLayoutString(),
+                            layoutStringHeaderSize + sizeof(uint64_t)};
+
+  auto tagBytes = reader.readBytes<size_t>();
+  auto numPayloads = reader.readBytes<size_t>();
+  reader.skip(sizeof(size_t));
+  auto enumSize = reader.readBytes<size_t>();
+  auto payloadSize = enumSize - tagBytes;
+
+  auto enumTag = (unsigned)readTagBytes(addr + payloadSize, tagBytes);
+  if (enumTag < numPayloads) {
+    return enumTag;
+  }
+
+  auto payloadValue = loadEnumElement(addr, payloadSize);
+
+  if (payloadSize >= 4) {
+    return numPayloads + payloadValue;
+  } else {
+    unsigned numPayloadBits = payloadSize * CHAR_BIT;
+    return (payloadValue | (enumTag - numPayloads) << numPayloadBits) +
+           numPayloads;
+  }
+}
+
 void swift::swift_resolve_resilientAccessors(uint8_t *layoutStr,
                                              size_t layoutStrOffset,
                                              const uint8_t *fieldLayoutStr,

stdlib/public/runtime/BytecodeLayouts.h

@@ -94,17 +94,30 @@ struct LayoutStringWriter {
 };
 
 SWIFT_RUNTIME_EXPORT
-void swift_generic_destroy(swift::OpaqueValue *address, const Metadata *metadata);
+void swift_generic_destroy(swift::OpaqueValue *address,
+                           const Metadata *metadata);
 SWIFT_RUNTIME_EXPORT
-swift::OpaqueValue *swift_generic_assignWithCopy(swift::OpaqueValue *dest, swift::OpaqueValue *src, const Metadata *metadata);
+swift::OpaqueValue *swift_generic_assignWithCopy(swift::OpaqueValue *dest,
+                                                 swift::OpaqueValue *src,
+                                                 const Metadata *metadata);
 SWIFT_RUNTIME_EXPORT
-swift::OpaqueValue *swift_generic_assignWithTake(swift::OpaqueValue *dest, swift::OpaqueValue *src, const Metadata *metadata);
+swift::OpaqueValue *swift_generic_assignWithTake(swift::OpaqueValue *dest,
+                                                 swift::OpaqueValue *src,
+                                                 const Metadata *metadata);
 SWIFT_RUNTIME_EXPORT
-swift::OpaqueValue *swift_generic_initWithCopy(swift::OpaqueValue *dest, swift::OpaqueValue *src, const Metadata *metadata);
+swift::OpaqueValue *swift_generic_initWithCopy(swift::OpaqueValue *dest,
+                                               swift::OpaqueValue *src,
+                                               const Metadata *metadata);
 SWIFT_RUNTIME_EXPORT
-swift::OpaqueValue *swift_generic_initWithTake(swift::OpaqueValue *dest, swift::OpaqueValue *src, const Metadata *metadata);
+swift::OpaqueValue *swift_generic_initWithTake(swift::OpaqueValue *dest,
+                                               swift::OpaqueValue *src,
+                                               const Metadata *metadata);
 SWIFT_RUNTIME_EXPORT
-void swift_generic_instantiateLayoutString(const uint8_t *layoutStr, Metadata *type);
+unsigned swift_multiPayloadEnumGeneric_getEnumTag(swift::OpaqueValue *address,
+                                                  const Metadata *metadata);
+SWIFT_RUNTIME_EXPORT
+void swift_generic_instantiateLayoutString(const uint8_t *layoutStr,
+                                           Metadata *type);
 
 void swift_resolve_resilientAccessors(uint8_t *layoutStr,
                                       size_t layoutStrOffset,

test/Interpreter/Inputs/layout_string_witnesses_types_resilient.swift

@@ -17,4 +17,15 @@ public struct GenericResilient<C, T> {
         self.x = x
         self.y = y
     }
-}
+}
+
+public enum ResilientMultiPayloadEnum<T> {
+    case empty0
+    case empty1
+    case nonEmpty0(AnyObject)
+    case nonEmpty1(T)
+}
+
+public func getResilientMultiPayloadEnumEmpty0<T>(_ t: T.Type) -> ResilientMultiPayloadEnum<T> {
+    return .empty0
+}

test/Interpreter/layout_string_witnesses_dynamic.swift

@@ -516,6 +516,20 @@ func testGenericSinglePayloadEnumManyXI() {
 
 testGenericSinglePayloadEnumManyXI()
 
+func testResilientMultiPayloadEnumTag() {
+    let x = switch getResilientMultiPayloadEnumEmpty0(AnyObject.self) {
+    case .nonEmpty0: 0
+    case .nonEmpty1: 1
+    case .empty0: 2
+    case .empty1: 3
+    }
+
+    // CHECK: Enum case: 2
+    print("Enum case: \(x)")
+}
+
+testResilientMultiPayloadEnumTag()
+
 #if os(macOS)
 
 import Foundation