[IRGen+Runtime] Layout string getEnumTag for fixed size enums subset (#66899)

* [IRGen+Runtime] Layout string getEnumTag for fixed size enums subset

getEnumTag impl for layout strings of fixed sized enums that use a function to fetch the enum tag

* Fix potential UB in IRGen

Author: drexin

include/swift/Runtime/RuntimeFunctions.def

@@ -2315,6 +2315,16 @@ FUNCTION(GenericInitWithTake,
          ATTRS(NoUnwind),
          EFFECT(Refcounting))
 
+// unsigned swift_enumFn_getEnumTag(swift::OpaqueValue *address,
+//                                  const Metadata *metadata);
+FUNCTION(EnumFnGetEnumTag,
+         swift_enumFn_getEnumTag,
+         C_CC, AlwaysAvailable,
+         RETURNS(Int32Ty),
+         ARGS(Int8PtrTy, TypeMetadataPtrTy),
+         ATTRS(NoUnwind),
+         EFFECT(NoEffect))
+
 // unsigned swift_multiPayloadEnumGeneric_getEnumTag(opaque* address,
 //                                                   const Metadata *type);
 FUNCTION(MultiPayloadEnumGenericGetEnumTag,
@@ -2325,7 +2335,8 @@ FUNCTION(MultiPayloadEnumGenericGetEnumTag,
          ATTRS(NoUnwind),
          EFFECT(NoEffect))
 
-// void swift_generic_instantiateLayoutString(const uint8_t* layoutStr, Metadata* type);
+// void swift_generic_instantiateLayoutString(const uint8_t* layoutStr,
+//                                            Metadata* type);
 FUNCTION(GenericInstantiateLayoutString,
          swift_generic_instantiateLayoutString,
          C_CC, AlwaysAvailable,

lib/IRGen/GenValueWitness.cpp

@@ -885,8 +885,8 @@ void addStride(ConstantStructBuilder &B, const TypeInfo *TI, IRGenModule &IGM) {
 }
 } // end anonymous namespace
 
-bool isRuntimeInstatiatedLayoutString(IRGenModule &IGM,
-                                      const TypeLayoutEntry *typeLayoutEntry) {
+static bool isRuntimeInstatiatedLayoutString(IRGenModule &IGM,
+                                       const TypeLayoutEntry *typeLayoutEntry) {
   if (IGM.Context.LangOpts.hasFeature(Feature::LayoutStringValueWitnesses) &&
       IGM.Context.LangOpts.hasFeature(
           Feature::LayoutStringValueWitnessesInstantiation) &&
@@ -899,6 +899,36 @@ bool isRuntimeInstatiatedLayoutString(IRGenModule &IGM,
   return false;
 }
 
+static llvm::Constant *getEnumTagFunction(IRGenModule &IGM,
+                                   const EnumTypeLayoutEntry *typeLayoutEntry) {
+  if (typeLayoutEntry->isSingleton()) {
+    return nullptr;
+  } else if (!typeLayoutEntry->isFixedSize(IGM)) {
+    if (typeLayoutEntry->isMultiPayloadEnum()) {
+      return IGM.getMultiPayloadEnumGenericGetEnumTagFn();
+    } else {
+      // TODO: implement support for single payload generic enums
+      return nullptr;
+    }
+  } else if (typeLayoutEntry->isMultiPayloadEnum()) {
+    return IGM.getEnumFnGetEnumTagFn();
+  } 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 IGM.getEnumFnGetEnumTagFn();
+    } else {
+      return nullptr;
+    }
+  }
+}
+
 static bool
 valueWitnessRequiresCopyability(ValueWitness index) {
   switch (index) {
@@ -1127,9 +1157,8 @@ static void addValueWitness(IRGenModule &IGM, ConstantStructBuilder &B,
       if (auto *typeLayoutEntry = typeInfo.buildTypeLayoutEntry(
               IGM, ty, /*useStructLayouts*/ true)) {
         if (auto *enumLayoutEntry = typeLayoutEntry->getAsEnum()) {
-          if (enumLayoutEntry->isMultiPayloadEnum() &&
-              !typeLayoutEntry->isFixedSize(IGM)) {
-            return addFunction(IGM.getMultiPayloadEnumGenericGetEnumTagFn());
+          if (auto *fn = getEnumTagFunction(IGM, enumLayoutEntry)) {
+            return addFunction(fn);
           }
         }
       }

stdlib/public/runtime/BytecodeLayouts.cpp

@@ -652,6 +652,17 @@ swift_generic_assignWithTake(swift::OpaqueValue *dest, swift::OpaqueValue *src,
   return swift_generic_initWithTake(dest, src, metadata);
 }
 
+extern "C"
+unsigned swift_enumFn_getEnumTag(swift::OpaqueValue *address,
+                                 const Metadata *metadata) {
+  auto addr = reinterpret_cast<const uint8_t *>(address);
+  LayoutStringReader reader{metadata->getLayoutString(),
+                            layoutStringHeaderSize + sizeof(uint64_t)};
+  auto getEnumTag = readRelativeFunctionPointer<GetEnumTagFn>(reader);
+
+  return getEnumTag(addr);
+}
+
 extern "C" unsigned
 swift_multiPayloadEnumGeneric_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_enumFn_getEnumTag(swift::OpaqueValue *address,
+                                 const Metadata *metadata);
+SWIFT_RUNTIME_EXPORT
 unsigned swift_multiPayloadEnumGeneric_getEnumTag(swift::OpaqueValue *address,
                                                   const Metadata *metadata);
 SWIFT_RUNTIME_EXPORT

test/Interpreter/Inputs/layout_string_witnesses_types_resilient.swift

@@ -19,13 +19,34 @@ public struct GenericResilient<C, T> {
     }
 }
 
-public enum ResilientMultiPayloadEnum<T> {
+public enum ResilientMultiPayloadEnumGeneric<T> {
     case empty0
     case empty1
     case nonEmpty0(AnyObject)
     case nonEmpty1(T)
 }
 
-public func getResilientMultiPayloadEnumEmpty0<T>(_ t: T.Type) -> ResilientMultiPayloadEnum<T> {
+public enum ResilientMultiPayloadEnum {
+    case empty0
+    case empty1
+    case nonEmpty0(AnyObject, Bool)
+    case nonEmpty1(AnyObject)
+}
+
+public enum ResilientSinglePayloadEnumComplex {
+    case empty0
+    case empty1
+    case nonEmpty(ResilientMultiPayloadEnum)
+}
+
+public func getResilientMultiPayloadEnumGenericEmpty0<T>(_ t: T.Type) -> ResilientMultiPayloadEnumGeneric<T> {
+    return .empty0
+}
+
+public func getResilientMultiPayloadEnumEmpty0() -> ResilientMultiPayloadEnum {
+    return .empty0
+}
+
+public func getResilientSinglePayloadEnumComplexEmpty0() -> ResilientSinglePayloadEnumComplex {
     return .empty0
 }

test/Interpreter/layout_string_witnesses_dynamic.swift

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