Fix AssertionRewriter to split the block and use precise types

Bug: 142064862
Change-Id: I66e045d55419408c0d85185360709f91387ce4cc
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 1f36d25..7c1590c 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
@@ -3747,31 +3747,43 @@
       return;
     }
 
-    InstructionListIterator iterator = code.instructionListIterator();
-    while (iterator.hasNext()) {
-      Instruction current = iterator.next();
-      if (current.isInvokeMethod()) {
-        DexMethod invokedMethod = current.asInvokeMethod().getInvokedMethod();
-        if (invokedMethod == dexItemFactory.assertionErrorMethods.initMessageAndCause) {
-          // Rewrite calls to new AssertionError(message, cause) to new AssertionError(message)
-          // and then initCause(cause).
-          List<Value> inValues = current.inValues();
-          assert inValues.size() == 3; // receiver, message, cause
+    ListIterator<BasicBlock> blockIterator = code.listIterator();
+    while (blockIterator.hasNext()) {
+      BasicBlock block = blockIterator.next();
+      InstructionListIterator insnIterator = block.listIterator(code);
+      while (insnIterator.hasNext()) {
+        Instruction current = insnIterator.next();
+        if (current.isInvokeMethod()) {
+          DexMethod invokedMethod = current.asInvokeMethod().getInvokedMethod();
+          if (invokedMethod == dexItemFactory.assertionErrorMethods.initMessageAndCause) {
+            // Rewrite calls to new AssertionError(message, cause) to new AssertionError(message)
+            // and then initCause(cause).
+            List<Value> inValues = current.inValues();
+            assert inValues.size() == 3; // receiver, message, cause
 
-          // Remove cause from the constructor call
-          List<Value> newInitInValues = inValues.subList(0, 2);
-          iterator.replaceCurrentInstruction(
-              new InvokeDirect(dexItemFactory.assertionErrorMethods.initMessage, null,
-                  newInitInValues));
+            // Remove cause from the constructor call
+            List<Value> newInitInValues = inValues.subList(0, 2);
+            insnIterator.replaceCurrentInstruction(
+                new InvokeDirect(
+                    dexItemFactory.assertionErrorMethods.initMessage, null, newInitInValues));
 
-          // On API 15 and older we cannot use initCause because of a bug in AssertionError.
-          if (options.canInitCauseAfterAssertionErrorObjectConstructor()) {
-            // Add a call to Throwable.initCause(cause)
-            List<Value> initCauseArguments = Arrays.asList(inValues.get(0), inValues.get(2));
-            InvokeVirtual initCause = new InvokeVirtual(dexItemFactory.throwableMethods.initCause,
-                code.createValue(TypeLatticeElement.SINGLE), initCauseArguments);
-            initCause.setPosition(current.getPosition());
-            iterator.add(initCause);
+            // On API 15 and older we cannot use initCause because of a bug in AssertionError.
+            if (options.canInitCauseAfterAssertionErrorObjectConstructor()) {
+              // Add a call to Throwable.initCause(cause)
+              if (block.hasCatchHandlers()) {
+                insnIterator = insnIterator.split(code, blockIterator).listIterator(code);
+              }
+              List<Value> initCauseArguments = Arrays.asList(inValues.get(0), inValues.get(2));
+              InvokeVirtual initCause =
+                  new InvokeVirtual(
+                      dexItemFactory.throwableMethods.initCause,
+                      code.createValue(
+                          TypeLatticeElement.fromDexType(
+                              dexItemFactory.throwableType, Nullability.maybeNull(), appView)),
+                      initCauseArguments);
+              initCause.setPosition(current.getPosition());
+              insnIterator.add(initCause);
+            }
           }
         }
       }