Split edges after transformations that might introduce new critical edges.

Register allocation needs all critical edges to be split. Additionally,
some of our optimizations rely on edge split form (such as switch
rewriting). Therefore, we should maintain edge split form throughout.

This change adds asserts to isConsistentSSA to verify that we
have no critical edges. Additionally, it resplits edges after the
transformations that can introduce new critical edges.

The other option is to update all rewriting to not rely on
edge split form and only split critical edges before register
allocations. We can consider going that route as well, but this
is the quick fix for critical edges making its way to the
register allocator resulting in wrong code generation.

R=herhut@google.com, sgjesse@google.com

Change-Id: Ie7b80503cce91bf07f15b8e66eeb8c5bd59a19a8
diff --git a/src/main/java/com/android/tools/r8/ir/code/BasicBlock.java b/src/main/java/com/android/tools/r8/ir/code/BasicBlock.java
index 2324f8e..54f961d 100644
--- a/src/main/java/com/android/tools/r8/ir/code/BasicBlock.java
+++ b/src/main/java/com/android/tools/r8/ir/code/BasicBlock.java
@@ -1171,11 +1171,9 @@
     // to be changed.
     for (BasicBlock catchSuccessor : catchSuccessors) {
       catchSuccessor.splitCriticalExceptionEdges(
+          code.getHighestBlockNumber() + 1,
           code.valueNumberGenerator,
-          newBlock -> {
-            newBlock.setNumber(code.getHighestBlockNumber() + 1);
-            blockIterator.add(newBlock);
-          });
+          blockIterator::add);
     }
   }
 
@@ -1191,11 +1189,11 @@
    * Note that if there are any other phis defined on this block, they remain valid, since
    * this method does not affect incoming edges in any way, and just adds new blocks with
    * MoveException and Goto.
-   *
-   * NOTE: onNewBlock must assign block number to the newly created block.
    */
-  public void splitCriticalExceptionEdges(
-      ValueNumberGenerator valueNumberGenerator, Consumer<BasicBlock> onNewBlock) {
+  public int splitCriticalExceptionEdges(
+      int nextBlockNumber,
+      ValueNumberGenerator valueNumberGenerator,
+      Consumer<BasicBlock> onNewBlock) {
     List<BasicBlock> predecessors = this.getPredecessors();
     boolean hasMoveException = entry().isMoveException();
     MoveException move = null;
@@ -1216,6 +1214,7 @@
             "Invalid block structure: catch block reachable via non-exceptional flow.");
       }
       BasicBlock newBlock = new BasicBlock();
+      newBlock.setNumber(nextBlockNumber++);
       newPredecessors.add(newBlock);
       if (hasMoveException) {
         Value value = new Value(valueNumberGenerator.next(), MoveType.OBJECT, move.getLocalInfo());
@@ -1244,6 +1243,7 @@
       phi.addOperands(values);
       move.outValue().replaceUsers(phi);
     }
+    return nextBlockNumber;
   }
 
   /**
diff --git a/src/main/java/com/android/tools/r8/ir/code/IRCode.java b/src/main/java/com/android/tools/r8/ir/code/IRCode.java
index d9fdaf2..477c9d7 100644
--- a/src/main/java/com/android/tools/r8/ir/code/IRCode.java
+++ b/src/main/java/com/android/tools/r8/ir/code/IRCode.java
@@ -40,6 +40,42 @@
     this.valueNumberGenerator = valueNumberGenerator;
   }
 
+  public void splitCriticalEdges() {
+    List<BasicBlock> newBlocks = new ArrayList<>();
+    int nextBlockNumber = getHighestBlockNumber() + 1;
+    for (BasicBlock block : blocks) {
+      // We are using a spilling register allocator that might need to insert moves at
+      // all critical edges, so we always split them all.
+      List<BasicBlock> predecessors = block.getPredecessors();
+      if (predecessors.size() <= 1) {
+        continue;
+      }
+      // If any of the edges to the block are critical, we need to insert new blocks on each
+      // containing the move-exception instruction which must remain the first instruction.
+      if (block.entry() instanceof MoveException) {
+        nextBlockNumber = block.splitCriticalExceptionEdges(
+            nextBlockNumber, valueNumberGenerator, newBlocks::add);
+        continue;
+      }
+      for (int predIndex = 0; predIndex < predecessors.size(); predIndex++) {
+        BasicBlock pred = predecessors.get(predIndex);
+        if (!pred.hasOneNormalExit()) {
+          // Critical edge: split it and inject a new block into which the
+          // phi moves can be inserted. The new block is created with the
+          // correct predecessor and successor structure. It is inserted
+          // at the end of the list of blocks disregarding branching
+          // structure.
+          BasicBlock newBlock = BasicBlock.createGotoBlock(block, nextBlockNumber++);
+          newBlocks.add(newBlock);
+          pred.replaceSuccessor(block, newBlock);
+          newBlock.getPredecessors().add(pred);
+          predecessors.set(predIndex, newBlock);
+        }
+      }
+    }
+    blocks.addAll(newBlocks);
+  }
+
   /**
    * Trace blocks and attempt to put fallthrough blocks immediately after the block that
    * falls through. When we fail to do that we create a new fallthrough block with an explicit
@@ -158,7 +194,10 @@
   }
 
   public boolean isConsistentSSA() {
-    assert isConsistentGraph() && consistentDefUseChains() && validThrowingInstructions();
+    assert isConsistentGraph();
+    assert consistentDefUseChains();
+    assert validThrowingInstructions();
+    assert noCriticalEdges();
     return true;
   }
 
@@ -170,6 +209,26 @@
     return true;
   }
 
+  private boolean noCriticalEdges() {
+    for (BasicBlock block : blocks) {
+      List<BasicBlock> predecessors = block.getPredecessors();
+      if (predecessors.size() <= 1) {
+        continue;
+      }
+      if (block.entry() instanceof MoveException) {
+        assert false;
+        return false;
+      }
+      for (int predIndex = 0; predIndex < predecessors.size(); predIndex++) {
+        if (!predecessors.get(predIndex).hasOneNormalExit()) {
+          assert false;
+          return false;
+        }
+      }
+    }
+    return true;
+  }
+
   private boolean consistentDefUseChains() {
     Set<Value> values = new HashSet<>();
 
diff --git a/src/main/java/com/android/tools/r8/ir/conversion/IRBuilder.java b/src/main/java/com/android/tools/r8/ir/conversion/IRBuilder.java
index 35cd7fa..08bc75e 100644
--- a/src/main/java/com/android/tools/r8/ir/conversion/IRBuilder.java
+++ b/src/main/java/com/android/tools/r8/ir/conversion/IRBuilder.java
@@ -383,13 +383,13 @@
     // same phi moves in multiple blocks.
     joinPredecessorsWithIdenticalPhis();
 
-    // Split critical edges to make sure that we have a place to insert phi moves if
-    // necessary.
-    splitCriticalEdges();
-
     // Package up the IR code.
     IRCode ir = new IRCode(method, blocks, valueNumberGenerator);
 
+    // Split critical edges to make sure that we have a place to insert phi moves if
+    // necessary.
+    ir.splitCriticalEdges();
+
     if (options.testing.invertConditionals) {
       invertConditionalsForTesting(ir);
     }
@@ -1943,45 +1943,6 @@
     blocks.addAll(blocksToAdd);
   }
 
-  private void splitCriticalEdges() {
-    List<BasicBlock> newBlocks = new ArrayList<>();
-    for (BasicBlock block : blocks) {
-      // We are using a spilling register allocator that might need to insert moves at
-      // all critical edges, so we always split them all.
-      List<BasicBlock> predecessors = block.getPredecessors();
-      if (predecessors.size() <= 1) {
-        continue;
-      }
-      // If any of the edges to the block are critical, we need to insert new blocks on each
-      // containing the move-exception instruction which must remain the first instruction.
-      if (block.entry() instanceof MoveException) {
-        block.splitCriticalExceptionEdges(valueNumberGenerator,
-            newBlock -> {
-              newBlock.setNumber(blocks.size() + newBlocks.size());
-              newBlocks.add(newBlock);
-            });
-        continue;
-      }
-      for (int predIndex = 0; predIndex < predecessors.size(); predIndex++) {
-        BasicBlock pred = predecessors.get(predIndex);
-        if (!pred.hasOneNormalExit()) {
-          // Critical edge: split it and inject a new block into which the
-          // phi moves can be inserted. The new block is created with the
-          // correct predecessor and successor structure. It is inserted
-          // at the end of the list of blocks disregarding branching
-          // structure.
-          int blockNumber = blocks.size() + newBlocks.size();
-          BasicBlock newBlock = BasicBlock.createGotoBlock(block, blockNumber);
-          newBlocks.add(newBlock);
-          pred.replaceSuccessor(block, newBlock);
-          newBlock.getPredecessors().add(pred);
-          predecessors.set(predIndex, newBlock);
-        }
-      }
-    }
-    blocks.addAll(newBlocks);
-  }
-
   // Other stuff.
 
   boolean isIntegerType(NumericType type) {
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/CodeRewriter.java b/src/main/java/com/android/tools/r8/ir/optimize/CodeRewriter.java
index 7cb186f..193770d 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/CodeRewriter.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/CodeRewriter.java
@@ -643,6 +643,12 @@
         }
       }
     }
+    // Rewriting of switches introduces new branching structure. It relies on critical edges
+    // being split on the way in but does not maintain this property. We therefore split
+    // critical edges at exit.
+    code.splitCriticalEdges();
+    code.traceBlocks();
+    assert code.isConsistentSSA();
   }
 
   /**
diff --git a/src/main/java/com/android/tools/r8/shaking/protolite/ProtoLitePruner.java b/src/main/java/com/android/tools/r8/shaking/protolite/ProtoLitePruner.java
index e7baec9..9ccdedc 100644
--- a/src/main/java/com/android/tools/r8/shaking/protolite/ProtoLitePruner.java
+++ b/src/main/java/com/android/tools/r8/shaking/protolite/ProtoLitePruner.java
@@ -351,6 +351,11 @@
     rewriteVisitCase(visitCase, code);
     assert code.isConsistentSSA();
     rewriteMergeCase(mergeCase, instanceType, code);
+    // Rewriting of the merge case changes branching structure and adds more blocks that target
+    // the fallthrough label. This can introduce critical edges. Therefore, we split critical
+    // edges to maintain our edge-split form.
+    code.splitCriticalEdges();
+    code.traceBlocks();
     assert code.isConsistentSSA();
   }