Fix Switch for 128bit integers

cranelift/frontend/src/switch.rs

@@ -1,11 +1,12 @@
 use super::HashMap;
 use crate::frontend::FunctionBuilder;
 use alloc::vec::Vec;
+use core::convert::TryFrom;
 use cranelift_codegen::ir::condcodes::IntCC;
 use cranelift_codegen::ir::*;
 use log::debug;
 
-type EntryIndex = u64;
+type EntryIndex = u128;
 
 /// Unlike with `br_table`, `Switch` cases may be sparse or non-0-based.
 /// They emit efficient code using branches, jump tables, or a combination of both.
@@ -152,11 +153,17 @@ impl Switch {
                 let left_block = bx.create_block();
                 let right_block = bx.create_block();
 
-                let should_take_right_side = bx.ins().icmp_imm(
-                    IntCC::UnsignedGreaterThanOrEqual,
-                    val,
-                    right[0].first_index as i64,
-                );
+                let first_index = right[0].first_index;
+                let should_take_right_side = if let Ok(index) = u64::try_from(first_index) {
+                    bx.ins()
+                        .icmp_imm(IntCC::UnsignedGreaterThanOrEqual, val, index as i64)
+                } else {
+                    let (lsb, msb) = (first_index as u64, (first_index >> 64) as u64);
+                    let lsb = bx.ins().iconst(types::I64, lsb as i64);
+                    let msb = bx.ins().iconst(types::I64, msb as i64);
+                    let index = bx.ins().iconcat(lsb, msb);
+                    bx.ins().icmp(IntCC::UnsignedGreaterThanOrEqual, val, index)
+                };
                 bx.ins().brnz(should_take_right_side, right_block, &[]);
                 bx.ins().jump(left_block, &[]);
 
@@ -200,7 +207,15 @@ impl Switch {
                 }
                 (1, _) => {
                     ins_fallthrough_jump(was_branch, bx);
-                    let is_good_val = bx.ins().icmp_imm(IntCC::Equal, val, first_index as i64);
+                    let is_good_val = if let Ok(first_index) = u64::try_from(first_index) {
+                        bx.ins().icmp_imm(IntCC::Equal, val, first_index as i64)
+                    } else {
+                        let (lsb, msb) = (first_index as u64, (first_index >> 64) as u64);
+                        let lsb = bx.ins().iconst(types::I64, lsb as i64);
+                        let msb = bx.ins().iconst(types::I64, msb as i64);
+                        let index = bx.ins().iconcat(lsb, msb);
+                        bx.ins().icmp(IntCC::Equal, val, index)
+                    };
                     bx.ins().brnz(is_good_val, blocks[0], &[]);
                 }
                 (_, 0) => {
@@ -217,11 +232,19 @@ impl Switch {
                 (_, _) => {
                     ins_fallthrough_jump(was_branch, bx);
                     let jt_block = bx.create_block();
-                    let is_good_val = bx.ins().icmp_imm(
-                        IntCC::UnsignedGreaterThanOrEqual,
-                        val,
-                        first_index as i64,
-                    );
+                    let is_good_val = if let Ok(first_index) = u64::try_from(first_index) {
+                        bx.ins().icmp_imm(
+                            IntCC::UnsignedGreaterThanOrEqual,
+                            val,
+                            first_index as i64,
+                        )
+                    } else {
+                        let (lsb, msb) = (first_index as u64, (first_index >> 64) as u64);
+                        let lsb = bx.ins().iconst(types::I64, lsb as i64);
+                        let msb = bx.ins().iconst(types::I64, msb as i64);
+                        let index = bx.ins().iconcat(lsb, msb);
+                        bx.ins().icmp(IntCC::UnsignedGreaterThanOrEqual, val, index)
+                    };
                     bx.ins().brnz(is_good_val, jt_block, &[]);
                     bx.seal_block(jt_block);
                     cases_and_jt_blocks.push((first_index, jt_block, blocks));
@@ -241,6 +264,13 @@ impl Switch {
         cases_and_jt_blocks: Vec<(EntryIndex, Block, Vec<Block>)>,
     ) {
         for (first_index, jt_block, blocks) in cases_and_jt_blocks.into_iter().rev() {
+            // There are currently no 128bit systems supported by rustc, but once we do ensure that
+            // we don't silently ignore a part of the jump table for 128bit integers on 128bit systems.
+            assert!(
+                core::mem::size_of::<usize>() <= 64,
+                "128bit jump tables are not yet supported"
+            );
+
             let mut jt_data = JumpTableData::new();
             for block in blocks {
                 jt_data.push_entry(block);
@@ -251,8 +281,33 @@ impl Switch {
             let discr = if first_index == 0 {
                 val
             } else {
-                bx.ins().iadd_imm(val, (first_index as i64).wrapping_neg())
+                if let Ok(first_index) = u64::try_from(first_index) {
+                    bx.ins().iadd_imm(val, (first_index as i64).wrapping_neg())
+                } else {
+                    let (lsb, msb) = (first_index as u64, (first_index >> 64) as u64);
+                    let lsb = bx.ins().iconst(types::I64, lsb as i64);
+                    let msb = bx.ins().iconst(types::I64, msb as i64);
+                    let index = bx.ins().iconcat(lsb, msb);
+                    bx.ins().isub(val, index)
+                }
+            };
+
+            let discr = if bx.func.dfg.value_type(discr).bits() > 64 {
+                // Check for overflow of cast to u64.
+                let new_block = bx.create_block();
+                let bigger_than_u64 =
+                    bx.ins()
+                        .icmp_imm(IntCC::UnsignedGreaterThan, discr, u64::max_value() as i64);
+                bx.ins().brnz(bigger_than_u64, otherwise, &[]);
+                bx.ins().jump(new_block, &[]);
+                bx.switch_to_block(new_block);
+
+                // Cast to u64, as br_table is not implemented for integers bigger than 64bits.
+                bx.ins().ireduce(types::I64, discr)
+            } else {
+                discr
             };
+
             bx.ins().br_table(discr, otherwise, jump_table);
         }
     }
@@ -440,7 +495,7 @@ block10:
 
     #[test]
     fn switch_min_index_value() {
-        let func = setup!(0, [::core::i64::MIN as u64, 1,]);
+        let func = setup!(0, [::core::i64::MIN as u64 as u128, 1,]);
         assert_eq!(
             func,
             "block0:
@@ -459,7 +514,7 @@ block3:
 
     #[test]
     fn switch_max_index_value() {
-        let func = setup!(0, [::core::i64::MAX as u64, 1,]);
+        let func = setup!(0, [::core::i64::MAX as u64 as u128, 1,]);
         assert_eq!(
             func,
             "block0:
@@ -478,7 +533,7 @@ block3:
 
     #[test]
     fn switch_optimal_codegen() {
-        let func = setup!(0, [-1i64 as u64, 0, 1,]);
+        let func = setup!(0, [-1i64 as u64 as u128, 0, 1,]);
         assert_eq!(
             func,
             "    jt0 = jump_table [block2, block3]
@@ -530,4 +585,45 @@ block4:
 
         builder.finalize(); // Will panic if some blocks are not sealed
     }
+
+    #[test]
+    fn switch_128bit() {
+        let mut func = Function::new();
+        let mut func_ctx = FunctionBuilderContext::new();
+        {
+            let mut bx = FunctionBuilder::new(&mut func, &mut func_ctx);
+            let block0 = bx.create_block();
+            bx.switch_to_block(block0);
+            let val = bx.ins().iconst(types::I128, 0);
+            let mut switch = Switch::new();
+            let block1 = bx.create_block();
+            switch.set_entry(1, block1);
+            let block2 = bx.create_block();
+            switch.set_entry(0, block2);
+            let block3 = bx.create_block();
+            switch.emit(&mut bx, val, block3);
+        }
+        let func = func
+            .to_string()
+            .trim_start_matches("function u0:0() fast {\n")
+            .trim_end_matches("\n}\n")
+            .to_string();
+        assert_eq!(
+            func,
+            "    jt0 = jump_table [block2, block1]
+
+block0:
+    v0 = iconst.i128 0
+    jump block4
+
+block4:
+    v1 = icmp_imm.i128 ugt v0, -1
+    brnz v1, block3
+    jump block5
+
+block5:
+    v2 = ireduce.i64 v0
+    br_table v2, block3, jt0"
+        );
+    }
 }