cranelift: mask shift amount in (x << k) >> k mid-end rules

cranelift/codegen/src/opts/shifts.isle

@@ -79,15 +79,20 @@
 ;; (x << N) >> N == x as T_SMALL as T_LARGE
 ;; if N == bytesizeof(T_LARGE) - bytesizeof(T_SMALL)
 ;;
+;; Shift amounts are taken modulo the bit width of the shifted type, so `N` is
+;; masked before it's used to compute `T_SMALL`.
+;;
 ;; Note that the shift is required to be >0 to ensure this doesn't accidentally
 ;; try to `ireduce` a type to itself, which isn't a valid use of `ireduce`.
 (rule (simplify (sshr (ty_int ty) (ishl ty x (iconst _ shift)) (iconst _ shift)))
-      (if-let (u64_from_imm64 (u64_extract_non_zero shift_u64)) shift)
-      (if-let ty_small (shift_amt_to_type (u64_wrapping_sub (ty_bits ty) shift_u64)))
+      (if-let (u64_from_imm64 shift_u64) shift)
+      (if-let (u64_extract_non_zero shift_amt) (u64_and shift_u64 (ty_shift_mask ty)))
+      (if-let ty_small (shift_amt_to_type (u64_wrapping_sub (ty_bits ty) shift_amt)))
       (sextend ty (ireduce ty_small x)))
 (rule (simplify (ushr (ty_int ty) (ishl ty x (iconst _ shift)) (iconst _ shift)))
-      (if-let (u64_from_imm64 (u64_extract_non_zero shift_u64)) shift)
-      (if-let ty_small (shift_amt_to_type (u64_wrapping_sub (ty_bits ty) shift_u64)))
+      (if-let (u64_from_imm64 shift_u64) shift)
+      (if-let (u64_extract_non_zero shift_amt) (u64_and shift_u64 (ty_shift_mask ty)))
+      (if-let ty_small (shift_amt_to_type (u64_wrapping_sub (ty_bits ty) shift_amt)))
       (uextend ty (ireduce ty_small x)))
 
 (decl pure partial shift_amt_to_type (u64) Type)

cranelift/filetests/filetests/egraph/shifts.clif

@@ -305,6 +305,55 @@ block0(v0: i64):
     ; check: v19 = uextend.i64 v18
     ; check: return v19
 }
+function %i8_shl_sshr_neg8(i8) -> i8 {
+block0(v0: i8):
+    v1 = iconst.i64 -8
+    v2 = ishl v0, v1
+    v3 = sshr v2, v1
+    return v3
+    ; check: v4 = iconst.i64 0
+    ; check: v5 = ishl v0, v4
+    ; check: v7 = sshr v5, v4
+    ; check: return v7
+}
+function %i8_shl_ushr_neg8(i8) -> i8 {
+block0(v0: i8):
+    v1 = iconst.i64 -8
+    v2 = ishl v0, v1
+    v3 = ushr v2, v1
+    return v3
+    ; check: return v0
+}
+function %i8_shl_sshr_neg24(i8) -> i8 {
+block0(v0: i8):
+    v1 = iconst.i64 -24
+    v2 = ishl v0, v1
+    v3 = sshr v2, v1
+    return v3
+    ; check: v4 = iconst.i64 0
+    ; check: v5 = ishl v0, v4
+    ; check: v7 = sshr v5, v4
+    ; check: return v7
+}
+function %i16_shl_sshr_neg16(i16) -> i16 {
+block0(v0: i16):
+    v1 = iconst.i64 -16
+    v2 = ishl v0, v1
+    v3 = sshr v2, v1
+    return v3
+    ; check: v4 = iconst.i64 0
+    ; check: v5 = ishl v0, v4
+    ; check: v7 = sshr v5, v4
+    ; check: return v7
+}
+function %i16_shl_ushr_neg16(i16) -> i16 {
+block0(v0: i16):
+    v1 = iconst.i64 -16
+    v2 = ishl v0, v1
+    v3 = ushr v2, v1
+    return v3
+    ; check: return v0
+}
 function %ishl_amt_type_ireduce(i8, i16) -> i8 {
 block0(v0: i8, v1: i16):
     v2 = ireduce.i8 v1

cranelift/filetests/filetests/runtests/shift-left-right-same-amount.clif

@@ -0,0 +1,96 @@
+;; Test that the `(x << k) >> k` rewrite into an extend of a reduce is correct,
+;; including when `k` is out of range for the type and therefore wraps.
+
+test interpret
+test run
+set opt_level=speed
+target aarch64
+target x86_64
+target x86_64 has_bmi2
+target riscv64
+target riscv64 has_c has_zcb
+target s390x
+target pulley32
+target pulley32be
+target pulley64
+target pulley64be
+
+function %sshr_ishl_24_i32(i32) -> i32 {
+block0(v0: i32):
+    v1 = iconst.i32 24
+    v2 = ishl v0, v1
+    v3 = sshr v2, v1
+    return v3
+}
+; run: %sshr_ishl_24_i32(0x12345678) == 0x78
+; run: %sshr_ishl_24_i32(0x12345680) == 0xffffff80
+
+function %ushr_ishl_24_i32(i32) -> i32 {
+block0(v0: i32):
+    v1 = iconst.i32 24
+    v2 = ishl v0, v1
+    v3 = ushr v2, v1
+    return v3
+}
+; run: %ushr_ishl_24_i32(0x12345678) == 0x78
+; run: %ushr_ishl_24_i32(0x12345680) == 0x80
+
+function %sshr_ishl_56_i32(i32) -> i32 {
+block0(v0: i32):
+    v1 = iconst.i32 56
+    v2 = ishl v0, v1
+    v3 = sshr v2, v1
+    return v3
+}
+; run: %sshr_ishl_56_i32(0x12345678) == 0x78
+; run: %sshr_ishl_56_i32(0x12345680) == 0xffffff80
+
+function %sshr_ishl_neg8_i8(i8) -> i8 {
+block0(v0: i8):
+    v1 = iconst.i64 -8
+    v2 = ishl v0, v1
+    v3 = sshr v2, v1
+    return v3
+}
+; run: %sshr_ishl_neg8_i8(0x7f) == 0x7f
+; run: %sshr_ishl_neg8_i8(0x80) == 0x80
+
+function %ushr_ishl_neg8_i8(i8) -> i8 {
+block0(v0: i8):
+    v1 = iconst.i64 -8
+    v2 = ishl v0, v1
+    v3 = ushr v2, v1
+    return v3
+}
+; run: %ushr_ishl_neg8_i8(0x7f) == 0x7f
+; run: %ushr_ishl_neg8_i8(0x80) == 0x80
+
+function %sshr_ishl_neg24_i8(i8) -> i8 {
+block0(v0: i8):
+    v1 = iconst.i64 -24
+    v2 = ishl v0, v1
+    v3 = sshr v2, v1
+    return v3
+}
+; run: %sshr_ishl_neg24_i8(0x7f) == 0x7f
+; run: %sshr_ishl_neg24_i8(0x80) == 0x80
+
+function %sshr_ishl_neg16_i16(i16) -> i16 {
+block0(v0: i16):
+    v1 = iconst.i64 -16
+    v2 = ishl v0, v1
+    v3 = sshr v2, v1
+    return v3
+}
+; run: %sshr_ishl_neg16_i16(0x7fff) == 0x7fff
+; run: %sshr_ishl_neg16_i16(0x8001) == 0x8001
+
+function %ushr_ishl_neg16_i16(i16) -> i16 {
+block0(v0: i16):
+    v1 = iconst.i64 -16
+    v2 = ishl v0, v1
+    v3 = ushr v2, v1
+    return v3
+}
+; run: %ushr_ishl_neg16_i16(0x7fff) == 0x7fff
+; run: %ushr_ishl_neg16_i16(0x8001) == 0x8001