[BitmaskEnum] Add support for shift operators.

[sdesmalen-arm]

For enums that describe a bitmask where successive bits within that mask describe some enum value (as described in the same enum), it is useful to support operator<< and operator>> as well.

For example:

  enum class E : unsigned {
    // 2 bits per option
    OptionA = 0,
    OptionB = 1,
    OptionC = 2,
    OptionD = 3,
    OptionMask = 3,

    // Given 3 values in the bitmask X, Y and Z, each is 2 bits in size
    // and represents a choice of OptionA..OptionD.
    ShiftX = 0,
    ShiftY = 2,
    ShiftZ = 4,
  };

  // The mask can be encoded with:
  E mask;
  mask |= getOptionFor(X) << E::ShiftX;
  mask |= getOptionFor(Y) << E::ShiftY;
  mask |= getOptionFor(Z) << E::ShiftZ;

  // And to extract a value:
  E OptionForX =  (mask >> E::ShiftX) & E::OptionMask;
  E OptionForY =  (mask >> E::ShiftY) & E::OptionMask;
  E OptionForZ =  (mask >> E::ShiftZ) & E::OptionMask;

[llvmbot]

@llvm/pr-subscribers-llvm-adt

Author: Sander de Smalen (sdesmalen-arm)

Changes

For enums that describe a bitmask where successive bits within that mask describe some enum value (as described in the same enum), it is useful to support operator<< and operator>> as well.

For example:

enum class E : unsigned {
// 2 bits per option
OptionA = 0,
OptionB = 1,
OptionC = 2,
OptionD = 3,
OptionMask = 3,

// Given 3 values in the bitmask X, Y and Z, each is 2 bits in size
// and represents a choice of OptionA..OptionD.
ShiftX = 0,
ShiftY = 2,
ShiftZ = 4,

};

// The mask can be encoded with:
E mask;
mask |= getOptionFor(X) << E::ShiftX;
mask |= getOptionFor(Y) << E::ShiftY;
mask |= getOptionFor(Z) << E::ShiftZ;

// And to extract a value:
E OptionForX = (mask >> E::ShiftX) & E::OptionMask;
E OptionForY = (mask >> E::ShiftY) & E::OptionMask;
E OptionForZ = (mask >> E::ShiftZ) & E::OptionMask;

---

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

2 Files Affected:

• (modified) llvm/include/llvm/ADT/BitmaskEnum.h (+13)
• (modified) llvm/unittests/ADT/BitmaskEnumTest.cpp (+9)

diff --git a/llvm/include/llvm/ADT/BitmaskEnum.h b/llvm/include/llvm/ADT/BitmaskEnum.h
index c87e7cac65a5b1..41e43b2ada621d 100644
--- a/llvm/include/llvm/ADT/BitmaskEnum.h
+++ b/llvm/include/llvm/ADT/BitmaskEnum.h
@@ -85,9 +85,11 @@
   using ::llvm::BitmaskEnumDetail::operator|;                                  \
   using ::llvm::BitmaskEnumDetail::operator&;                                  \
   using ::llvm::BitmaskEnumDetail::operator^;                                  \
+  using ::llvm::BitmaskEnumDetail::operator<<;                                 \
   using ::llvm::BitmaskEnumDetail::operator|=;                                 \
   using ::llvm::BitmaskEnumDetail::operator&=;                                 \
   using ::llvm::BitmaskEnumDetail::operator^=;                                 \
+  using ::llvm::BitmaskEnumDetail::operator<<=;                                \
   /* Force a semicolon at the end of this macro. */                            \
   using ::llvm::BitmaskEnumDetail::any
 
@@ -162,6 +164,11 @@ constexpr E operator^(E LHS, E RHS) {
   return static_cast<E>(Underlying(LHS) ^ Underlying(RHS));
 }
 
+template <typename E, typename = std::enable_if_t<is_bitmask_enum<E>::value>>
+constexpr E operator<<(E LHS, E RHS) {
+  return static_cast<E>(Underlying(LHS) << Underlying(RHS));
+}
+
 // |=, &=, and ^= return a reference to LHS, to match the behavior of the
 // operators on builtin types.
 
@@ -183,6 +190,12 @@ E &operator^=(E &LHS, E RHS) {
   return LHS;
 }
 
+template <typename E, typename = std::enable_if_t<is_bitmask_enum<E>::value>>
+E &operator<<=(E &LHS, E RHS) {
+  LHS = LHS << RHS;
+  return LHS;
+}
+
 } // namespace BitmaskEnumDetail
 
 // Enable bitmask enums in namespace ::llvm and all nested namespaces.
diff --git a/llvm/unittests/ADT/BitmaskEnumTest.cpp b/llvm/unittests/ADT/BitmaskEnumTest.cpp
index c78937c3571fd1..5927e0bb58c71f 100644
--- a/llvm/unittests/ADT/BitmaskEnumTest.cpp
+++ b/llvm/unittests/ADT/BitmaskEnumTest.cpp
@@ -130,6 +130,15 @@ TEST(BitmaskEnumTest, BitwiseXorEquals) {
   EXPECT_EQ(V3, f2);
 }
 
+TEST(BitmaskEnumTest, BitwiseShift) {
+  Flags f = (F1 << F1);
+  EXPECT_EQ(f, F2);
+
+  Flags f2 = F1;
+  f2 <<= F1;
+  EXPECT_EQ(f2, F2);
+}
+
 TEST(BitmaskEnumTest, ConstantExpression) {
   constexpr Flags f1 = ~F1;
   constexpr Flags f2 = F1 | F2;

[kmclaughlin-arm]

In the commit message you mention operator>>, should this be supported here too?

[sdesmalen-arm]

Sorry, I forgot to git add those changes. Fixed now.