[support] Make ScopedPrinter compatible with bitmask enums

[labath]

.. produced by ADT/BitmaskEnum.h.

These aren't implicitly convertible to an integer, so I needed to tweak a couple of bitwise operations and add an explicit constructor for HexNumber and FlagEntry.

Motivation: I'd like to use this in the SFrame data structures (#147264)

[llvmbot]

@llvm/pr-subscribers-llvm-support

Author: Pavel Labath (labath)

Changes

.. produced by ADT/BitmaskEnum.h.

These aren't implicitly convertible to an integer, so I needed to tweak a couple of bitwise operations and add an explicit constructor for HexNumber and FlagEntry.

Motivation: I'd like to use this in the SFrame data structures (#147264)

---

Full diff: https://github.com/llvm/llvm-project/pull/147512.diff

2 Files Affected:

• (modified) llvm/include/llvm/Support/ScopedPrinter.h (+16-7)
• (modified) llvm/unittests/Support/ScopedPrinterTest.cpp (+29)

diff --git a/llvm/include/llvm/Support/ScopedPrinter.h b/llvm/include/llvm/Support/ScopedPrinter.h
index 753fd3f55bbdd..e6c4cc4a4ea13 100644
--- a/llvm/include/llvm/Support/ScopedPrinter.h
+++ b/llvm/include/llvm/Support/ScopedPrinter.h
@@ -19,6 +19,7 @@
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/raw_ostream.h"
+#include <type_traits>
 
 namespace llvm {
 
@@ -41,8 +42,8 @@ template <typename T> struct EnumEntry {
 struct HexNumber {
   // To avoid sign-extension we have to explicitly cast to the appropriate
   // unsigned type. The overloads are here so that every type that is implicitly
-  // convertible to an integer (including enums and endian helpers) can be used
-  // without requiring type traits or call-site changes.
+  // convertible to an integer (including endian helpers) can be used without
+  // requiring type traits or call-site changes.
   HexNumber(char Value) : Value(static_cast<unsigned char>(Value)) {}
   HexNumber(signed char Value) : Value(static_cast<unsigned char>(Value)) {}
   HexNumber(signed short Value) : Value(static_cast<unsigned short>(Value)) {}
@@ -55,6 +56,10 @@ struct HexNumber {
   HexNumber(unsigned int Value) : Value(Value) {}
   HexNumber(unsigned long Value) : Value(Value) {}
   HexNumber(unsigned long long Value) : Value(Value) {}
+  template <typename EnumT, typename = std::enable_if_t<std::is_enum_v<EnumT>>>
+  HexNumber(EnumT Value)
+      : HexNumber(static_cast<std::underlying_type_t<EnumT>>(Value)) {}
+
   uint64_t Value;
 };
 
@@ -77,6 +82,10 @@ struct FlagEntry {
   FlagEntry(StringRef Name, unsigned long Value) : Name(Name), Value(Value) {}
   FlagEntry(StringRef Name, unsigned long long Value)
       : Name(Name), Value(Value) {}
+  template <typename EnumT, typename = std::enable_if_t<std::is_enum_v<EnumT>>>
+  FlagEntry(StringRef Name, EnumT Value)
+      : FlagEntry(Name, static_cast<std::underlying_type_t<EnumT>>(Value)) {}
+
   StringRef Name;
   uint64_t Value;
 };
@@ -165,17 +174,17 @@ class LLVM_ABI ScopedPrinter {
     SmallVector<FlagEntry, 10> SetFlags(ExtraFlags);
 
     for (const auto &Flag : Flags) {
-      if (Flag.Value == 0)
+      if (Flag.Value == TFlag{})
         continue;
 
       TFlag EnumMask{};
-      if (Flag.Value & EnumMask1)
+      if ((Flag.Value & EnumMask1) != TFlag{})
         EnumMask = EnumMask1;
-      else if (Flag.Value & EnumMask2)
+      else if ((Flag.Value & EnumMask2) != TFlag{})
         EnumMask = EnumMask2;
-      else if (Flag.Value & EnumMask3)
+      else if ((Flag.Value & EnumMask3) != TFlag{})
         EnumMask = EnumMask3;
-      bool IsEnum = (Flag.Value & EnumMask) != 0;
+      bool IsEnum = (Flag.Value & EnumMask) != TFlag{};
       if ((!IsEnum && (Value & Flag.Value) == Flag.Value) ||
           (IsEnum && (Value & EnumMask) == Flag.Value)) {
         SetFlags.emplace_back(Flag.Name, Flag.Value);
diff --git a/llvm/unittests/Support/ScopedPrinterTest.cpp b/llvm/unittests/Support/ScopedPrinterTest.cpp
index 72af331a2b009..1e2b13865f68a 100644
--- a/llvm/unittests/Support/ScopedPrinterTest.cpp
+++ b/llvm/unittests/Support/ScopedPrinterTest.cpp
@@ -8,11 +8,21 @@
 
 #include "llvm/Support/ScopedPrinter.h"
 #include "llvm/ADT/APSInt.h"
+#include "llvm/ADT/BitmaskEnum.h"
 #include "llvm/Support/Format.h"
 #include "gtest/gtest.h"
 #include <cmath>
 #include <vector>
 
+namespace {
+enum class BitmaskEnum : uint8_t {
+  F1 = 0x1,
+  F2 = 0x2,
+  LLVM_MARK_AS_BITMASK_ENUM(/*LargestValue=*/F2),
+};
+LLVM_ENABLE_BITMASK_ENUMS_IN_NAMESPACE();
+} // namespace
+
 using namespace llvm;
 
 TEST(JSONScopedPrinterTest, PrettyPrintCtor) {
@@ -261,6 +271,12 @@ TEST_F(ScopedPrinterTest, PrintFlag) {
         {"SecondByte2", "Second2", 0x20u}, {"SecondByte3", "Second3", 0x30u},
         {"ThirdByte1", "Third1", 0x100u},  {"ThirdByte2", "Third2", 0x200u},
         {"ThirdByte3", "Third3", 0x300u}};
+
+    const EnumEntry<BitmaskEnum> ScopedFlags[] = {
+        {"F1", "AltF1", BitmaskEnum::F1},
+        {"F2", "AltF2", BitmaskEnum::F2},
+    };
+
     W.printFlags("ZeroFlag", 0, ArrayRef(SingleBitFlags));
     W.printFlags("NoFlag", 1 << 3, ArrayRef(SingleBitFlags));
     W.printFlags("Flag1", SingleBitFlags[1].Value, ArrayRef(SingleBitFlags));
@@ -286,6 +302,7 @@ TEST_F(ScopedPrinterTest, PrintFlag) {
                  FirstByteMask, SecondByteMask);
     W.printFlags("FirstSecondThirdByteMask", 0x333u, ArrayRef(EnumFlags),
                  FirstByteMask, SecondByteMask, ThirdByteMask);
+    W.printFlags("BitmaskEnum::F1", BitmaskEnum::F1, ArrayRef(ScopedFlags));
   };
 
   const char *ExpectedOut = R"(ZeroFlag [ (0x0)
@@ -343,6 +360,9 @@ FirstSecondThirdByteMask [ (0x333)
   SecondByte3 (0x30)
   ThirdByte3 (0x300)
 ]
+BitmaskEnum::F1 [ (0x1)
+  F1 (0x1)
+]
 )";
 
   const char *JSONExpectedOut = R"({
@@ -504,6 +524,15 @@ FirstSecondThirdByteMask [ (0x333)
         "Value": 768
       }
     ]
+  },
+  "BitmaskEnum::F1": {
+    "Value": 1,
+    "Flags": [
+      {
+        "Name": "F1",
+        "Value": 1
+      }
+    ]
   }
 })";
   verifyAll(ExpectedOut, JSONExpectedOut, PrintFunc);

[ilovepi]

LGTM