From 7e947042644cd2956b2e963b0724abe3e16e2112 Mon Sep 17 00:00:00 2001
From: Afonso Bordado <afonso360@users.noreply.github.com>
Date: Tue, 3 Jan 2023 19:59:14 +0000
Subject: [PATCH] riscv64: Add masking for small types when lowering select
 (#5504)

When lowering `select+icmp` we have an optimization that allows us to
avoid materializing the icmp result.

We were accidentally not masking the high bits for i8 and i16 in this case.

Issue #5498 reported this as an illegal instruction but what was happening
there was that the invalid select caused a division by zero.
---
 cranelift/codegen/src/isa/riscv64/lower.isle   |  6 ++++--
 .../filetests/isa/riscv64/condops.clif         | 12 ++++++++----
 .../filetests/runtests/issue-5498.clif         | 18 ++++++++++++++++++
 3 files changed, 30 insertions(+), 6 deletions(-)
 create mode 100644 cranelift/filetests/filetests/runtests/issue-5498.clif

diff --git a/cranelift/codegen/src/isa/riscv64/lower.isle b/cranelift/codegen/src/isa/riscv64/lower.isle
index 060b8a5892..e951378a40 100644
--- a/cranelift/codegen/src/isa/riscv64/lower.isle
+++ b/cranelift/codegen/src/isa/riscv64/lower.isle
@@ -615,8 +615,10 @@
   (gen_select ty (truthy_to_reg cty (normalize_cmp_value cty c)) x y))
 
 (rule 1
-  (lower (has_type ty (select (icmp cc a b) x y)))
-  (gen_select_reg cc a b x y))
+  (lower (has_type ty (select (icmp cc a b @ (value_type in_ty)) x y)))
+  (let ((a Reg (normalize_cmp_value in_ty a))
+        (b Reg (normalize_cmp_value in_ty b)))
+    (gen_select_reg cc a b x y)))
 
 ;;;;;  Rules for `bitselect`;;;;;;;;;
 
diff --git a/cranelift/filetests/filetests/isa/riscv64/condops.clif b/cranelift/filetests/filetests/isa/riscv64/condops.clif
index 6077e2b8ac..baa894f28f 100644
--- a/cranelift/filetests/filetests/isa/riscv64/condops.clif
+++ b/cranelift/filetests/filetests/isa/riscv64/condops.clif
@@ -11,8 +11,10 @@ block0(v0: i8, v1: i64, v2: i64):
 }
 
 ; block0:
-;   li a3,42
-;   select_reg a0,a1,a2##condition=(a0 eq a3)
+;   andi a3,a0,255
+;   li a4,42
+;   andi a5,a4,255
+;   select_reg a0,a1,a2##condition=(a3 eq a5)
 ;   ret
 
 function %g(i8) -> i8 {
@@ -62,8 +64,10 @@ block0(v0: i32, v1: i8, v2: i8):
 }
 
 ; block0:
-;   li a3,42
-;   select_reg a0,a1,a2##condition=(a0 eq a3)
+;   addiw a3,a0,0
+;   li a4,42
+;   addiw a5,a4,0
+;   select_reg a0,a1,a2##condition=(a3 eq a5)
 ;   ret
 
 function %i128_select(i8, i128, i128) -> i128 {
diff --git a/cranelift/filetests/filetests/runtests/issue-5498.clif b/cranelift/filetests/filetests/runtests/issue-5498.clif
new file mode 100644
index 0000000000..c17f2a2c84
--- /dev/null
+++ b/cranelift/filetests/filetests/runtests/issue-5498.clif
@@ -0,0 +1,18 @@
+test interpret
+test run
+target aarch64
+target s390x
+target x86_64
+target riscv64
+
+function %a(i16, i8) -> i16 {
+block0(v0: i16, v1: i8):
+    v2 = iconst.i16 0
+    v3 = iconst.i16 1
+
+    v4 = ishl v0, v1
+    v5 = icmp eq v4, v2
+    v6 = select v5, v3, v4
+    return v6
+}
+; run: %a(514, -1) == 1