From 761fe891cab5d9b87c8eb00542b49462332655ea Mon Sep 17 00:00:00 2001
From: Afonso Bordado <afonsobordado@az8.co>
Date: Fri, 22 Jul 2022 00:58:58 +0100
Subject: [PATCH] cranelift: Fix Switch bug when emitting indexes larger than
 val type

In #4502 we discovered a bug in the switch api where it would emit
`icmp_imm`'s with types that were not able to fully represent the
destination index.

We now reject these inputs. The index val must always have a
type that is capable of addressing the entire range of inputs.
---
 cranelift/frontend/src/switch.rs | 99 +++++++++++++++++---------------
 1 file changed, 54 insertions(+), 45 deletions(-)

diff --git a/cranelift/frontend/src/switch.rs b/cranelift/frontend/src/switch.rs
index 7be8c8d0aa31..4ff852c5d26f 100644
--- a/cranelift/frontend/src/switch.rs
+++ b/cranelift/frontend/src/switch.rs
@@ -293,11 +293,16 @@ impl Switch {
     /// * The value to switch on
     /// * The default block
     pub fn emit(self, bx: &mut FunctionBuilder, val: Value, otherwise: Block) {
-        // FIXME icmp(_imm) doesn't have encodings for i8 and i16 on x86(_64) yet
-        let val = match bx.func.dfg.value_type(val) {
-            types::I8 | types::I16 => bx.ins().uextend(types::I32, val),
-            _ => val,
-        };
+        // Validate that the type of `val` is sufficiently wide to address all cases.
+        let max = self.cases.keys().max().copied().unwrap_or(0);
+        let val_ty = bx.func.dfg.value_type(val);
+        let val_ty_max = val_ty.bounds(false).1;
+        if max > val_ty_max {
+            panic!(
+                "The index type {} does not fit the maximum switch entry of {}",
+                val_ty, max
+            );
+        }
 
         let contiguous_case_ranges = self.collect_contiguous_case_ranges();
         let cases_and_jt_blocks =
@@ -384,8 +389,7 @@ mod tests {
             func,
             "block0:
     v0 = iconst.i8 0
-    v1 = uextend.i32 v0
-    brz v1, block1
+    brz v0, block1
     jump block0"
         );
     }
@@ -397,9 +401,8 @@ mod tests {
             func,
             "block0:
     v0 = iconst.i8 0
-    v1 = uextend.i32 v0
-    v2 = icmp_imm eq v1, 1
-    brnz v2, block1
+    v1 = icmp_imm eq v0, 1
+    brnz v1, block1
     jump block0"
         );
     }
@@ -413,11 +416,10 @@ mod tests {
 
 block0:
     v0 = iconst.i8 0
-    v1 = uextend.i32 v0
     jump block3
 
 block3:
-    br_table.i32 v1, block0, jt0"
+    br_table.i8 v0, block0, jt0"
         );
     }
 
@@ -428,13 +430,12 @@ block3:
             func,
             "block0:
     v0 = iconst.i8 0
-    v1 = uextend.i32 v0
-    v2 = icmp_imm eq v1, 2
-    brnz v2, block2
+    v1 = icmp_imm eq v0, 2
+    brnz v1, block2
     jump block3
 
 block3:
-    brz.i32 v1, block1
+    brz.i8 v0, block1
     jump block0"
         );
     }
@@ -449,92 +450,100 @@ block3:
 
 block0:
     v0 = iconst.i8 0
-    v1 = uextend.i32 v0
-    v2 = icmp_imm uge v1, 7
-    brnz v2, block9
+    v1 = icmp_imm uge v0, 7
+    brnz v1, block9
     jump block8
 
 block9:
-    v3 = icmp_imm.i32 uge v1, 10
-    brnz v3, block10
+    v2 = icmp_imm.i8 uge v0, 10
+    brnz v2, block10
     jump block11
 
 block11:
-    v4 = icmp_imm.i32 eq v1, 7
-    brnz v4, block4
+    v3 = icmp_imm.i8 eq v0, 7
+    brnz v3, block4
     jump block0
 
 block8:
-    v5 = icmp_imm.i32 eq v1, 5
-    brnz v5, block3
+    v4 = icmp_imm.i8 eq v0, 5
+    brnz v4, block3
     jump block12
 
 block12:
-    br_table.i32 v1, block0, jt0
+    br_table.i8 v0, block0, jt0
 
 block10:
-    v6 = iadd_imm.i32 v1, -10
-    br_table v6, block0, jt1"
+    v5 = iadd_imm.i8 v0, -10
+    br_table v5, block0, jt1"
         );
     }
 
     #[test]
     fn switch_min_index_value() {
-        let func = setup!(0, [::core::i64::MIN as u64 as u128, 1,]);
+        let func = setup!(0, [i8::MIN as u8 as u128, 1,]);
         assert_eq!(
             func,
             "block0:
     v0 = iconst.i8 0
-    v1 = uextend.i32 v0
-    v2 = icmp_imm eq v1, 0x8000_0000_0000_0000
-    brnz v2, block1
+    v1 = icmp_imm eq v0, 128
+    brnz v1, block1
     jump block3
 
 block3:
-    v3 = icmp_imm.i32 eq v1, 1
-    brnz v3, block2
+    v2 = icmp_imm.i8 eq v0, 1
+    brnz v2, block2
     jump block0"
         );
     }
 
     #[test]
     fn switch_max_index_value() {
-        let func = setup!(0, [::core::i64::MAX as u64 as u128, 1,]);
+        let func = setup!(0, [i8::MAX as u8 as u128, 1,]);
         assert_eq!(
             func,
             "block0:
     v0 = iconst.i8 0
-    v1 = uextend.i32 v0
-    v2 = icmp_imm eq v1, 0x7fff_ffff_ffff_ffff
-    brnz v2, block1
+    v1 = icmp_imm eq v0, 127
+    brnz v1, block1
     jump block3
 
 block3:
-    v3 = icmp_imm.i32 eq v1, 1
-    brnz v3, block2
+    v2 = icmp_imm.i8 eq v0, 1
+    brnz v2, block2
     jump block0"
         )
     }
 
     #[test]
     fn switch_optimal_codegen() {
-        let func = setup!(0, [-1i64 as u64 as u128, 0, 1,]);
+        let func = setup!(0, [-1i8 as u8 as u128, 0, 1,]);
         assert_eq!(
             func,
             "    jt0 = jump_table [block2, block3]
 
 block0:
     v0 = iconst.i8 0
-    v1 = uextend.i32 v0
-    v2 = icmp_imm eq v1, -1
-    brnz v2, block1
+    v1 = icmp_imm eq v0, 255
+    brnz v1, block1
     jump block4
 
 block4:
-    br_table.i32 v1, block0, jt0"
+    br_table.i8 v0, block0, jt0"
         );
     }
 
+    #[test]
+    #[should_panic(
+        expected = "The index type i8 does not fit the maximum switch entry of 4683743612477887600"
+    )]
+    fn switch_rejects_small_inputs() {
+        // This is a regression test for a bug that we found where we would emit a cmp
+        // with a type that was not able to fully represent a large index.
+        //
+        // See: https://github.com/bytecodealliance/wasmtime/pull/4502#issuecomment-1191961677
+        setup!(1, [0x4100_0000_00bf_d470,]);
+    }
+
     #[test]
     fn switch_seal_generated_blocks() {
         let cases = &[vec![0, 1, 2], vec![0, 1, 2, 10, 11, 12, 20, 30, 40, 50]];