absl int128: avoid shifting signed integer by a number of bits greater than or equal to the precision of the operand

PiperOrigin-RevId: 519808237
Change-Id: I9123b167b606d609b8f3924d6f4fd298fa866a90

2 files changed, 70 insertions, 30 deletions

diff --git a/absl/numeric/int128_no_intrinsic.inc b/absl/numeric/int128_no_intrinsic.inc
index 8834804c..6f5d8377 100644
--- a/absl/numeric/int128_no_intrinsic.inc
+++ b/absl/numeric/int128_no_intrinsic.inc
@@ -23,8 +23,7 @@ constexpr int64_t Int128High64(int128 v) { return v.hi_; }
 
 #if defined(ABSL_IS_LITTLE_ENDIAN)
 
-constexpr int128::int128(int64_t high, uint64_t low) :
-    lo_(low), hi_(high) {}
+constexpr int128::int128(int64_t high, uint64_t low) : lo_(low), hi_(high) {}
 
 constexpr int128::int128(int v)
     : lo_{static_cast<uint64_t>(v)}, hi_{v < 0 ? ~int64_t{0} : 0} {}
@@ -44,8 +43,7 @@ constexpr int128::int128(uint128 v)
 
 #elif defined(ABSL_IS_BIG_ENDIAN)
 
-constexpr int128::int128(int64_t high, uint64_t low) :
-    hi_{high}, lo_{low} {}
+constexpr int128::int128(int64_t high, uint64_t low) : hi_{high}, lo_{low} {}
 
 constexpr int128::int128(int v)
     : hi_{v < 0 ? ~int64_t{0} : 0}, lo_{static_cast<uint64_t>(v)} {}
@@ -279,33 +277,52 @@ constexpr int128 operator^(int128 lhs, int128 rhs) {
 }
 
 constexpr int128 operator<<(int128 lhs, int amount) {
-  // int64_t shifts of >= 64 are undefined, so we need some special-casing.
-  return amount >= 64
-             ? MakeInt128(
-                   static_cast<int64_t>(Int128Low64(lhs) << (amount - 64)), 0)
-         : amount == 0
-             ? lhs
-             : MakeInt128(
-                   (Int128High64(lhs) << amount) |
-                       static_cast<int64_t>(Int128Low64(lhs) >> (64 - amount)),
-                   Int128Low64(lhs) << amount);
+  // int64_t shifts of >= 63 are undefined, so we need some special-casing.
+  assert(amount >= 0 && amount < 127);
+  if (amount <= 0) {
+    return lhs;
+  } else if (amount < 63) {
+    return MakeInt128(
+        (Int128High64(lhs) << amount) |
+            static_cast<int64_t>(Int128Low64(lhs) >> (64 - amount)),
+        Int128Low64(lhs) << amount);
+  } else if (amount == 63) {
+    return MakeInt128(((Int128High64(lhs) << 32) << 31) |
+                          static_cast<int64_t>(Int128Low64(lhs) >> 1),
+                      (Int128Low64(lhs) << 32) << 31);
+  } else if (amount == 127) {
+    return MakeInt128(static_cast<int64_t>(Int128Low64(lhs) << 63), 0);
+  } else if (amount > 127) {
+    return MakeInt128(0, 0);
+  } else {
+    // amount >= 64 && amount < 127
+    return MakeInt128(static_cast<int64_t>(Int128Low64(lhs) << (amount - 64)),
+                      0);
+  }
 }
 
 constexpr int128 operator>>(int128 lhs, int amount) {
-  // int64_t shifts of >= 64 are undefined, so we need some special-casing.
-  // The (Int128High64(lhs) >> 32) >> 32 "trick" causes the the most significant
-  // int64 to be inititialized with all zeros or all ones correctly. It takes
-  // into account whether the number is negative or positive, and whether the
-  // current architecture does arithmetic or logical right shifts for negative
-  // numbers.
-  return amount >= 64
-             ? MakeInt128(
-                   (Int128High64(lhs) >> 32) >> 32,
-                   static_cast<uint64_t>(Int128High64(lhs) >> (amount - 64)))
-         : amount == 0
-             ? lhs
-             : MakeInt128(Int128High64(lhs) >> amount,
-                          (Int128Low64(lhs) >> amount) |
-                              (static_cast<uint64_t>(Int128High64(lhs))
-                               << (64 - amount)));
+  // int64_t shifts of >= 63 are undefined, so we need some special-casing.
+  assert(amount >= 0 && amount < 127);
+  if (amount <= 0) {
+    return lhs;
+  } else if (amount < 63) {
+    return MakeInt128(
+        Int128High64(lhs) >> amount,
+        Int128Low64(lhs) >> amount | static_cast<uint64_t>(Int128High64(lhs))
+                                         << (64 - amount));
+  } else if (amount == 63) {
+    return MakeInt128((Int128High64(lhs) >> 32) >> 31,
+                      static_cast<uint64_t>(Int128High64(lhs) << 1) |
+                          (Int128Low64(lhs) >> 32) >> 31);
+
+  } else if (amount >= 127) {
+    return MakeInt128((Int128High64(lhs) >> 32) >> 31,
+                      static_cast<uint64_t>((Int128High64(lhs) >> 32) >> 31));
+  } else {
+    // amount >= 64 && amount < 127
+    return MakeInt128(
+        (Int128High64(lhs) >> 32) >> 31,
+        static_cast<uint64_t>(Int128High64(lhs) >> (amount - 64)));
+  }
 }
diff --git a/absl/numeric/int128_test.cc b/absl/numeric/int128_test.cc
index dd9425d7..a6fed5cf 100644
--- a/absl/numeric/int128_test.cc
+++ b/absl/numeric/int128_test.cc
@@ -32,6 +32,8 @@
 #pragma warning(disable:4146)
 #endif
 
+#define MAKE_INT128(HI, LO) absl::MakeInt128(static_cast<int64_t>(HI), LO)
+
 namespace {
 
 template <typename T>
@@ -1245,6 +1247,27 @@ TEST(Int128, BitwiseShiftTest) {
                 absl::MakeInt128(uint64_t{1} << j, 0) >>= (j - i));
     }
   }
+
+  // Manually calculated cases with shift count for positive (val1) and negative
+  // (val2) values
+  absl::int128 val1 = MAKE_INT128(0x123456789abcdef0, 0x123456789abcdef0);
+  absl::int128 val2 = MAKE_INT128(0xfedcba0987654321, 0xfedcba0987654321);
+
+  EXPECT_EQ(val1 << 63, MAKE_INT128(0x91a2b3c4d5e6f78, 0x0));
+  EXPECT_EQ(val1 << 64, MAKE_INT128(0x123456789abcdef0, 0x0));
+  EXPECT_EQ(val2 << 63, MAKE_INT128(0xff6e5d04c3b2a190, 0x8000000000000000));
+  EXPECT_EQ(val2 << 64, MAKE_INT128(0xfedcba0987654321, 0x0));
+
+  EXPECT_EQ(val1 << 126, MAKE_INT128(0x0, 0x0));
+  EXPECT_EQ(val2 << 126, MAKE_INT128(0x4000000000000000, 0x0));
+
+  EXPECT_EQ(val1 >> 63, MAKE_INT128(0x0, 0x2468acf13579bde0));
+  EXPECT_EQ(val1 >> 64, MAKE_INT128(0x0, 0x123456789abcdef0));
+  EXPECT_EQ(val2 >> 63, MAKE_INT128(0xffffffffffffffff, 0xfdb974130eca8643));
+  EXPECT_EQ(val2 >> 64, MAKE_INT128(0xffffffffffffffff, 0xfedcba0987654321));
+
+  EXPECT_EQ(val1 >> 126, MAKE_INT128(0x0, 0x0));
+  EXPECT_EQ(val2 >> 126, MAKE_INT128(0xffffffffffffffff, 0xffffffffffffffff));
 }
 
 TEST(Int128, NumericLimitsTest) {