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();
   }