From 1e49b0266c73e2b930b9de4975da0bdd1ed745a5 Mon Sep 17 00:00:00 2001
From: Zachary Turner <zturner@google.com>
Date: Wed, 19 Apr 2017 15:45:31 +0000
Subject: [PATCH] [Support] Add some helpers to generate bitmasks.

Frequently you you want a bitmask consisting of a specified
number of 1s, either at the beginning or end of a word.

The naive way to do this is to write

template<typename T>
T leadingBitMask(unsigned N) {
  return (T(1) << N) - 1;
}

but using this function you cannot produce a word with every
bit set to 1 (i.e. leadingBitMask<uint8_t>(8)) because left
shift is undefined when N is greater than or equal to the
number of bits in the word.

This patch provides an efficient, branch-free implementation
that works for all values of N in [0, CHAR_BIT*sizeof(T)]

Differential Revision: https://reviews.llvm.org/D32212

llvm-svn: 300710
---
 include/llvm/Support/MathExtras.h    | 15 +++++++++++++++
 unittests/Support/MathExtrasTest.cpp | 20 ++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/include/llvm/Support/MathExtras.h b/include/llvm/Support/MathExtras.h
index 19380b23d9d..e389a7e3b95 100644
--- a/include/llvm/Support/MathExtras.h
+++ b/include/llvm/Support/MathExtras.h
@@ -198,6 +198,21 @@ template <typename T> T findFirstSet(T Val, ZeroBehavior ZB = ZB_Max) {
   return countTrailingZeros(Val, ZB_Undefined);
 }
 
+/// \brief Create a bitmask with the N right-most bits set to 1, and all other
+/// bits set to 0.  Only unsigned types are allowed.
+template <typename T> T maskTrailingOnes(unsigned N) {
+  static_assert(std::is_unsigned<T>::value, "Invalid type!");
+  const unsigned Bits = CHAR_BIT * sizeof(T);
+  assert(N <= Bits && "Invalid bit index");
+  return -T(N != 0) & (T(-1) >> (Bits - N));
+}
+
+/// \brief Create a bitmask with the N left-most bits set to 1, and all other
+/// bits set to 0.  Only unsigned types are allowed.
+template <typename T> T maskLeadingOnes(unsigned N) {
+  return ~maskTrailingOnes<T>(CHAR_BIT * sizeof(T) - N);
+}
+
 /// \brief Get the index of the last set bit starting from the least
 ///   significant bit.
 ///
diff --git a/unittests/Support/MathExtrasTest.cpp b/unittests/Support/MathExtrasTest.cpp
index b2c37797887..de827217083 100644
--- a/unittests/Support/MathExtrasTest.cpp
+++ b/unittests/Support/MathExtrasTest.cpp
@@ -66,6 +66,26 @@ TEST(MathExtras, countLeadingZeros) {
   }
 }
 
+TEST(MathExtras, onesMask) {
+  EXPECT_EQ(0U, maskLeadingOnes<uint8_t>(0));
+  EXPECT_EQ(0U, maskTrailingOnes<uint8_t>(0));
+  EXPECT_EQ(0U, maskLeadingOnes<uint16_t>(0));
+  EXPECT_EQ(0U, maskTrailingOnes<uint16_t>(0));
+  EXPECT_EQ(0U, maskLeadingOnes<uint32_t>(0));
+  EXPECT_EQ(0U, maskTrailingOnes<uint32_t>(0));
+  EXPECT_EQ(0U, maskLeadingOnes<uint64_t>(0));
+  EXPECT_EQ(0U, maskTrailingOnes<uint64_t>(0));
+
+  EXPECT_EQ(0x00000003U, maskTrailingOnes<uint32_t>(2U));
+  EXPECT_EQ(0xC0000000U, maskLeadingOnes<uint32_t>(2U));
+
+  EXPECT_EQ(0x000007FFU, maskTrailingOnes<uint32_t>(11U));
+  EXPECT_EQ(0xFFE00000U, maskLeadingOnes<uint32_t>(11U));
+
+  EXPECT_EQ(0xFFFFFFFFU, maskTrailingOnes<uint32_t>(32U));
+  EXPECT_EQ(0xFFFFFFFFU, maskLeadingOnes<uint32_t>(32U));
+}
+
 TEST(MathExtras, findFirstSet) {
   uint8_t Z8 = 0;
   uint16_t Z16 = 0;