Cranelift: Correctly wrap shifts in constant propagation (#5695)

Fixes #5690
Fixes #5696

Co-authored-by: Jamey Sharp <jsharp@fastly.com>

cranelift/codegen/src/isle_prelude.rs

@@ -81,27 +81,42 @@ macro_rules! isle_common_prelude_methods {
         }
 
         #[inline]
-        fn u64_shl(&mut self, x: u64, y: u64) -> u64 {
-            x << y
+        fn imm64_shl(&mut self, ty: Type, x: Imm64, y: Imm64) -> Imm64 {
+            // Mask off any excess shift bits.
+            let shift_mask = (ty.bits() - 1) as u64;
+            let y = (y.bits() as u64) & shift_mask;
+
+            // Mask the result to `ty` bits.
+            let ty_mask = self.ty_mask(ty) as i64;
+            Imm64::new((x.bits() << y) & ty_mask)
         }
 
         #[inline]
         fn imm64_ushr(&mut self, ty: Type, x: Imm64, y: Imm64) -> Imm64 {
-            let shift = u32::checked_sub(64, ty.bits()).unwrap_or(0);
-            let mask = u64::MAX >> shift;
-            let x = (x.bits() as u64) & mask;
-            let y = (y.bits() as u64) & mask;
+            let ty_mask = self.ty_mask(ty);
+            let x = (x.bits() as u64) & ty_mask;
+
+            // Mask off any excess shift bits.
+            let shift_mask = (ty.bits() - 1) as u64;
+            let y = (y.bits() as u64) & shift_mask;
+
+            // NB: No need to mask off high bits because they are already zero.
             Imm64::new((x >> y) as i64)
         }
 
         #[inline]
         fn imm64_sshr(&mut self, ty: Type, x: Imm64, y: Imm64) -> Imm64 {
+            // Sign extend `x` from `ty.bits()`-width to the full 64 bits.
             let shift = u32::checked_sub(64, ty.bits()).unwrap_or(0);
-            let mask = u64::MAX >> shift;
-            let x = (x.bits() as u64) & mask;
-            let y = (y.bits() as u64) & mask;
-            let sext = |v| ((v << shift) as i64) >> shift;
-            Imm64::new(sext(x) >> sext(y))
+            let x = (x.bits() << shift) >> shift;
+
+            // Mask off any excess shift bits.
+            let shift_mask = (ty.bits() - 1) as i64;
+            let y = y.bits() & shift_mask;
+
+            // Mask off sign bits that aren't part of `ty`.
+            let ty_mask = self.ty_mask(ty) as i64;
+            Imm64::new((x >> y) & ty_mask)
         }
 
         #[inline]
@@ -179,14 +194,12 @@ macro_rules! isle_common_prelude_methods {
 
         #[inline]
         fn ty_mask(&mut self, ty: Type) -> u64 {
-            match ty.bits() {
-                1 => 1,
-                8 => 0xff,
-                16 => 0xffff,
-                32 => 0xffff_ffff,
-                64 => 0xffff_ffff_ffff_ffff,
-                _ => unimplemented!(),
-            }
+            let ty_bits = ty.bits();
+            debug_assert_ne!(ty_bits, 0);
+            let shift = 64_u64
+                .checked_sub(ty_bits.into())
+                .expect("unimplemented for > 64 bits");
+            u64::MAX >> shift
         }
 
         fn fits_in_16(&mut self, ty: Type) -> Option<Type> {

cranelift/codegen/src/opts/algebraic.isle

@@ -191,15 +191,15 @@
 ;; `(x >> k) << k` is the same as masking off the bottom `k` bits (regardless if
 ;; this is a signed or unsigned shift right).
 (rule (simplify (ishl (fits_in_64 ty)
-                      (ushr ty x (iconst _ (u64_from_imm64 k)))
-                      (iconst _ (u64_from_imm64 k))))
-      (let ((mask u64 (u64_shl 0xFFFFFFFFFFFFFFFF k)))
-        (band ty x (iconst ty (imm64_masked ty mask)))))
+                      (ushr ty x (iconst _ k))
+                      (iconst _ k)))
+      (let ((mask Imm64 (imm64_shl ty (imm64 0xFFFF_FFFF_FFFF_FFFF) k)))
+        (band ty x (iconst ty mask))))
 (rule (simplify (ishl (fits_in_64 ty)
-                      (sshr ty x (iconst _ (u64_from_imm64 k)))
-                      (iconst _ (u64_from_imm64 k))))
-      (let ((mask u64 (u64_shl 0xFFFFFFFFFFFFFFFF k)))
-        (band ty x (iconst ty (imm64_masked ty mask)))))
+                      (sshr ty x (iconst _ k))
+                      (iconst _ k)))
+      (let ((mask Imm64 (imm64_shl ty (imm64 0xFFFF_FFFF_FFFF_FFFF) k)))
+        (band ty x (iconst ty mask))))
 
 ;; Rematerialize ALU-op-with-imm and iconsts in each block where they're
 ;; used. This is neutral (add-with-imm) or positive (iconst) for

cranelift/codegen/src/opts/cprop.isle

@@ -56,18 +56,18 @@
       (subsume (iconst ty (imm64_masked ty (u64_not k)))))
 
 (rule (simplify (ishl (fits_in_64 ty)
-                      (iconst ty (u64_from_imm64 k1))
-                      (iconst ty (u64_from_imm64 k2))))
-      (subsume (iconst ty (imm64_masked ty (u64_shl k1 k2)))))
+                      (iconst ty k1)
+                      (iconst _ k2)))
+      (subsume (iconst ty (imm64_shl ty k1 k2))))
 
 (rule (simplify (ushr (fits_in_64 ty)
                       (iconst ty k1)
-                      (iconst ty k2)))
+                      (iconst _ k2)))
       (subsume (iconst ty (imm64_ushr ty k1 k2))))
 
 (rule (simplify (sshr (fits_in_64 ty)
                       (iconst ty k1)
-                      (iconst ty k2)))
+                      (iconst _ k2)))
       (subsume (iconst ty (imm64_sshr ty k1 k2))))
 
 (rule (simplify

cranelift/codegen/src/prelude.isle

@@ -129,8 +129,8 @@
 (decl pure u64_xor (u64 u64) u64)
 (extern constructor u64_xor u64_xor)
 
-(decl pure u64_shl (u64 u64) u64)
-(extern constructor u64_shl u64_shl)
+(decl pure imm64_shl (Type Imm64 Imm64) Imm64)
+(extern constructor imm64_shl imm64_shl)
 
 (decl pure imm64_ushr (Type Imm64 Imm64) Imm64)
 (extern constructor imm64_ushr imm64_ushr)