Budget overlapping Lua inline results

Author: Frotty

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/translation/imoptimizer/ImInliner.java

@@ -506,6 +506,21 @@ public void visit(ImVarAccess access) {
         });
     }
 
+    private static int statementExpressionResultSlots(ImStmt statement) {
+        int[] result = {0};
+        statement.accept(new ImStmt.DefaultVisitor() {
+            @Override
+            public void visit(ImStatementExpr expression) {
+                super.visit(expression);
+                ImType type = expression.getExpr().attrTyp();
+                if (!(type instanceof ImVoid)) {
+                    result[0] += ImHelper.flattenedJassArity(type);
+                }
+            }
+        });
+        return result[0];
+    }
+
     private LuaPressure pressureOf(Iterable<ImVar> variables) {
         LuaPressure result = new LuaPressure();
         for (ImVar variable : variables) {
@@ -646,7 +661,15 @@ private LuaPressure pressureAt(Element element) {
                         java.util.Set<ImVar> active = Collections.newSetFromMap(new IdentityHashMap<>());
                         active.addAll(live.toJavaSet());
                         collectReadLocals(statement, active);
-                        return pressureOf(active);
+                        LuaPressure pressure = pressureOf(active);
+                        int stagedResults = statementExpressionResultSlots(statement);
+                        if (stagedResults > 0) {
+                            // Flattening stages each already-inlined sibling result until the
+                            // surrounding expression consumes it. The pre-inline liveness map
+                            // cannot contain those future backend temporaries yet.
+                            pressure.add("statement-expression-results", stagedResults);
+                        }
+                        return pressure;
                     }
                 }
                 current = current.getParent();

de.peeeq.wurstscript/src/test/java/tests/wurstscript/tests/OptimizerTests.java

@@ -1686,6 +1686,71 @@ public void luaArithmeticHelperRetryReusesSequentialSlots() {
             "sequential helper sites should share the same peak allocation slots");
     }
 
+    @Test
+    public void luaArithmeticHelperRetryBudgetsOverlappingArgumentResults() {
+        WurstModel model = Ast.WurstModel();
+        ImTranslator translator = new ImTranslator(model, false,
+            new RunArgs().with("-lua", "-localOptimizations"));
+        ImProg prog = translator.getImProg();
+        ImVar helperA = JassIm.ImVar(model, TypesHelper.imInt(), "a", false);
+        ImVar helperB = JassIm.ImVar(model, TypesHelper.imInt(), "b", false);
+        ImFunction helper = JassIm.ImFunction(model, "__wurst_modInt", JassIm.ImTypeVars(),
+            JassIm.ImVars(helperA, helperB), TypesHelper.imInt(), JassIm.ImVars(),
+            JassIm.ImStmts(JassIm.ImReturn(model, JassIm.ImVarAccess(helperA))),
+            Collections.emptyList());
+        translator.luaModIntFunc = helper;
+
+        ImVars callerParameters = JassIm.ImVars();
+        for (int i = 0; i < 187; i++) {
+            callerParameters.add(JassIm.ImVar(model, TypesHelper.imInt(), "p" + i, false));
+        }
+        ImVars fiveParameters = JassIm.ImVars();
+        ImExprs overlappingArguments = JassIm.ImExprs();
+        for (int i = 0; i < 5; i++) {
+            fiveParameters.add(JassIm.ImVar(model, TypesHelper.imInt(), "arg" + i, false));
+            overlappingArguments.add(JassIm.ImFunctionCall(model, helper, JassIm.ImTypeArguments(),
+                JassIm.ImExprs(JassIm.ImVarAccess(callerParameters.get(i)), JassIm.ImIntVal(3)),
+                false, de.peeeq.wurstscript.translation.imtranslation.CallType.NORMAL));
+        }
+        ImFunction takesFive = JassIm.ImFunction(model, "takesFive", JassIm.ImTypeVars(), fiveParameters,
+            JassIm.ImVoid(), JassIm.ImVars(), JassIm.ImStmts(), Collections.emptyList());
+        ImVars keepAliveParameters = JassIm.ImVars();
+        ImExprs keepAliveArguments = JassIm.ImExprs();
+        for (int i = 0; i < callerParameters.size(); i++) {
+            keepAliveParameters.add(JassIm.ImVar(model, TypesHelper.imInt(), "value" + i, false));
+            keepAliveArguments.add(JassIm.ImVarAccess(callerParameters.get(i)));
+        }
+        ImFunction keepAlive = JassIm.ImFunction(model, "keepAlive", JassIm.ImTypeVars(),
+            keepAliveParameters, JassIm.ImVoid(), JassIm.ImVars(), JassIm.ImStmts(),
+            Collections.emptyList());
+        ImFunction caller = JassIm.ImFunction(model, "caller", JassIm.ImTypeVars(), callerParameters,
+            JassIm.ImVoid(), JassIm.ImVars(), JassIm.ImStmts(
+                JassIm.ImFunctionCall(model, takesFive, JassIm.ImTypeArguments(), overlappingArguments,
+                    false, de.peeeq.wurstscript.translation.imtranslation.CallType.NORMAL),
+                JassIm.ImFunctionCall(model, keepAlive, JassIm.ImTypeArguments(), keepAliveArguments,
+                    false, de.peeeq.wurstscript.translation.imtranslation.CallType.NORMAL)),
+            Collections.emptyList());
+        prog.getFunctions().add(helper);
+        prog.getFunctions().add(takesFive);
+        prog.getFunctions().add(keepAlive);
+        prog.getFunctions().add(caller);
+
+        int changed = new ImInliner(translator).inlineLuaDivModHelpersWithinLocalBudget();
+        assertTrue(changed < 5,
+            "overlapping argument results must stop helper inlining at the register budget");
+        int[] remaining = {0};
+        caller.getBody().accept(new ImStmts.DefaultVisitor() {
+            @Override
+            public void visit(ImFunctionCall call) {
+                super.visit(call);
+                if (call.getFunc() == helper) {
+                    remaining[0]++;
+                }
+            }
+        });
+        assertTrue(remaining[0] > 0, "some overlapping helper calls must remain after the budget is reached");
+    }
+
     @Test
     public void luaArithmeticHelperRetryPreservesLocalPlayerAllocationClasses() {
         WurstModel model = Ast.WurstModel();