C: fix insertion of shadow stack instructions

core/src/main/java/org/teavm/model/lowlevel/GCShadowStackContributor.java

@@ -22,14 +22,11 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.BitSet;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.Comparator;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
-import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import org.teavm.common.DisjointSet;
 import org.teavm.common.DominatorTree;
 import org.teavm.common.Graph;
@@ -102,14 +99,11 @@ public class GCShadowStackContributor {
         // If a variable is spilled to stack, then phi which takes this variable as input also spilled to stack
         // If all of phi inputs are spilled to stack, then we don't need to insert spilling instruction
         // for this phi.
-        List<Set<Phi>> destinationPhis = getDestinationPhis(program);
+        Graph cfg = ProgramUtils.buildControlFlowGraph(program);
-        int[] inputCount = getInputCount(program);
+        DominatorTree dom = GraphUtils.buildDominatorTree(cfg);
-        boolean[] autoSpilled = new boolean[spilled.length];
+        boolean[] autoSpilled = new SpilledPhisFinder(liveInInformation, dom, program).find();
-        for (int i = 0; i < spilled.length; ++i) {
-            findAutoSpilledPhis(spilled, destinationPhis, inputCount, autoSpilled, i);
-        }
 
-        List<Map<Instruction, int[]>> liveInStores = reduceGCRootStores(program, usedColors, liveInInformation,
+        List<Map<Instruction, int[]>> liveInStores = reduceGCRootStores(dom, program, usedColors, liveInInformation,
                 colors, autoSpilled);
         putLiveInGCRoots(program, liveInStores);
 
@@ -131,26 +125,6 @@ public class GCShadowStackContributor {
         return variableClasses.pack(program.variableCount());
     }
 
-    private void findAutoSpilledPhis(boolean[] spilled, List<Set<Phi>> destinationPhis, int[] inputCount,
-            boolean[] autoSpilled, int i) {
-        if (!spilled[i]) {
-            return;
-        }
-        Set<Phi> phis = destinationPhis.get(i);
-        if (phis != null) {
-            for (Phi phi : destinationPhis.get(i)) {
-                int destination = phi.getReceiver().getIndex();
-                autoSpilled[destination] = --inputCount[destination] == 0;
-                if (!spilled[destination]) {
-                    spilled[destination] = true;
-                    if (i > destination) {
-                        findAutoSpilledPhis(spilled, destinationPhis, inputCount, autoSpilled, destination);
-                    }
-                }
-            }
-        }
-    }
-
     private List<Map<Instruction, BitSet>> findCallSiteLiveIns(Program program, MethodReader method) {
         boolean[] nativePointers = nativePointerFinder.findNativePointers(method.getReference(), program);
 
@@ -243,41 +217,7 @@ public class GCShadowStackContributor {
         return affectedVariables;
     }
 
-    private List<Set<Phi>> getDestinationPhis(Program program) {
+    private List<Map<Instruction, int[]>> reduceGCRootStores(DominatorTree dom, Program program, int usedColors,
-        List<Set<Phi>> destinationPhis = new ArrayList<>();
-        destinationPhis.addAll(Collections.nCopies(program.variableCount(), null));
-
-        for (int i = 0; i < program.basicBlockCount(); ++i) {
-            BasicBlock block = program.basicBlockAt(i);
-            for (Phi phi : block.getPhis()) {
-                for (Incoming incoming : phi.getIncomings()) {
-                    Set<Phi> phis = destinationPhis.get(incoming.getValue().getIndex());
-                    if (phis == null) {
-                        phis = new LinkedHashSet<>();
-                        destinationPhis.set(incoming.getValue().getIndex(), phis);
-                    }
-                    phis.add(phi);
-                }
-            }
-        }
-
-        return destinationPhis;
-    }
-
-    private int[] getInputCount(Program program) {
-        int[] inputCount = new int[program.variableCount()];
-
-        for (int i = 0; i < program.basicBlockCount(); ++i) {
-            BasicBlock block = program.basicBlockAt(i);
-            for (Phi phi : block.getPhis()) {
-                inputCount[phi.getReceiver().getIndex()] = phi.getIncomings().size();
-            }
-        }
-
-        return inputCount;
-    }
-
-    private List<Map<Instruction, int[]>> reduceGCRootStores(Program program, int usedColors,
             List<Map<Instruction, BitSet>> liveInInformation, int[] colors, boolean[] autoSpilled) {
         class Step {
             private final int node;
@@ -292,9 +232,7 @@ public class GCShadowStackContributor {
             slotsToUpdate.add(new LinkedHashMap<>());
         }
 
-        Graph cfg = ProgramUtils.buildControlFlowGraph(program);
+        Graph domGraph = GraphUtils.buildDominatorGraph(dom, program.basicBlockCount());
-        DominatorTree dom = GraphUtils.buildDominatorTree(cfg);
-        Graph domGraph = GraphUtils.buildDominatorGraph(dom, cfg.size());
 
         Step[] stack = new Step[program.basicBlockCount() * 2];
         int head = 0;
@@ -374,7 +312,7 @@ public class GCShadowStackContributor {
             Instruction[] callSiteLocations = updatesByIndex.keySet().toArray(new Instruction[0]);
             ObjectIntMap<Instruction> instructionIndexes = getInstructionIndexes(block);
             Arrays.sort(callSiteLocations, Comparator.comparing(instructionIndexes::get));
-            for (Instruction callSiteLocation : updatesByIndex.keySet()) {
+            for (Instruction callSiteLocation : callSiteLocations) {
                 int[] updates = updatesByIndex.get(callSiteLocation);
                 storeLiveIns(block, callSiteLocation, updates);
             }

core/src/main/java/org/teavm/model/lowlevel/SpilledPhisFinder.java

@@ -0,0 +1,131 @@
+/*
+ *  Copyright 2019 konsoletyper.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+package org.teavm.model.lowlevel;
+
+import com.carrotsearch.hppc.IntHashSet;
+import com.carrotsearch.hppc.IntSet;
+import java.util.Arrays;
+import java.util.BitSet;
+import java.util.List;
+import java.util.Map;
+import org.teavm.common.DominatorTree;
+import org.teavm.model.BasicBlock;
+import org.teavm.model.Incoming;
+import org.teavm.model.Instruction;
+import org.teavm.model.Phi;
+import org.teavm.model.Program;
+
+class SpilledPhisFinder {
+    private static final byte VISITING = 1;
+    private static final byte VISITED = 2;
+    private DominatorTree dom;
+    boolean[] autoSpilled;
+    byte[] status;
+    int[][] variableSpilledBlocks;
+    Phi[] definingPhis;
+    int variableCount;
+
+    SpilledPhisFinder(List<Map<Instruction, BitSet>> liveInInformation, DominatorTree dom, Program program) {
+        this.dom = dom;
+        variableCount = program.variableCount();
+        autoSpilled = new boolean[variableCount];
+        status = new byte[variableCount];
+        variableSpilledBlocks = variableSpilledBlocks(liveInInformation, variableCount);
+        definingPhis = findPhis(program);
+    }
+
+    private static int[][] variableSpilledBlocks(List<Map<Instruction, BitSet>> liveInInformation, int count) {
+        IntSet[] builder = new IntSet[count];
+        for (int b = 0; b < liveInInformation.size(); b++) {
+            Map<Instruction, BitSet> blockLiveIn = liveInInformation.get(b);
+            for (BitSet liveVarsSet : blockLiveIn.values()) {
+                for (int v = liveVarsSet.nextSetBit(0); v >= 0; v = liveVarsSet.nextSetBit(v + 1)) {
+                    if (builder[v] == null) {
+                        builder[v] = new IntHashSet();
+                    }
+                    builder[v].add(b);
+                }
+            }
+        }
+
+        int[][] result = new int[builder.length][];
+        for (int v = 0; v < result.length; ++v) {
+            if (builder[v] != null) {
+                result[v] = builder[v].toArray();
+                Arrays.sort(result[v]);
+            }
+        }
+        return result;
+    }
+
+    private static Phi[] findPhis(Program program) {
+        Phi[] result = new Phi[program.variableCount()];
+        for (BasicBlock block : program.getBasicBlocks()) {
+            for (Phi phi : block.getPhis()) {
+                result[phi.getReceiver().getIndex()] = phi;
+            }
+        }
+        return result;
+    }
+
+    boolean[] find() {
+        for (int v = 0; v < variableCount; ++v) {
+            isAutoSpilled(v);
+        }
+        return autoSpilled;
+    }
+
+    private boolean isAutoSpilled(int v) {
+        if (status[v] == VISITED) {
+            return autoSpilled[v];
+        }
+
+        if (status[v] == VISITING) {
+            return false;
+        }
+
+        Phi definingPhi = definingPhis[v];
+        if (definingPhi == null) {
+            status[v] = VISITED;
+            return false;
+        }
+
+        status[v] = VISITING;
+
+        boolean result = true;
+        for (Incoming incoming : definingPhi.getIncomings()) {
+            if (!isAutoSpilled(incoming.getValue().getIndex())) {
+                int[] spilledAt = variableSpilledBlocks[incoming.getValue().getIndex()];
+                result = false;
+                if (spilledAt != null) {
+                    for (int spilledAtBlock : spilledAt) {
+                        if (dom.dominates(spilledAtBlock, incoming.getSource().getIndex())) {
+                            result = true;
+                            break;
+                        }
+                    }
+                }
+                if (!result) {
+                    break;
+                }
+            }
+        }
+
+        autoSpilled[v] = result;
+        status[v] = VISITED;
+        return result;
+    }
+}