From 669a40006bf3494b0d320fb77bf39cc19a6c5bc4 Mon Sep 17 00:00:00 2001
From: Ralf Jung <post@ralfj.de>
Date: Fri, 14 Jun 2024 11:28:18 +0200
Subject: [PATCH] MIR Shl/Shr: the offset can be computed with rem_euclid

---
 .../src/interpret/operator.rs                 | 19 +++++++------------
 compiler/rustc_middle/src/mir/syntax.rs       |  6 ++++--
 2 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/compiler/rustc_const_eval/src/interpret/operator.rs b/compiler/rustc_const_eval/src/interpret/operator.rs
index a6eef9f5662ca..4d0f35618ef27 100644
--- a/compiler/rustc_const_eval/src/interpret/operator.rs
+++ b/compiler/rustc_const_eval/src/interpret/operator.rs
@@ -112,25 +112,20 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
 
         // Shift ops can have an RHS with a different numeric type.
         if matches!(bin_op, Shl | ShlUnchecked | Shr | ShrUnchecked) {
-            let size = left.layout.size.bits();
+            let l_bits = left.layout.size.bits();
             // Compute the equivalent shift modulo `size` that is in the range `0..size`. (This is
             // the one MIR operator that does *not* directly map to a single LLVM operation.)
             let (shift_amount, overflow) = if right.layout.abi.is_signed() {
                 let shift_amount = r_signed();
-                let overflow = shift_amount < 0 || shift_amount >= i128::from(size);
-                // Deliberately wrapping `as` casts: shift_amount *can* be negative, but the result
-                // of the `as` will be equal modulo `size` (since it is a power of two).
-                let masked_amount = (shift_amount as u128) % u128::from(size);
-                assert_eq!(overflow, shift_amount != i128::try_from(masked_amount).unwrap());
-                (masked_amount, overflow)
+                let rem = shift_amount.rem_euclid(l_bits.into());
+                // `rem` is guaranteed positive, so the `unwrap` cannot fail
+                (u128::try_from(rem).unwrap(), rem != shift_amount)
             } else {
                 let shift_amount = r_unsigned();
-                let overflow = shift_amount >= u128::from(size);
-                let masked_amount = shift_amount % u128::from(size);
-                assert_eq!(overflow, shift_amount != masked_amount);
-                (masked_amount, overflow)
+                let rem = shift_amount.rem_euclid(l_bits.into());
+                (rem, rem != shift_amount)
             };
-            let shift_amount = u32::try_from(shift_amount).unwrap(); // we masked so this will always fit
+            let shift_amount = u32::try_from(shift_amount).unwrap(); // we brought this in the range `0..size` so this will always fit
             // Compute the shifted result.
             let result = if left.layout.abi.is_signed() {
                 let l = l_signed();
diff --git a/compiler/rustc_middle/src/mir/syntax.rs b/compiler/rustc_middle/src/mir/syntax.rs
index ebe77a1abfd8b..3edc5fe36cd2f 100644
--- a/compiler/rustc_middle/src/mir/syntax.rs
+++ b/compiler/rustc_middle/src/mir/syntax.rs
@@ -1490,7 +1490,8 @@ pub enum BinOp {
     BitOr,
     /// The `<<` operator (shift left)
     ///
-    /// The offset is (uniquely) determined as follows:
+    /// The offset is given by `RHS.rem_euclid(LHS::BITS)`.
+    /// In other words, it is (uniquely) determined as follows:
     /// - it is "equal modulo LHS::BITS" to the RHS
     /// - it is in the range `0..LHS::BITS`
     Shl,
@@ -1498,7 +1499,8 @@ pub enum BinOp {
     ShlUnchecked,
     /// The `>>` operator (shift right)
     ///
-    /// The offset is (uniquely) determined as follows:
+    /// The offset is given by `RHS.rem_euclid(LHS::BITS)`.
+    /// In other words, it is (uniquely) determined as follows:
     /// - it is "equal modulo LHS::BITS" to the RHS
     /// - it is in the range `0..LHS::BITS`
     ///