From 70d0658b473d89aa3948f1ef1e2b6d101733ea19 Mon Sep 17 00:00:00 2001
From: Alexey Andreev <konsoletyper@gmail.com>
Date: Mon, 28 Aug 2023 17:16:31 +0200
Subject: [PATCH] wasm/c: fix bugs in coroutine transformation

---
 .../transform/CoroutineTransformation.java    | 19 +++++++++++++++----
 .../RedundantJumpElimination.java             | 17 +++++++++--------
 2 files changed, 24 insertions(+), 12 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 5335330e7..1be2a83f3 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
@@ -17,7 +17,6 @@ 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;
@@ -57,6 +56,7 @@ import org.teavm.model.instructions.MonitorEnterInstruction;
 import org.teavm.model.instructions.NullConstantInstruction;
 import org.teavm.model.instructions.SwitchInstruction;
 import org.teavm.model.instructions.SwitchTableEntry;
+import org.teavm.model.optimization.RedundantJumpElimination;
 import org.teavm.model.util.BasicBlockMapper;
 import org.teavm.model.util.BasicBlockSplitter;
 import org.teavm.model.util.DefinitionExtractor;
@@ -125,6 +125,7 @@ public class CoroutineTransformation {
         }
         splitter.fixProgram();
         processIrreducibleCfg();
+        RedundantJumpElimination.optimize(program);
     }
 
     private void createSplitPrologue() {
@@ -485,9 +486,9 @@ public class CoroutineTransformation {
         @Override
         public int[] split(int[] domain, int[] nodes) {
             int[] copies = new int[nodes.length];
-            IntIntMap map = new IntIntHashMap();
+            var map = new IntIntHashMap();
             IntSet nodeSet = IntHashSet.from(nodes);
-            List<List<Incoming>> outputs = ProgramUtils.getPhiOutputs(program);
+            var outputs = ProgramUtils.getPhiOutputs(program);
             for (int i = 0; i < nodes.length; ++i) {
                 int node = nodes[i];
                 BasicBlock block = program.basicBlockAt(node);
@@ -497,7 +498,7 @@ public class CoroutineTransformation {
                 map.put(node, copies[i] + 1);
             }
 
-            BasicBlockMapper copyBlockMapper = new BasicBlockMapper((int block) -> {
+            var copyBlockMapper = new BasicBlockMapper((int block) -> {
                 int mappedIndex = map.get(block);
                 return mappedIndex == 0 ? block : mappedIndex - 1;
             });
@@ -508,6 +509,7 @@ public class CoroutineTransformation {
                 copyBlockMapper.transformWithoutPhis(program.basicBlockAt(domainNode));
             }
 
+            var domainSet = IntHashSet.from(domain);
             for (int i = 0; i < nodes.length; ++i) {
                 int node = nodes[i];
                 BasicBlock blockCopy = program.basicBlockAt(copies[i]);
@@ -519,6 +521,15 @@ public class CoroutineTransformation {
                         output.getPhi().getIncomings().add(outputCopy);
                     }
                 }
+
+                var block = program.basicBlockAt(node);
+                for (var phi : block.getPhis()) {
+                    phi.getIncomings().removeIf(incoming -> domainSet.contains(incoming.getSource().getIndex()));
+                }
+
+                for (var phi : blockCopy.getPhis()) {
+                    phi.getIncomings().removeIf(incoming -> !domainSet.contains(incoming.getSource().getIndex()));
+                }
             }
 
             return copies;
diff --git a/core/src/main/java/org/teavm/model/optimization/RedundantJumpElimination.java b/core/src/main/java/org/teavm/model/optimization/RedundantJumpElimination.java
index 3ec80b959..54f56c61b 100644
--- a/core/src/main/java/org/teavm/model/optimization/RedundantJumpElimination.java
+++ b/core/src/main/java/org/teavm/model/optimization/RedundantJumpElimination.java
@@ -71,18 +71,13 @@ public class RedundantJumpElimination implements MethodOptimization {
                 assign.setAssignee(incoming.getValue());
                 block.add(assign);
             }
-            while (target.getFirstInstruction() != null) {
-                Instruction instruction = target.getFirstInstruction();
-                instruction.delete();
-                block.add(instruction);
-            }
 
-            Instruction lastInsn = block.getLastInstruction();
+            var lastInsn = target.getLastInstruction();
             if (lastInsn != null) {
                 lastInsn.acceptVisitor(transitionExtractor);
-                BasicBlock[] successors = transitionExtractor.getTargets();
+                var successors = transitionExtractor.getTargets();
                 if (successors != null) {
-                    for (BasicBlock successor : successors) {
+                    for (var successor : successors) {
                         successor.getPhis().stream()
                                 .flatMap(phi -> phi.getIncomings().stream())
                                 .filter(incoming -> incoming.getSource() == target)
@@ -91,6 +86,12 @@ public class RedundantJumpElimination implements MethodOptimization {
                 }
             }
 
+            while (target.getFirstInstruction() != null) {
+                Instruction instruction = target.getFirstInstruction();
+                instruction.delete();
+                block.add(instruction);
+            }
+
             for (TryCatchBlock tryCatch : target.getTryCatchBlocks()) {
                 for (Phi phi : tryCatch.getHandler().getPhis()) {
                     phi.getIncomings().removeIf(incoming -> incoming.getSource() == target);