From 0fed78e06365f5df225278f4ab8d513d08c7009c Mon Sep 17 00:00:00 2001
From: Dan Gohman <sunfish@mozilla.com>
Date: Mon, 12 Nov 2018 05:17:54 -0800
Subject: [PATCH] Don't allow loop preheaders to have critical edges.

If the block which would be a preheader for a loop has other successors,
don't hoist instructions into it. Instead create a dedicated preheader.
---
 cranelift/filetests/licm/complex.clif       | 76 +++++++++++----------
 cranelift/filetests/licm/critical-edge.clif | 42 ++++++++++++
 lib/codegen/src/licm.rs                     | 19 +++---
 3 files changed, 94 insertions(+), 43 deletions(-)
 create mode 100644 cranelift/filetests/licm/critical-edge.clif

diff --git a/cranelift/filetests/licm/complex.clif b/cranelift/filetests/licm/complex.clif
index 0f825e3a8e..c7d26f9a33 100644
--- a/cranelift/filetests/licm/complex.clif
+++ b/cranelift/filetests/licm/complex.clif
@@ -3,43 +3,43 @@ target riscv32
 
 function %complex(i32) -> i32 system_v {
 ebb0(v0: i32):
-    jump ebb1(v0)
+[UJ#1b]    jump ebb1(v0)
 
-ebb1(v1: i32):
-    v2 = iconst.i32 1
-    v3 = iconst.i32 4
-    v4 = iadd v2, v1
-    brz v1, ebb2(v2)
-    jump ebb4(v4)
+         ebb1(v1: i32):
+            v2 = iconst.i32 1
+            v3 = iconst.i32 4
+            v4 = iadd v2, v1
+[SBzero#18] brz v1, ebb2(v2)
+[UJ#1b]     jump ebb4(v4)
 
-ebb2(v5: i32):
-    v6 = iconst.i32 2
-    v7 = iadd v5, v4
-    v8 = iadd v6, v1
-    jump ebb3(v8)
+         ebb2(v5: i32):
+            v6 = iconst.i32 2
+            v7 = iadd v5, v4
+            v8 = iadd v6, v1
+[UJ#1b]   jump ebb3(v8)
 
-ebb3(v9: i32):
-    v10 = iadd v9, v5
-    v11 = iadd.i32 v1, v4
-    brz.i32 v1, ebb2(v9)
-    jump ebb6(v10)
+         ebb3(v9: i32):
+            v10 = iadd v9, v5
+            v11 = iadd.i32 v1, v4
+[SBzero#18] brz.i32 v1, ebb2(v9)
+[UJ#1b]   jump ebb6(v10)
 
-ebb4(v12: i32):
-    v13 = iconst.i32 3
-    v14 = iadd v12, v13
-    v15 = iadd.i32 v4, v13
-    jump ebb5(v13)
+         ebb4(v12: i32):
+            v13 = iconst.i32 3
+            v14 = iadd v12, v13
+            v15 = iadd.i32 v4, v13
+[UJ#1b]   jump ebb5(v13)
 
-ebb5(v16: i32):
-    v17 = iadd.i32 v14, v4
-    brz.i32 v1, ebb4(v16)
-    jump ebb6(v16)
+         ebb5(v16: i32):
+            v17 = iadd.i32 v14, v4
+[SBzero#18] brz.i32 v1, ebb4(v16)
+[UJ#1b]   jump ebb6(v16)
 
-ebb6(v18: i32):
-    v19 = iadd v18, v2
-    v20 = iadd.i32 v2, v3
-    brz.i32 v1, ebb1(v20)
-    return v19
+         ebb6(v18: i32):
+            v19 = iadd v18, v2
+            v20 = iadd.i32 v2, v3
+[SBzero#18] brz.i32 v1, ebb1(v20)
+[Iret#19]   return v19
 }
 
 ; sameln: function %complex
@@ -53,11 +53,13 @@ ebb6(v18: i32):
 ; nextln: 
 ; nextln: ebb1(v1: i32):
 ; nextln:     v4 = iadd.i32 v2, v1
+; nextln:     brz v1, ebb7(v2)
+; nextln:     jump ebb8(v4)
+; nextln: 
+; nextln: ebb7(v21: i32):
 ; nextln:     v8 = iadd.i32 v6, v1
-; nextln:     v11 = iadd v1, v4
-; nextln:     brz v1, ebb2(v2)
-; nextln:     v15 = iadd v4, v13
-; nextln:     jump ebb4(v4)
+; nextln:     v11 = iadd.i32 v1, v4
+; nextln:     jump ebb2(v21)
 ; nextln: 
 ; nextln: ebb2(v5: i32):
 ; nextln:     v7 = iadd v5, v4
@@ -68,6 +70,10 @@ ebb6(v18: i32):
 ; nextln:     brz.i32 v1, ebb2(v9)
 ; nextln:     jump ebb6(v10)
 ; nextln: 
+; nextln: ebb8(v22: i32):
+; nextln:     v15 = iadd.i32 v4, v13
+; nextln:     jump ebb4(v22)
+; nextln: 
 ; nextln: ebb4(v12: i32):
 ; nextln:     v14 = iadd v12, v13
 ; nextln:     jump ebb5(v13)
diff --git a/cranelift/filetests/licm/critical-edge.clif b/cranelift/filetests/licm/critical-edge.clif
new file mode 100644
index 0000000000..1c69f6364e
--- /dev/null
+++ b/cranelift/filetests/licm/critical-edge.clif
@@ -0,0 +1,42 @@
+test licm
+target riscv32
+
+; The loop in this function is entered from a critical edge.
+
+function %critical_edge(i32, i32) -> i32 {
+
+            ebb0(v0: i32, v7: i32):
+[SBzero#38]   brnz v7, ebb1(v0)
+[Iret#19]     return v0
+
+            ebb1(v1: i32):
+              v2 = iconst.i32 1
+              v3 = iconst.i32 2
+              v4 = iadd v2, v3
+[SBzero#18]   brz v1, ebb2(v1)
+              v5 = isub v1, v2
+[UJ#1b]       jump ebb1(v5)
+
+            ebb2(v6: i32):
+[Iret#19]     return v6
+
+}
+; sameln: function %critical_edge
+; nextln: ebb0(v0: i32, v7: i32):
+; nextln:     brnz v7, ebb3(v0)
+; nextln:     return v0
+; nextln: 
+; nextln: ebb3(v8: i32):
+; nextln:     v2 = iconst.i32 1
+; nextln:     v3 = iconst.i32 2
+; nextln:     v4 = iadd v2, v3
+; nextln:     jump ebb1(v8)
+; nextln: 
+; nextln: ebb1(v1: i32):
+; nextln:     brz v1, ebb2(v1)
+; nextln:     v5 = isub v1, v2
+; nextln:     jump ebb1(v5)
+; nextln: 
+; nextln: ebb2(v6: i32):
+; nextln:     return v6
+; nextln: }
diff --git a/lib/codegen/src/licm.rs b/lib/codegen/src/licm.rs
index 0cba66b53b..b203fdfe9e 100644
--- a/lib/codegen/src/licm.rs
+++ b/lib/codegen/src/licm.rs
@@ -1,6 +1,6 @@
 //! A Loop Invariant Code Motion optimization pass
 
-use cursor::{Cursor, FuncCursor, EncCursor};
+use cursor::{Cursor, EncCursor, FuncCursor};
 use dominator_tree::DominatorTree;
 use entity::{EntityList, ListPool};
 use flowgraph::{BasicBlock, ControlFlowGraph};
@@ -111,21 +111,24 @@ fn has_pre_header(
     header: Ebb,
 ) -> Option<(Ebb, Inst)> {
     let mut result = None;
-    let mut found = false;
     for BasicBlock {
         ebb: pred_ebb,
-        inst: last_inst,
+        inst: branch_inst,
     } in cfg.pred_iter(header)
     {
         // We only count normal edges (not the back edges)
-        if !domtree.dominates(header, last_inst, layout) {
-            if found {
+        if !domtree.dominates(header, branch_inst, layout) {
+            if result.is_some() {
                 // We have already found one, there are more than one
                 return None;
-            } else {
-                result = Some((pred_ebb, last_inst));
-                found = true;
             }
+            if branch_inst != layout.last_inst(pred_ebb).unwrap()
+                || cfg.succ_iter(pred_ebb).nth(1).is_some()
+            {
+                // It's along a critical edge, so don't use it.
+                return None;
+            }
+            result = Some((pred_ebb, branch_inst));
         }
     }
     result