From 217434b474b76f855b237dee921d2841d2bea8cb Mon Sep 17 00:00:00 2001
From: Jakob Stoklund Olesen <jolesen@mozilla.com>
Date: Fri, 25 Aug 2017 12:40:43 -0700
Subject: [PATCH] Add custom legalization for conditional traps.

The expansion of these instructions requires the CFG to be modified,
something the Python XForms can't yet do.
---
 .../filetests/isa/intel/legalize-custom.cton  | 30 ++++++++++++
 lib/cretonne/meta/base/legalize.py            |  5 ++
 lib/cretonne/src/legalizer/mod.rs             | 48 ++++++++++++++++++-
 3 files changed, 82 insertions(+), 1 deletion(-)
 create mode 100644 cranelift/filetests/isa/intel/legalize-custom.cton

diff --git a/cranelift/filetests/isa/intel/legalize-custom.cton b/cranelift/filetests/isa/intel/legalize-custom.cton
new file mode 100644
index 0000000000..ae63a56e0a
--- /dev/null
+++ b/cranelift/filetests/isa/intel/legalize-custom.cton
@@ -0,0 +1,30 @@
+; Test the custom legalizations.
+test legalizer
+isa intel
+set is_64bit
+isa intel
+
+; regex: V=v\d+
+; regex: EBB=ebb\d+
+
+function %cond_trap(i32) {
+ebb0(v1: i32):
+    trapz v1
+    return
+    ; check: $ebb0($v1: i32):
+    ; nextln: brnz $v1, $(new=$EBB)
+    ; nextln: trap
+    ; check: $new:
+    ; nextln: return
+}
+
+function %cond_trap2(i32) {
+ebb0(v1: i32):
+    trapnz v1
+    return
+    ; check: $ebb0($v1: i32):
+    ; nextln: brz $v1, $(new=$EBB)
+    ; nextln: trap
+    ; check: $new:
+    ; nextln: return
+}
diff --git a/lib/cretonne/meta/base/legalize.py b/lib/cretonne/meta/base/legalize.py
index 5db7e6b7c8..2bbdd7247a 100644
--- a/lib/cretonne/meta/base/legalize.py
+++ b/lib/cretonne/meta/base/legalize.py
@@ -48,6 +48,11 @@ expand = XFormGroup('expand', """
 expand.custom_legalize(insts.global_addr, 'expand_global_addr')
 expand.custom_legalize(insts.heap_addr, 'expand_heap_addr')
 
+# Custom expansions that need to change the CFG.
+# TODO: Add sufficient XForm syntax that we don't need to hand-code these.
+expand.custom_legalize(insts.trapz, 'expand_cond_trap')
+expand.custom_legalize(insts.trapnz, 'expand_cond_trap')
+
 x = Var('x')
 y = Var('y')
 a = Var('a')
diff --git a/lib/cretonne/src/legalizer/mod.rs b/lib/cretonne/src/legalizer/mod.rs
index 8909366e01..da4a2b93fa 100644
--- a/lib/cretonne/src/legalizer/mod.rs
+++ b/lib/cretonne/src/legalizer/mod.rs
@@ -15,7 +15,7 @@
 
 use cursor::{Cursor, FuncCursor};
 use flowgraph::ControlFlowGraph;
-use ir;
+use ir::{self, InstBuilder};
 use isa::TargetIsa;
 use bitset::BitSet;
 
@@ -96,3 +96,49 @@ pub fn legalize_function(func: &mut ir::Function, cfg: &mut ControlFlowGraph, is
 //
 // Concretely, this defines private functions `narrow()`, and `expand()`.
 include!(concat!(env!("OUT_DIR"), "/legalizer.rs"));
+
+/// Custom expansion for conditional trap instructions.
+/// TODO: Add CFG support to the Python patterns so we won't have to do this.
+fn expand_cond_trap(inst: ir::Inst, func: &mut ir::Function, cfg: &mut ControlFlowGraph) {
+    // Parse the instruction.
+    let trapz;
+    let arg = match func.dfg[inst] {
+        ir::InstructionData::Unary { opcode, arg } => {
+            // We want to branch *over* an unconditional trap.
+            trapz = match opcode {
+                ir::Opcode::Trapz => true,
+                ir::Opcode::Trapnz => false,
+                _ => panic!("Expected cond trap: {}", func.dfg.display_inst(inst, None)),
+            };
+            arg
+        }
+        _ => panic!("Expected cond trap: {}", func.dfg.display_inst(inst, None)),
+    };
+
+    // Split the EBB after `inst`:
+    //
+    //     trapnz arg
+    //
+    // Becomes:
+    //
+    //     brz arg, new_ebb
+    //     trap
+    //   new_ebb:
+    //
+    let old_ebb = func.layout.pp_ebb(inst);
+    let new_ebb = func.dfg.make_ebb();
+    if trapz {
+        func.dfg.replace(inst).brnz(arg, new_ebb, &[]);
+    } else {
+        func.dfg.replace(inst).brz(arg, new_ebb, &[]);
+    }
+
+    let mut pos = FuncCursor::new(func).at_inst(inst);
+    pos.next_inst();
+    pos.ins().trap();
+    pos.insert_ebb(new_ebb);
+
+    // Finally update the CFG.
+    cfg.recompute_ebb(pos.func, old_ebb);
+    cfg.recompute_ebb(pos.func, new_ebb);
+}