MathExtras: template'ize alignToPowerOf2 (#97814)

Follow up on 5627794 (MathExtras: avoid unnecessarily widening types) to
change the overflow behavior of alignToPowerOf2 to only overflow if the
result is not representable in the return type. This allows us to
template'ize it, and avoid unnecessarily widening the types of
arguments.

Author: artagnon

llvm/include/llvm/Support/MathExtras.h

@@ -497,13 +497,21 @@ constexpr uint64_t alignTo(uint64_t Value, uint64_t Align) {
   return CeilDiv * Align;
 }
 
+/// Will overflow only if result is not representable in T.
+template <typename U, typename V, typename T = common_uint<U, V>>
+constexpr T alignToPowerOf2(U Value, V Align) {
+  assert(Align != 0 && (Align & (Align - 1)) == 0 &&
+         "Align must be a power of 2");
+  T NegAlign = static_cast<T>(0) - Align;
+  return (Value + (Align - 1)) & NegAlign;
+}
+
+/// Fallback when arguments aren't integral.
 constexpr uint64_t alignToPowerOf2(uint64_t Value, uint64_t Align) {
   assert(Align != 0 && (Align & (Align - 1)) == 0 &&
          "Align must be a power of 2");
-  // Replace unary minus to avoid compilation error on Windows:
-  // "unary minus operator applied to unsigned type, result still unsigned"
-  uint64_t NegAlign = (~Align) + 1;
-  return (Value + Align - 1) & NegAlign;
+  uint64_t NegAlign = 0 - Align;
+  return (Value + (Align - 1)) & NegAlign;
 }
 
 /// If non-zero \p Skew is specified, the return value will be a minimal integer

llvm/unittests/Support/MathExtrasTest.cpp

@@ -218,8 +218,11 @@ TEST(MathExtras, AlignToPowerOf2) {
   EXPECT_EQ(24u, alignToPowerOf2(17, 8));
   EXPECT_EQ(0u, alignToPowerOf2(~0LL, 8));
   EXPECT_EQ(240u, alignToPowerOf2(240, 16));
-  EXPECT_EQ(static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()) + 1,
-            alignToPowerOf2(std::numeric_limits<uint32_t>::max(), 2));
+
+  // Overflow.
+  EXPECT_EQ(0u, alignToPowerOf2(static_cast<uint8_t>(200),
+                                static_cast<uint8_t>(128)));
+  EXPECT_EQ(0u, alignToPowerOf2(std::numeric_limits<uint32_t>::max(), 2));
 }
 
 TEST(MathExtras, AlignDown) {