Merge pull request #71428 from augusto2112/nfc-sbm

[NFC] Move MultiPayloadEnum spare bits mask calculation into a function

Author: augusto2112

lib/IRGen/GenEnum.cpp

@@ -5836,6 +5836,32 @@ namespace {
       tagBits.append(APInt(extraTagSize, (1U << ExtraTagBitCount) - 1U));
       return tagBits.build();
     }
+
+    std::optional<SpareBitsMaskInfo> calculateSpareBitsMask() const override {
+      SpareBitVector spareBits;
+      for (auto enumCase : getElementsWithPayload()) {
+        cast<FixedTypeInfo>(enumCase.ti)
+            ->applyFixedSpareBitsMask(IGM, spareBits);
+      }
+      // Trim leading/trailing zero bytes, then pad to a multiple of 32 bits
+      llvm::APInt bits = spareBits.asAPInt();
+      uint32_t byteOffset = bits.countTrailingZeros() / 8;
+      bits.lshrInPlace(byteOffset * 8); // Trim zero bytes from bottom end
+
+      auto bitsInMask = bits.getActiveBits(); // Ignore high-order zero bits
+      uint32_t bytesInMask = (bitsInMask + 7) / 8;
+      auto wordsInMask = (bytesInMask + 3) / 4;
+      bits = bits.zextOrTrunc(wordsInMask * 32);
+
+      // Never write an MPE descriptor bigger than 16k
+      // The runtime will fall back on its own internal
+      // spare bits calculation for this (very rare) case.
+      if (bytesInMask > 16384) {
+        return {};
+      }
+
+      return {{bits, byteOffset, bytesInMask}};
+    }
   };
 
   class ResilientEnumImplStrategy final

lib/IRGen/GenEnum.h

@@ -518,6 +518,20 @@ class EnumImplStrategy {
     return false;
   }
 
+  struct SpareBitsMaskInfo {
+    const llvm::APInt bits;
+    const uint32_t byteOffset;
+    const uint32_t bytesInMask;
+
+    uint64_t wordsInMask() const { return (bytesInMask + 3) / 4; }
+  };
+
+  /// Calculates the spare bits mask for the enum. Returns none if the type
+  /// should not emit the spare bits.
+  virtual std::optional<SpareBitsMaskInfo> calculateSpareBitsMask() const {
+    return {};
+  };
+
 private:
   EnumImplStrategy(const EnumImplStrategy &) = delete;
   EnumImplStrategy &operator=(const EnumImplStrategy &) = delete;

lib/IRGen/GenReflection.cpp

@@ -1125,54 +1125,44 @@ class MultiPayloadEnumDescriptorBuilder : public ReflectionMetadataBuilder {
   void layout() override {
     auto &strategy = getEnumImplStrategy(IGM, typeInContext);
     bool isMPE = strategy.getElementsWithPayload().size() > 1;
-    assert(isMPE && "Cannot emit Multi-Payload Enum data for an enum that doesn't have multiple payloads");
+    assert(isMPE && "Cannot emit Multi-Payload Enum data for an enum that "
+                    "doesn't have multiple payloads");
 
     const TypeInfo &TI = strategy.getTypeInfo();
     auto fixedTI = dyn_cast<FixedTypeInfo>(&TI);
-    assert(fixedTI != nullptr
-           && "MPE reflection records can only be emitted for fixed-layout enums");
+    assert(fixedTI != nullptr &&
+           "MPE reflection records can only be emitted for fixed-layout enums");
 
-    // Get the spare bits mask for the enum payloads.
-    SpareBitVector spareBits;
-    for (auto enumCase : strategy.getElementsWithPayload()) {
-      cast<FixedTypeInfo>(enumCase.ti)->applyFixedSpareBitsMask(IGM, spareBits);
-    }
-
-    // Trim leading/trailing zero bytes, then pad to a multiple of 32 bits
-    llvm::APInt bits = spareBits.asAPInt();
-    uint32_t byteOffset = bits.countTrailingZeros() / 8;
-    bits.lshrInPlace(byteOffset * 8); // Trim zero bytes from bottom end
-
-    auto bitsInMask = bits.getActiveBits(); // Ignore high-order zero bits
-    auto usesPayloadSpareBits = bitsInMask > 0;
-    uint32_t bytesInMask = (bitsInMask + 7) / 8;
-    auto wordsInMask = (bytesInMask + 3) / 4;
-    bits = bits.zextOrTrunc(wordsInMask * 32);
+    auto spareBitsMaskInfo = strategy.calculateSpareBitsMask();
 
     // Never write an MPE descriptor bigger than 16k
     // The runtime will fall back on its own internal
     // spare bits calculation for this (very rare) case.
-    if (bytesInMask > 16384) {
+    if (!spareBitsMaskInfo)
       return;
-    }
+
+    auto bits = spareBitsMaskInfo->bits;
 
     addTypeRef(type, CanGenericSignature());
 
+    bool usesPayloadSpareBits = spareBitsMaskInfo->bytesInMask > 0;
+
     // MPE record contents are a multiple of 32-bits
     uint32_t contentsSizeInWords = 1; /* Size + flags is mandatory */
-    if (wordsInMask > 0) {
-      contentsSizeInWords +=
-        1 /* SpareBits byte count */
-        + wordsInMask;
+    if (usesPayloadSpareBits) {
+      contentsSizeInWords += 1 /* SpareBits byte count */
+                             + spareBitsMaskInfo->wordsInMask();
     }
+
     uint32_t flags = usesPayloadSpareBits ? 1 : 0;
 
     B.addInt32((contentsSizeInWords << 16) | flags);
 
-    if (bytesInMask > 0) {
-      B.addInt32((byteOffset << 16) | bytesInMask);
+    if (usesPayloadSpareBits) {
+      B.addInt32((spareBitsMaskInfo->byteOffset << 16) |
+                 spareBitsMaskInfo->bytesInMask);
       // TODO: Endianness??
-      for (unsigned i = 0; i < wordsInMask; ++i) {
+      for (unsigned i = 0; i < spareBitsMaskInfo->wordsInMask(); ++i) {
         uint32_t nextWord = bits.extractBitsAsZExtValue(32, 0);
         B.addInt32(nextWord);
         bits.lshrInPlace(32);