From f2bac088f1df1df9cecd89c0ff1d47b97420f32a Mon Sep 17 00:00:00 2001
From: Alexey Andreev <konsoletyper@gmail.com>
Date: Mon, 22 Mar 2021 22:09:41 +0300
Subject: [PATCH] Fix more complex cases of irreducible CFGs during coroutine
 transformation

---
 .../transform/CoroutineTransformation.java    | 25 +++++++++++++++++--
 .../teavm/model/util/BasicBlockMapper.java    | 14 +++++++----
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/core/src/main/java/org/teavm/backend/lowlevel/transform/CoroutineTransformation.java b/core/src/main/java/org/teavm/backend/lowlevel/transform/CoroutineTransformation.java
index 9bdd7bb39..644aeb3fa 100644
--- a/core/src/main/java/org/teavm/backend/lowlevel/transform/CoroutineTransformation.java
+++ b/core/src/main/java/org/teavm/backend/lowlevel/transform/CoroutineTransformation.java
@@ -15,8 +15,10 @@
  */
 package org.teavm.backend.lowlevel.transform;
 
+import com.carrotsearch.hppc.IntHashSet;
 import com.carrotsearch.hppc.IntIntHashMap;
 import com.carrotsearch.hppc.IntIntMap;
+import com.carrotsearch.hppc.IntSet;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.BitSet;
@@ -31,6 +33,7 @@ import org.teavm.common.GraphUtils;
 import org.teavm.model.BasicBlock;
 import org.teavm.model.ClassReader;
 import org.teavm.model.ClassReaderSource;
+import org.teavm.model.Incoming;
 import org.teavm.model.Instruction;
 import org.teavm.model.MethodDescriptor;
 import org.teavm.model.MethodReader;
@@ -475,6 +478,7 @@ public class CoroutineTransformation {
             weights[i] = program.basicBlockAt(i).instructionCount();
         }
         GraphUtils.splitIrreducibleGraph(graph, weights, splittingBackend);
+        new PhiUpdater().updatePhis(program, parameterCount + 1);
     }
 
     class SplittingBackend implements GraphSplittingBackend {
@@ -482,14 +486,17 @@ public class CoroutineTransformation {
         public int[] split(int[] domain, int[] nodes) {
             int[] copies = new int[nodes.length];
             IntIntMap map = new IntIntHashMap();
+            IntSet nodeSet = IntHashSet.from(nodes);
+            List<List<Incoming>> outputs = ProgramUtils.getPhiOutputs(program);
             for (int i = 0; i < nodes.length; ++i) {
                 int node = nodes[i];
                 BasicBlock block = program.basicBlockAt(node);
                 BasicBlock blockCopy = program.createBasicBlock();
                 ProgramUtils.copyBasicBlock(block, blockCopy);
                 copies[i] = blockCopy.getIndex();
-                map.put(nodes[i], copies[i] + 1);
+                map.put(node, copies[i] + 1);
             }
+
             BasicBlockMapper copyBlockMapper = new BasicBlockMapper((int block) -> {
                 int mappedIndex = map.get(block);
                 return mappedIndex == 0 ? block : mappedIndex - 1;
@@ -498,8 +505,22 @@ public class CoroutineTransformation {
                 copyBlockMapper.transform(program.basicBlockAt(copy));
             }
             for (int domainNode : domain) {
-                copyBlockMapper.transform(program.basicBlockAt(domainNode));
+                copyBlockMapper.transformWithoutPhis(program.basicBlockAt(domainNode));
             }
+
+            for (int i = 0; i < nodes.length; ++i) {
+                int node = nodes[i];
+                BasicBlock blockCopy = program.basicBlockAt(copies[i]);
+                for (Incoming output : outputs.get(node)) {
+                    if (!nodeSet.contains(output.getPhi().getBasicBlock().getIndex())) {
+                        Incoming outputCopy = new Incoming();
+                        outputCopy.setSource(blockCopy);
+                        outputCopy.setValue(output.getValue());
+                        output.getPhi().getIncomings().add(outputCopy);
+                    }
+                }
+            }
+
             return copies;
         }
     }
diff --git a/core/src/main/java/org/teavm/model/util/BasicBlockMapper.java b/core/src/main/java/org/teavm/model/util/BasicBlockMapper.java
index c2b92591e..931aac1f1 100644
--- a/core/src/main/java/org/teavm/model/util/BasicBlockMapper.java
+++ b/core/src/main/java/org/teavm/model/util/BasicBlockMapper.java
@@ -42,16 +42,20 @@ public class BasicBlockMapper extends AbstractInstructionVisitor {
     }
 
     public void transform(BasicBlock block) {
-        Instruction lastInsn = block.getLastInstruction();
-        if (lastInsn != null) {
-            lastInsn.acceptVisitor(this);
-        }
-
+        transformWithoutPhis(block);
         for (Phi phi : block.getPhis()) {
             for (Incoming incoming : phi.getIncomings()) {
                 incoming.setSource(map(incoming.getSource()));
             }
         }
+    }
+
+    public void transformWithoutPhis(BasicBlock block) {
+        Instruction lastInsn = block.getLastInstruction();
+        if (lastInsn != null) {
+            lastInsn.acceptVisitor(this);
+        }
+
         for (TryCatchBlock tryCatch : block.getTryCatchBlocks()) {
             tryCatch.setHandler(map(tryCatch.getHandler()));
         }