[IRGen] Disable simple single payload enum in layout strings (#70529)

drexin · web-flow · commit b11899de606f · 2023-12-20T19:46:08.000-08:00
rdar://119792426

There are a few issues with wrong assumptions around extra inhabitants that cause tags to not be identified properly in some cases. Until a proper fix is identified, we emit tag functions instead.
diff --git a/lib/IRGen/GenEnum.cpp b/lib/IRGen/GenEnum.cpp
@@ -282,6 +282,12 @@ EnumImplStrategy::emitResilientTagIndices(IRGenModule &IGM) const {
   }
 }
 
+llvm::Value *
+EnumImplStrategy::emitFixedGetEnumTag(IRGenFunction &IGF, SILType T, Address enumAddr) const {
+  assert(TIK >= Fixed);
+  return emitGetEnumTag(IGF, T, enumAddr);
+}
+
 namespace {
   /// Implementation strategy for singleton enums, with zero or one cases.
   class SingletonEnumImplStrategy final : public EnumImplStrategy {
@@ -1908,13 +1914,32 @@ namespace {
           llvm::ConstantInt::get(IGF.IGM.Int32Ty, ElementsWithNoPayload.size());
 
       auto PayloadT = getPayloadType(IGF.IGM, T);
+
       auto opaqueAddr = Address(
           IGF.Builder.CreateBitCast(enumAddr.getAddress(), IGF.IGM.OpaquePtrTy),
           IGF.IGM.OpaqueTy, enumAddr.getAlignment());
+
       return emitGetEnumTagSinglePayloadCall(IGF, PayloadT, numEmptyCases,
                                              opaqueAddr);
     }
 
+
+    llvm::Value *emitFixedGetEnumTag(IRGenFunction &IGF, SILType T,
+                                Address enumAddr) const override {
+      assert(TIK >= Fixed);
+      auto numEmptyCases =
+          llvm::ConstantInt::get(IGF.IGM.Int32Ty, ElementsWithNoPayload.size());
+
+      auto PayloadT = getPayloadType(IGF.IGM, T);
+
+      auto &fixedTI = getFixedPayloadTypeInfo();
+      auto addr = IGF.Builder.CreateBitCast(
+          enumAddr.getAddress(), fixedTI.getStorageType()->getPointerTo());
+      return fixedTI.getEnumTagSinglePayload(IGF, numEmptyCases,
+                                             fixedTI.getAddressForPointer(addr),
+                                             PayloadT, /*isOutlined*/ false);
+    }
+
     /// The payload for a single-payload enum is always placed in front and
     /// will never have interleaved tag bits, so we can just bitcast the enum
     /// address to the payload type for either injection or projection of the
diff --git a/lib/IRGen/GenEnum.h b/lib/IRGen/GenEnum.h
@@ -256,6 +256,14 @@ class EnumImplStrategy {
                                       SILType T,
                                       Address enumAddr) const = 0;
 
+  /// Return the enum case tag for the given value. Payload cases come first,
+  /// followed by non-payload cases. Used for the getEnumTag value witness.
+  ///
+  /// Only ever called for fixed types.
+  virtual llvm::Value *emitFixedGetEnumTag(IRGenFunction &IGF,
+                                           SILType T,
+                                           Address enumAddr) const;
+
   /// Project the address of the data for a case. Does not check or modify
   /// the referenced enum value.
   /// Corresponds to the SIL 'init_enum_data_addr' instruction.
diff --git a/lib/IRGen/GenValueWitness.cpp b/lib/IRGen/GenValueWitness.cpp
@@ -923,15 +923,15 @@ static llvm::Constant *getEnumTagFunction(IRGenModule &IGM,
     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))) {
+    // auto toCount = shiftedMask.countTrailingOnes();
+    // if (payloadTI.mayHaveExtraInhabitants(IGM) &&
+    //     (mask.popcount() > 64 ||
+    //      toCount != mask.popcount() ||
+    //      (tzCount % toCount != 0))) {
       return IGM.getEnumFnGetEnumTagFn();
-    } else {
-      return IGM.getEnumSimpleGetEnumTagFn();
-    }
+    // } else {
+    //   return IGM.getEnumSimpleGetEnumTagFn();
+    // }
   }
 }
 
@@ -960,14 +960,14 @@ getDestructiveInjectEnumTagFunction(IRGenModule &IGM,
     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))) {
+    // auto toCount = shiftedMask.countTrailingOnes();
+    // if (payloadTI.mayHaveExtraInhabitants(IGM) &&
+    //     (mask.popcount() > 64 || toCount != mask.popcount() ||
+    //      (tzCount % toCount != 0))) {
       return nullptr;
-    } else {
-      return IGM.getEnumSimpleDestructiveInjectEnumTagFn();
-    }
+    // } else {
+    //   return IGM.getEnumSimpleDestructiveInjectEnumTagFn();
+    // }
   }
 }
 
diff --git a/lib/IRGen/TypeLayout.cpp b/lib/IRGen/TypeLayout.cpp
@@ -520,7 +520,7 @@ llvm::Function *createFixedEnumLoadTag(IRGenModule &IGM,
         auto enumAddr = typeInfo->getAddressForPointer(castEnumPtr);
 
         auto &strategy = getEnumImplStrategy(IGM, entry.ty);
-        auto tag = strategy.emitGetEnumTag(IGF, entry.ty, enumAddr);
+        auto tag = strategy.emitFixedGetEnumTag(IGF, entry.ty, enumAddr);
         IGF.Builder.CreateRet(tag);
       });
 
@@ -2233,7 +2233,7 @@ bool EnumTypeLayoutEntry::buildSinglePayloadRefCountString(
   uint64_t zeroTagValue = 0;
   unsigned xiBitCount = 0;
   unsigned xiBitOffset = 0;
-  bool isSimple = true;
+  bool isSimple = false;
 
   auto &payloadTI = **cases[0]->getFixedTypeInfo();
 
@@ -2248,17 +2248,17 @@ bool EnumTypeLayoutEntry::buildSinglePayloadRefCountString(
 
       auto tzCount = mask.countTrailingZeros();
       auto shiftedMask = mask.lshr(tzCount);
-      auto toCount = shiftedMask.countTrailingOnes();
-      if (mask.countPopulation() > 64 || toCount != mask.countPopulation() ||
-          (tzCount % toCount != 0)) {
+      // auto toCount = shiftedMask.countTrailingOnes();
+      // if (mask.popcount() > 64 || toCount != mask.popcount() ||
+      //     (tzCount % toCount != 0)) {
         // We currently don't handle cases with non-contiguous or > 64 bits of
         // extra inhabitants
         isSimple = false;
-      } else {
-        xiBitCount = std::min(64u, mask.countPopulation());
-        xiBitOffset = mask.countTrailingZeros();
-        zeroTagValue = lowValue.extractBitsAsZExtValue(xiBitCount, xiBitOffset);
-      }
+      // } else {
+      //   xiBitCount = std::min(64u, mask.popcount());
+      //   xiBitOffset = mask.countTrailingZeros();
+      //   zeroTagValue = lowValue.extractBitsAsZExtValue(xiBitCount, xiBitOffset);
+      // }
     }
   }
 
diff --git a/test/Interpreter/Inputs/layout_string_witnesses_types_resilient.swift b/test/Interpreter/Inputs/layout_string_witnesses_types_resilient.swift
@@ -56,6 +56,25 @@ public enum ResilientSinglePayloadEnumSimple {
     case nonEmpty(AnyObject)
 }
 
+public enum MyString {
+    case immortal(UInt64)
+    case s(AnyObject)
+    case t(AnyObject)
+}
+
+public enum ResilientSinglePayloadEnumSimpleMultiExtraTagPayload {
+    case empty0
+    case empty1
+    case empty2
+    case empty3
+    case empty4
+    case empty5
+    case empty6
+    case empty7
+    case empty8
+    case nonEmpty(MyString)
+}
+
 public enum ResilientSingletonEnum {
     case nonEmpty(AnyObject)
 }
@@ -92,6 +111,10 @@ public func getResilientSingletonEnumNonEmpty(_ x: AnyObject) -> ResilientSingle
     return .nonEmpty(x)
 }
 
+public func getResilientSinglePayloadEnumSimpleMultiExtraTagPayloadEmpty3() -> ResilientSinglePayloadEnumSimpleMultiExtraTagPayload {
+    return .empty3
+}
+
 public enum ResilientSinglePayloadEnum {
     case empty0
     case empty1
diff --git a/test/Interpreter/layout_string_witnesses_dynamic.swift b/test/Interpreter/layout_string_witnesses_dynamic.swift
@@ -842,6 +842,26 @@ func testResilientSinglePayloadEnumSimpleTag() {
 
 testResilientSinglePayloadEnumSimpleTag()
 
+func testResilientSinglePayloadEnumSimpleTagMultiExtraTagPayload() {
+    let x = switch getResilientSinglePayloadEnumSimpleMultiExtraTagPayloadEmpty3() {
+    case .nonEmpty: 0
+    case .empty0: 1
+    case .empty1: 2
+    case .empty2: 3
+    case .empty3: 4
+    case .empty4: 5
+    case .empty5: 6
+    case .empty6: 7
+    case .empty7: 8
+    case .empty8: 9
+    }
+
+    // CHECK: Enum case: 4
+    print("Enum case: \(x)")
+}
+
+testResilientSinglePayloadEnumSimpleTagMultiExtraTagPayload()
+
 func testResilientSinglePayloadEnumIndirectTag() {
     let x = switch getResilientSinglePayloadEnumIndirectNonEmpty(SimpleClass(x: 23)) {
     case .nonEmpty: 0
