[5.0] IRGen: Fix extra inhabitants of multi-payload enums with more spare bits than tag bits.

lib/IRGen/GenEnum.cpp

@@ -4844,6 +4844,31 @@ namespace {
       }
       return addr;
     }
+
+    // If there are common spare bits we didn't use for tags, rotate the
+    // extra inhabitant values so that the used tag bits are at the bottom.
+    // This will cleanly separate the used tag values from the extra inhabitants
+    // so we can discriminate them with one comparison. The tag favors high
+    // bits, whereas extra inhabitants count down from -1 using all bits
+    // (capping out at up to 32 spare bits, in which case the lowest 32
+    // bits are used).
+    std::pair<unsigned, unsigned> getRotationAmountsForExtraInhabitants() const{
+      assert([&]{
+        auto maskedBits = PayloadTagBits;
+        maskedBits &= CommonSpareBits;
+        return maskedBits == PayloadTagBits;
+      }());
+
+      unsigned commonSpareBitsCount = CommonSpareBits.count();
+      unsigned payloadTagBitsCount = PayloadTagBits.count();
+      if (commonSpareBitsCount == payloadTagBitsCount
+          || commonSpareBitsCount - payloadTagBitsCount >= 32) {
+        return std::make_pair(0, 0);
+      }
+      unsigned shlAmount = commonSpareBitsCount - payloadTagBitsCount;
+      unsigned shrAmount = std::min(commonSpareBitsCount, 32u) - shlAmount;
+      return {shlAmount, shrAmount};
+    }
 
     llvm::Value *getExtraInhabitantIndex(IRGenFunction &IGF,
                                          Address src,
@@ -4859,6 +4884,32 @@ namespace {
         auto payload = EnumPayload::load(IGF, projectPayload(IGF, src),
                                          PayloadSchema);
         tag = payload.emitGatherSpareBits(IGF, CommonSpareBits, 0, 32);
+
+        // If there are common spare bits we didn't use for tags, rotate the
+        // tag value so that the used tag bits are at the bottom. This will
+        // cleanly separate the used tag values from the extra inhabitants so
+        // we can discriminate them with one comparison. The tag favors high
+        // bits, whereas extra inhabitants count down from -1 using all bits
+        // (capping out at up to 32 spare bits, in which case the lowest 32
+        // bits are used).
+        //
+        // Note that since this is the inverse operation--we're taking the bits
+        // out of a payload and mapping them back to an extra inhabitant index--
+        // the `shr` and `shl` amounts are intentionally swapped here.
+        unsigned shrAmount, shlAmount;
+        std::tie(shrAmount, shlAmount) = getRotationAmountsForExtraInhabitants();
+        if (shrAmount != 0) {
+          assert(getExtraTagBitCountForExtraInhabitants() == 0);
+          auto tagLo = IGF.Builder.CreateLShr(tag, shrAmount);
+          auto tagHi = IGF.Builder.CreateShl(tag, shlAmount);
+          tag = IGF.Builder.CreateOr(tagLo, tagHi);
+          if (CommonSpareBits.count() < 32) {
+            auto mask = llvm::ConstantInt::get(IGF.IGM.Int32Ty,
+                                          (1u << CommonSpareBits.count()) - 1u);
+            tag = IGF.Builder.CreateAnd(tag, mask);
+          }
+        }
+
         if (getExtraTagBitCountForExtraInhabitants()) {
           auto extraTagAddr = projectExtraTagBitsForExtraInhabitants(IGF, src);
           auto extraTag = IGF.Builder.CreateLoad(extraTagAddr);
@@ -4899,6 +4950,29 @@ namespace {
       }
 
       auto indexValue = IGF.Builder.CreateNot(index);
+
+      // If there are common spare bits we didn't use for tags, rotate the
+      // tag value so that the used tag bits are at the bottom. This will
+      // cleanly separate the used tag values from the extra inhabitants so
+      // we can discriminate them with one comparison. The tag favors high
+      // bits, whereas extra inhabitants count down from -1 using all bits
+      // (capping out at up to 32 spare bits, in which case the lowest 32
+      // bits are used).
+      unsigned shlAmount, shrAmount;
+      std::tie(shlAmount, shrAmount) = getRotationAmountsForExtraInhabitants();
+
+      if (shlAmount != 0) {
+        assert(getExtraTagBitCountForExtraInhabitants() == 0);
+        if (CommonSpareBits.count() < 32) {
+          auto mask = llvm::ConstantInt::get(IGF.IGM.Int32Ty,
+                                         (1u << CommonSpareBits.count()) - 1u);
+          indexValue = IGF.Builder.CreateAnd(indexValue, mask);
+        }
+        auto indexValueHi = IGF.Builder.CreateShl(indexValue, shlAmount);
+        auto indexValueLo = IGF.Builder.CreateLShr(indexValue, shrAmount);
+        indexValue = IGF.Builder.CreateOr(indexValueHi, indexValueLo);
+      }
+
       if (CommonSpareBits.count()) {
         // Factor the index value into parts to scatter into the payload and
         // to store in the extra tag bits, if any.
@@ -4952,6 +5026,25 @@ namespace {
       // Count down from all-ones since a small negative number constant is
       // likely to be easier to reify.
       auto mask = ~index;
+
+      // If there are common spare bits we didn't use for tags, rotate the
+      // tag value so that the used tag bits are at the bottom. This will
+      // cleanly separate the used tag values from the extra inhabitants so
+      // we can discriminate them with one comparison. The tag favors high
+      // bits, whereas extra inhabitants count down from -1 using all bits
+      // (capping out at up to 32 spare bits, in which case the lowest 32
+      // bits are used).
+      unsigned shlAmount, shrAmount;
+      std::tie(shlAmount, shrAmount) = getRotationAmountsForExtraInhabitants();
+
+      if (shlAmount != 0) {
+        assert(getExtraTagBitCountForExtraInhabitants() == 0);
+        if (CommonSpareBits.count() < 32) {
+          mask &= (1u << CommonSpareBits.count()) - 1;
+        }
+        mask = (mask >> shrAmount) | (mask << shlAmount);
+      }
+
       auto extraTagMask = getExtraTagBitCountForExtraInhabitants() >= 32
         ? ~0u : (1 << getExtraTagBitCountForExtraInhabitants()) - 1;
 

test/Interpreter/multi_payload_extra_inhabitant.swift

@@ -1,6 +1,6 @@
 // RUN: %empty-directory(%t)
 
-// RUN: %target-build-swift -parse-stdlib -Xfrontend -verify-type-layout -Xfrontend SpareBitExtraInhabitants -Xfrontend -verify-type-layout -Xfrontend SpareBitSingleExtraInhabitant -Xfrontend -verify-type-layout -Xfrontend SpareBitNoExtraInhabitant -Xfrontend -verify-type-layout -Xfrontend SpareBitNoExtraInhabitant2 -Xfrontend -verify-type-layout -Xfrontend TwoTagExtraInhabitants -Xfrontend -verify-type-layout -Xfrontend ThreeTagExtraInhabitants -Xfrontend -verify-type-layout -Xfrontend NoTagExtraInhabitants -Xfrontend -verify-type-layout -Xfrontend DynamicExtraInhabitantsNever -Xfrontend -verify-type-layout -Xfrontend DynamicExtraInhabitantsZeroBytes -Xfrontend -verify-type-layout -Xfrontend DynamicExtraInhabitantsOneByte -Xfrontend -verify-type-layout -Xfrontend DynamicExtraInhabitantsTwoBytes -O -o %t/a.out %s
+// RUN: %target-build-swift -parse-stdlib -Xfrontend -verify-type-layout -Xfrontend SpareBitExtraInhabitants -Xfrontend -verify-type-layout -Xfrontend SpareBitSingleExtraInhabitant -Xfrontend -verify-type-layout -Xfrontend SpareBitNoExtraInhabitant -Xfrontend -verify-type-layout -Xfrontend SpareBitNoExtraInhabitant2 -Xfrontend -verify-type-layout -Xfrontend TwoTagExtraInhabitants -Xfrontend -verify-type-layout -Xfrontend ThreeTagExtraInhabitants -Xfrontend -verify-type-layout -Xfrontend NoTagExtraInhabitants -Xfrontend -verify-type-layout -Xfrontend DynamicExtraInhabitantsNever -Xfrontend -verify-type-layout -Xfrontend DynamicExtraInhabitantsZeroBytes -Xfrontend -verify-type-layout -Xfrontend DynamicExtraInhabitantsOneByte -Xfrontend -verify-type-layout -Xfrontend DynamicExtraInhabitantsTwoBytes -Xfrontend -verify-type-layout -Xfrontend MoreSpareBitsThanTagsExtraInhabitants -Xfrontend -verify-type-layout -Xfrontend OptOptMoreSpareBitsThanTagsExtraInhabitants -O -o %t/a.out %s
 // RUN: %target-run %t/a.out 2>&1
 
 // Type layout verifier is only compiled into the runtime in asserts builds.
@@ -50,6 +50,21 @@ enum ThreeTagExtraInhabitants {
   case c(Builtin.Int32)
 }
 
+enum MoreSpareBitsThanTagsExtraInhabitants {
+  case a(Builtin.Int29)
+  case b(Builtin.Int29)
+  case c(Builtin.Int29)
+  case d(Builtin.Int29)
+}
+typealias OptOptMoreSpareBitsThanTagsExtraInhabitants =
+  Optional<Optional<MoreSpareBitsThanTagsExtraInhabitants>>
+
+enum MoreSpareBitsThanTagsExtraInhabitants2 {
+  case a(Builtin.Int29)
+  case b(Builtin.Int29)
+  case c(Builtin.Int29)
+}
+
 enum NoTagExtraInhabitants {
   case aaa(Builtin.Int32), aab(Builtin.Int32), aac(Builtin.Int32), aad(Builtin.Int32), aae(Builtin.Int32), aaf(Builtin.Int32), aag(Builtin.Int32), aah(Builtin.Int32)
   case aba(Builtin.Int32), abb(Builtin.Int32), abc(Builtin.Int32), abd(Builtin.Int32), abe(Builtin.Int32), abf(Builtin.Int32), abg(Builtin.Int32), abh(Builtin.Int32)
@@ -201,4 +216,37 @@ tests.test("types that have at least two extra inhabitants") {
   expectHasAtLeastTwoExtraInhabitants(DynamicExtraInhabitantsOneByte.self, nil: nil, someNil: .some(nil))
   expectHasAtLeastTwoExtraInhabitants(DynamicExtraInhabitantsTwoBytes.self, nil: nil, someNil: .some(nil))
 }
+tests.test("types with more spare bits than used by tags") {
+  expectHasAtLeastTwoExtraInhabitants(MoreSpareBitsThanTagsExtraInhabitants.self,
+                                      nil: nil, someNil: .some(nil))
+
+  for x in [MoreSpareBitsThanTagsExtraInhabitants.a(Builtin.zeroInitializer()),
+            MoreSpareBitsThanTagsExtraInhabitants.b(Builtin.zeroInitializer()),
+            MoreSpareBitsThanTagsExtraInhabitants.c(Builtin.zeroInitializer()),
+            MoreSpareBitsThanTagsExtraInhabitants.d(Builtin.zeroInitializer())]{
+    let opt = Optional(x)
+    expectNotNil(opt)
+    let opt2 = Optional(opt)
+    expectNotNil(opt2)
+    let opt3 = Optional(opt2)
+    expectNotNil(opt3)
+    let opt4 = Optional(opt3)
+    expectNotNil(opt4)
+  }
+
+  for x in [MoreSpareBitsThanTagsExtraInhabitants.a(Builtin.zeroInitializer()),
+            MoreSpareBitsThanTagsExtraInhabitants.b(Builtin.zeroInitializer()),
+            MoreSpareBitsThanTagsExtraInhabitants.c(Builtin.zeroInitializer())]{
+    let opt = Optional(x)
+    expectNotNil(opt)
+    let opt2 = Optional(opt)
+    expectNotNil(opt2)
+    let opt3 = Optional(opt2)
+    expectNotNil(opt3)
+    let opt4 = Optional(opt3)
+    expectNotNil(opt4)
+    let opt5 = Optional(opt4)
+    expectNotNil(opt5)
+  }
+}
 runAllTests()