Bound Lua inlining without local allocation

Author: Frotty

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstio/WurstCompilerJassImpl.java

@@ -946,10 +946,10 @@ public LuaCompilationUnit transformProgToLua() {
             timeTaker.endPhase();
         }
 
-        if (runArgs.isInline()) {
+        if (runArgs.isInline() && runArgs.isLocalOptimizations()) {
             beginPhase(10, "inline Lua arithmetic helpers within allocated local budget");
             int arithmeticHelpersInlined = optimizer.inlineLuaDivModHelpersWithinLocalBudget();
-            if (arithmeticHelpersInlined > 0 && runArgs.isLocalOptimizations()) {
+            if (arithmeticHelpersInlined > 0) {
                 optimizer.localOptimizations();
             }
             timeTaker.endPhase();

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

@@ -555,6 +555,7 @@ private final class LuaRegisterBudget {
         private Map<ImStmt, io.vavr.collection.Set<ImVar>> liveness;
         private LuaPressure peakPressure;
         private int backendLocals;
+        private int declarationsWithoutAllocation;
 
         private LuaRegisterBudget(ImFunction function) {
             this.function = function;
@@ -566,16 +567,31 @@ private LuaRegisterBudget(ImFunction function) {
             }
             peakPressure = cachedPressure.copy();
             backendLocals = backendGeneratedLuaLocals(function);
+            declarationsWithoutAllocation = function.getParameters().size() + function.getLocals().size()
+                + backendLocals;
         }
 
         private boolean fits(ImFunctionCall call, ImFunction callee) {
+            if (!translator.getRunArgs().isLocalOptimizations()) {
+                return declarationsWithoutAllocation + declarationsAddedByInline(callee)
+                    <= LUA_INLINE_REGISTER_BUDGET;
+            }
             return projectedPressure(call, callee) <= LUA_INLINE_REGISTER_BUDGET
                 - backendLocals - backendGeneratedLuaLocals(callee);
         }
 
+        private int declarationsAddedByInline(ImFunction callee) {
+            int controlLocals = maxOneReturn(callee)
+                ? 0
+                : 1 + (callee.getReturnType() instanceof ImVoid ? 0 : 1);
+            return callee.getParameters().size() + callee.getLocals().size() + controlLocals
+                + backendGeneratedLuaLocals(callee);
+        }
+
         private void recordInline(ImFunctionCall call, ImFunction callee) {
             peakPressure.keepMaximums(pressureDuringInline(call, callee));
             backendLocals += backendGeneratedLuaLocals(callee);
+            declarationsWithoutAllocation += declarationsAddedByInline(callee);
             // Callers are processed after their callees. Publish the expanded pressure so a
             // later caller budgets the body it will actually copy, not the pre-inline callee.
             luaRegisterPressure.put(function, peakPressure.copy());

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

@@ -2208,6 +2208,33 @@ public void luaInlinerKeepsCallWhenLiveValuesWouldExceedRegisterBudget() {
             callerBody.contains("__wurst_locals"));
     }
 
+    @Test
+    public void luaInliningWithoutLocalAllocationUsesDeclarationBudget() {
+        List<String> lines = new ArrayList<>();
+        lines.add("package Test");
+        lines.add("native takesInt(int i)");
+        lines.add("@noinline function caller(int value)");
+        for (int i = 0; i < 70; i++) {
+            lines.add("    takesInt((value + " + i + ") mod 3)");
+        }
+        lines.add("init");
+        lines.add("    caller(7)");
+
+        String compiled = compileLuaWithCUs(
+            "LuaTranslationTests_luaInliningWithoutLocalAllocationUsesDeclarationBudget",
+            false, Collections.emptyList(), new RunArgs().with("-lua", "-inline"),
+            lines.toArray(new String[0]));
+        int callerStart = compiled.indexOf("function caller(");
+        assertTrue("caller function not found", callerStart >= 0);
+        int callerEnd = compiled.indexOf("\nend", callerStart);
+        assertTrue("caller function end not found", callerEnd > callerStart);
+        String body = compiled.substring(callerStart, callerEnd);
+        assertFalse("inlining without allocation must not cross into whole-function spill mode:\n" + body,
+            body.contains("__wurst_locals"));
+        assertTrue("the exact declaration budget must retain residual helper calls near the limit:\n" + body,
+            body.contains("__wurst_modInt("));
+    }
+
     @Test
     public void luaInlinerReusesRegistersAcrossSequentialInlineSites() {
         List<String> lines = new ArrayList<>();
@@ -2354,7 +2381,8 @@ public void luaInlinerDoesNotTreatSequentialVarargLoopTempsAsConcurrent() {
 
         String compiled = compileLuaWithCUs(
             "LuaTranslationTests_luaInlinerDoesNotTreatSequentialVarargLoopTempsAsConcurrent",
-            false, Collections.emptyList(), new RunArgs().with("-lua", "-inline"),
+            false, Collections.emptyList(),
+            new RunArgs().with("-lua", "-inline", "-localOptimizations"),
             lines.toArray(new String[0]));
         assertFalse("sequential loop temporaries must not consume concurrent register budget:\n" + compiled,
             compiled.contains("small("));

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

@@ -1600,7 +1600,8 @@ public void repeatedLocalOptimizationStartsANewIteration() {
     @Test
     public void luaArithmeticHelperRetryRespectsFunctionLocalBudget() {
         WurstModel model = Ast.WurstModel();
-        ImTranslator translator = new ImTranslator(model, false, new RunArgs().with("-lua"));
+        ImTranslator translator = new ImTranslator(model, false,
+            new RunArgs().with("-lua", "-localOptimizations"));
         ImProg prog = translator.getImProg();
 
         ImVar helperA = JassIm.ImVar(model, TypesHelper.imInt(), "a", false);
@@ -1654,7 +1655,8 @@ public void luaArithmeticHelperRetryRespectsFunctionLocalBudget() {
     @Test
     public void luaArithmeticHelperRetryReusesSequentialSlots() {
         WurstModel model = Ast.WurstModel();
-        ImTranslator translator = new ImTranslator(model, false, new RunArgs().with("-lua"));
+        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);