Fix Switch for 128bit integers

Author: bjorn3

cranelift/frontend/src/switch.rs

@@ -1,11 +1,29 @@
 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;
+
+/*
+FIXME add test for
+
+fn main() {
+    let options = [1u128];
+    match options[0] {
+        1 => (),
+        0 => loop {},
+        v => panic(v),
+    };
+}
+
+fn panic(v: u128) -> !{
+    panic!("{}", v)
+}
+*/
 
 /// 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 +170,20 @@ 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 +227,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 +252,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 +284,10 @@ 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 +298,31 @@ 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);
         }
     }