cranelift: add a new resumable_trapnz instruction;

This is useful to have to allow resumable_trap to happen in loop
headers, for instance. This is the correct way to implement interrupt
checks in Spidermonkey, which are effectively resumable traps. Previous
implementation was using traps, which is wrong, since traps semantically
can't be resumed after.

Author: bnjbvr

cranelift/codegen/meta/src/shared/instructions.rs

@@ -340,7 +340,21 @@ fn define_control_flow(
                 r#"
         Trap when non-zero.
 
-        if ``c`` is zero, execution continues at the following instruction.
+        If ``c`` is zero, execution continues at the following instruction.
+        "#,
+                &formats.cond_trap,
+            )
+            .operands_in(vec![c, code])
+            .can_trap(true),
+        );
+
+        ig.push(
+            Inst::new(
+                "resumable_trapnz",
+                r#"
+        A resumable trap to be called when the passed condition is non-zero.
+
+        If ``c`` is zero, execution continues at the following instruction.
         "#,
                 &formats.cond_trap,
             )

cranelift/codegen/meta/src/shared/legalize.rs

@@ -99,6 +99,7 @@ pub(crate) fn define(insts: &InstructionGroup, imm: &Immediates) -> TransformGro
     let jump = insts.by_name("jump");
     let load = insts.by_name("load");
     let popcnt = insts.by_name("popcnt");
+    let resumable_trapnz = insts.by_name("resumable_trapnz");
     let rotl = insts.by_name("rotl");
     let rotl_imm = insts.by_name("rotl_imm");
     let rotr = insts.by_name("rotr");
@@ -138,6 +139,7 @@ pub(crate) fn define(insts: &InstructionGroup, imm: &Immediates) -> TransformGro
     // TODO: Add sufficient XForm syntax that we don't need to hand-code these.
     expand.custom_legalize(trapz, "expand_cond_trap");
     expand.custom_legalize(trapnz, "expand_cond_trap");
+    expand.custom_legalize(resumable_trapnz, "expand_cond_trap");
     expand.custom_legalize(br_table, "expand_br_table");
     expand.custom_legalize(select, "expand_select");
     widen.custom_legalize(select, "expand_select"); // small ints

cranelift/codegen/src/isa/aarch64/lower_inst.rs

@@ -1334,7 +1334,7 @@ pub(crate) fn lower_insn_to_regs<C: LowerCtx<I = Inst>>(
             ctx.emit(Inst::Brk);
         }
 
-        Opcode::Trap => {
+        Opcode::Trap | Opcode::ResumableTrap => {
             let trap_info = (ctx.srcloc(insn), inst_trapcode(ctx.data(insn)).unwrap());
             ctx.emit(Inst::Udf { trap_info })
         }
@@ -1385,12 +1385,8 @@ pub(crate) fn lower_insn_to_regs<C: LowerCtx<I = Inst>>(
             panic!("safepoint support not implemented!");
         }
 
-        Opcode::Trapz | Opcode::Trapnz => {
-            panic!("trapz / trapnz should have been removed by legalization!");
-        }
-
-        Opcode::ResumableTrap => {
-            panic!("Resumable traps not supported");
+        Opcode::Trapz | Opcode::Trapnz | Opcode::ResumableTrapnz => {
+            panic!("trapz / trapnz / resumable_trapnz should have been removed by legalization!");
         }
 
         Opcode::FuncAddr => {
@@ -2277,6 +2273,7 @@ pub(crate) fn lower_branch<C: LowerCtx<I = Inst>>(
                     dest: BranchTarget::Label(targets[0]),
                 });
             }
+
             Opcode::BrTable => {
                 // Expand `br_table index, default, JT` to:
                 //

cranelift/codegen/src/legalizer/mod.rs

@@ -214,6 +214,7 @@ pub fn simple_legalize(func: &mut ir::Function, cfg: &mut ControlFlowGraph, isa:
                 | ir::Opcode::TableAddr
                 | ir::Opcode::Trapnz
                 | ir::Opcode::Trapz
+                | ir::Opcode::ResumableTrapnz
                 | ir::Opcode::BandImm
                 | ir::Opcode::BorImm
                 | ir::Opcode::BxorImm
@@ -261,15 +262,15 @@ fn expand_cond_trap(
 ) {
     // Parse the instruction.
     let trapz;
-    let (arg, code) = match func.dfg[inst] {
+    let (arg, code, opcode) = match func.dfg[inst] {
         ir::InstructionData::CondTrap { opcode, arg, code } => {
             // We want to branch *over* an unconditional trap.
             trapz = match opcode {
                 ir::Opcode::Trapz => true,
-                ir::Opcode::Trapnz => false,
+                ir::Opcode::Trapnz | ir::Opcode::ResumableTrapnz => false,
                 _ => panic!("Expected cond trap: {}", func.dfg.display_inst(inst, None)),
             };
-            (arg, code)
+            (arg, code, opcode)
         }
         _ => panic!("Expected cond trap: {}", func.dfg.display_inst(inst, None)),
     };
@@ -307,7 +308,17 @@ fn expand_cond_trap(
 
     // Insert the new label and the unconditional trap terminator.
     pos.insert_block(new_block_trap);
-    pos.ins().trap(code);
+
+    match opcode {
+        ir::Opcode::Trapz | ir::Opcode::Trapnz => {
+            pos.ins().trap(code);
+        }
+        ir::Opcode::ResumableTrapnz => {
+            pos.ins().resumable_trap(code);
+            pos.ins().jump(new_block_resume, &[]);
+        }
+        _ => unreachable!(),
+    }
 
     // Insert the new label and resume the execution when the trap fails.
     pos.insert_block(new_block_resume);