Optimize generated Lua table access in hot loops (#1297)

* Optimize generated Lua table access in hot loops

* Update Lua inliner storage assertion

Author: Frotty

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/translation/lua/translation/LuaTranslator.java

@@ -24,6 +24,8 @@
 
 public class LuaTranslator {
     private static final int LUA_LOCALS_LIMIT = 200;
+    private static final int LUA_STORAGE_ALIAS_LOCAL_BUDGET = 190;
+    private static final int LUA_STORAGE_ALIAS_LIMIT = 16;
     private static final List<String> HASHTABLE_HANDLE_SAVE_NAMES = Arrays.asList(
         "SavePlayerHandle", "SaveWidgetHandle", "SaveDestructableHandle", "SaveItemHandle", "SaveUnitHandle",
         "SaveAbilityHandle", "SaveTimerHandle", "SaveTriggerHandle", "SaveTriggerConditionHandle",
@@ -80,6 +82,8 @@ public class LuaTranslator {
     private final LuaStatements deferredMainInit = LuaAst.LuaStatements();
     private final Map<String, Integer> uniqueNameCounters = new HashMap<>();
     private final Set<String> usedNames = LuaReservedNames.all();
+    private final Set<LuaVariable> localizableStorageTables = Collections.newSetFromMap(new IdentityHashMap<>());
+    private LuaFunction bootstrapFunction;
     private final Set<String> emittedDispatchSlots = new HashSet<>();
     private final Map<ImClass, Set<String>> emittedDispatchSlotsByClass = new IdentityHashMap<>();
     private final Map<ImClass, Map<String, Set<DispatchGroupIdentity>>> emittedDispatchSlotGroupsByClass = new IdentityHashMap<>();
@@ -222,8 +226,10 @@ public LuaVariable initFor(ImClass a) {
     GetAForB<ImVar, LuaVariable> luaFieldStorage = new GetAForB<ImVar, LuaVariable>() {
         @Override
         public LuaVariable initFor(ImVar field) {
-            return LuaAst.LuaVariable(uniqueName(field.getName() + "_storage"),
+            LuaVariable storage = LuaAst.LuaVariable(uniqueName(field.getName() + "_storage"),
                 LuaAst.LuaTableConstructor(LuaAst.LuaTableFields()));
+            localizableStorageTables.add(storage);
+            return storage;
         }
     };
 
@@ -382,6 +388,7 @@ public LuaCompilationUnit translate() {
 
         createBootstrapFunction();
         cleanStatements();
+        localizeHotStorageTables();
         enforceLuaLocalLimits();
 
         return luaModel;
@@ -509,6 +516,7 @@ private void createBootstrapFunction() {
         LuaVariable doneFlag = LuaAst.LuaVariable(uniqueName("__wurst_bootstrap_done"), LuaAst.LuaExprNull());
         luaModel.add(doneFlag);
         LuaFunction boot = LuaAst.LuaFunction(uniqueName("__wurst_init_bootstrap"), LuaAst.LuaParams(), LuaAst.LuaStatements());
+        bootstrapFunction = boot;
         boot.getBody().add(LuaAst.LuaLiteral("if " + doneFlag.getName() + " then return end"));
         boot.getBody().add(LuaAst.LuaAssignment(LuaAst.LuaExprVarAccess(doneFlag), LuaAst.LuaExprBoolVal(true)));
         List<LuaStatement> stmts = new ArrayList<>();
@@ -955,6 +963,107 @@ public void visit(LuaMethod m) {
         });
     }
 
+    /**
+     * Cache compiler-owned array and field-storage tables in function locals when they are indexed
+     * from a loop. This removes a global lookup from every dynamic iteration without changing cold
+     * paths or making assumptions about user-authored Lua tables. The cap and headroom avoid turning
+     * the optimization into register pressure or immediately triggering the locals-table fallback.
+     */
+    private void localizeHotStorageTables() {
+        luaModel.accept(new LuaModel.DefaultVisitor() {
+            @Override
+            public void visit(LuaFunction f) {
+                super.visit(f);
+                localizeHotStorageTables(f.getParams(), f.getBody());
+            }
+
+            @Override
+            public void visit(LuaMethod m) {
+                super.visit(m);
+                localizeHotStorageTables(m.getParams(), m.getBody());
+            }
+        });
+    }
+
+    private void localizeHotStorageTables(LuaParams params, LuaStatements body) {
+        Map<LuaVariable, Integer> loopAccessCounts = new IdentityHashMap<>();
+        collectLoopStorageAccesses(body, false, loopAccessCounts);
+        if (loopAccessCounts.isEmpty()) {
+            return;
+        }
+
+        int existingLocals = params.size() + collectFunctionScopeLocals(body).size();
+        int aliasCount = Math.min(LUA_STORAGE_ALIAS_LIMIT,
+            LUA_STORAGE_ALIAS_LOCAL_BUDGET - existingLocals);
+        if (aliasCount <= 0) {
+            return;
+        }
+
+        List<LuaVariable> selected = new ArrayList<>(loopAccessCounts.keySet());
+        selected.sort(Comparator.<LuaVariable>comparingInt(loopAccessCounts::get).reversed()
+            .thenComparing(LuaVariable::getName));
+        if (selected.size() > aliasCount) {
+            selected = new ArrayList<>(selected.subList(0, aliasCount));
+        }
+
+        Map<LuaVariable, LuaVariable> aliases = new IdentityHashMap<>();
+        for (LuaVariable storage : selected) {
+            aliases.put(storage, LuaAst.LuaVariable(uniqueName(storage.getName() + "_local"),
+                LuaAst.LuaExprVarAccess(storage)));
+        }
+        rewriteStorageAccesses(body, aliases);
+
+        int insertionIndex = bootstrapCallInsertionIndex(body);
+        for (LuaVariable storage : selected) {
+            body.add(insertionIndex++, aliases.get(storage));
+        }
+    }
+
+    private void collectLoopStorageAccesses(de.peeeq.wurstscript.luaAst.Element element, boolean insideLoop,
+                                            Map<LuaVariable, Integer> counts) {
+        if (element instanceof LuaExprFunctionAbstraction
+            || element instanceof LuaFunction || element instanceof LuaMethod) {
+            return;
+        }
+        boolean childInsideLoop = insideLoop || element instanceof LuaWhile;
+        if (childInsideLoop && element instanceof LuaExprArrayAccess) {
+            LuaExpr left = ((LuaExprArrayAccess) element).getLeft();
+            if (left instanceof LuaExprVarAccess) {
+                LuaVariable storage = ((LuaExprVarAccess) left).getVar();
+                if (localizableStorageTables.contains(storage)) {
+                    counts.merge(storage, 1, Integer::sum);
+                }
+            }
+        }
+        element.forEachElement(child -> collectLoopStorageAccesses(child, childInsideLoop, counts));
+    }
+
+    private void rewriteStorageAccesses(de.peeeq.wurstscript.luaAst.Element element,
+                                        Map<LuaVariable, LuaVariable> aliases) {
+        if (element instanceof LuaExprFunctionAbstraction
+            || element instanceof LuaFunction || element instanceof LuaMethod) {
+            return;
+        }
+        if (element instanceof LuaExprArrayAccess) {
+            LuaExpr left = ((LuaExprArrayAccess) element).getLeft();
+            if (left instanceof LuaExprVarAccess) {
+                LuaVariable alias = aliases.get(((LuaExprVarAccess) left).getVar());
+                if (alias != null) {
+                    left.replaceBy(LuaAst.LuaExprVarAccess(alias));
+                }
+            }
+        }
+        element.forEachElement(child -> rewriteStorageAccesses(child, aliases));
+    }
+
+    private int bootstrapCallInsertionIndex(LuaStatements body) {
+        if (bootstrapFunction != null && !body.isEmpty() && body.get(0) instanceof LuaExprFunctionCall
+            && ((LuaExprFunctionCall) body.get(0)).getFunc() == bootstrapFunction) {
+            return 1;
+        }
+        return 0;
+    }
+
     private void spillLocalsIntoTableIfNeeded(String functionName, LuaParams params, LuaStatements body) {
         List<LuaVariable> scopeLocals = collectFunctionScopeLocals(body);
         int localCount = params.size() + scopeLocals.size();
@@ -2037,6 +2146,9 @@ private void translateGlobal(ImVar v) {
             return;
         }
         LuaVariable lv = luaVar.getFor(v);
+        if (v.getType() instanceof ImArrayType || v.getType() instanceof ImArrayTypeMulti) {
+            localizableStorageTables.add(lv);
+        }
         lv.setInitialValue(LuaAst.LuaExprNull());
         luaModel.add(lv);
         deferMainInit(LuaAst.LuaAssignment(LuaAst.LuaExprVarAccess(lv), defaultValue(v.getType())));

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

@@ -2048,8 +2048,12 @@ public void tinyMonomorphicMethodsInlineOnLua() {
         String body = topLevelFunctionBodyWithPrefix(compiled, "hotLoop");
         assertFalse("a tiny method with exactly one implementation must inline in optimized Lua:\n" + body,
             body.contains("Accumulator_add("));
-        assertTrue("the inlined method must retain its field read:\n" + body,
-            body.contains("Accumulator_offset_storage["));
+        java.util.regex.Matcher offsetAlias = java.util.regex.Pattern
+            .compile("local (\\w+) = Accumulator_offset_storage").matcher(body);
+        assertTrue("the inlined method's field storage must be localized:\n" + body,
+            offsetAlias.find());
+        assertTrue("the inlined method must retain its field read through the local alias:\n" + body,
+            body.contains(offsetAlias.group(1) + "["));
         String dynamicBody = topLevelFunctionBodyWithPrefix(compiled, "dynamicCall");
         assertTrue("a genuinely virtual method call must retain dispatch:\n" + dynamicBody,
             dynamicBody.contains("dispatch_"));
@@ -2165,16 +2169,31 @@ public void optimizedUnitSpatialIndexInnerLoopUsesRawLuaOperations() {
         // Register-pressure-aware inlining may keep the range helper as the hot-loop owner instead
         // of folding it into query. Inspect whichever function actually retains the loop.
         String body = topLevelFunctionBodyWithPrefix(compiled, "query");
-        if (!body.contains("UnitSpatialIndex_nextInCell[")) {
+        if (!body.contains("= UnitSpatialIndex_nextInCell")) {
             assertTrue("query must call the retained range helper:\n" + body,
                 body.contains("addRangeMatches("));
             body = topLevelFunctionBodyWithPrefix(compiled, "addRangeMatches");
         }
-        assertTrue("spatial-index hot loop must read the next-link array directly:\n" + body,
+        java.util.regex.Matcher nextAlias = java.util.regex.Pattern
+            .compile("local (\\w+) = UnitSpatialIndex_nextInCell").matcher(body);
+        java.util.regex.Matcher xAlias = java.util.regex.Pattern
+            .compile("local (\\w+) = UnitSpatialIndex_lastX").matcher(body);
+        java.util.regex.Matcher yAlias = java.util.regex.Pattern
+            .compile("local (\\w+) = UnitSpatialIndex_lastY").matcher(body);
+        assertTrue("spatial-index hot loop must localize the next-link array:\n" + body, nextAlias.find());
+        assertTrue("spatial-index hot loop must localize cached X:\n" + body, xAlias.find());
+        assertTrue("spatial-index hot loop must localize cached Y:\n" + body, yAlias.find());
+        assertTrue("spatial-index hot loop must index the localized next-link array:\n" + body,
+            body.contains(nextAlias.group(1) + "["));
+        assertTrue("spatial-index hot loop must index localized cached X:\n" + body,
+            body.contains(xAlias.group(1) + "["));
+        assertTrue("spatial-index hot loop must index localized cached Y:\n" + body,
+            body.contains(yAlias.group(1) + "["));
+        assertFalse("spatial-index loop must not retain global next-link lookups:\n" + body,
             body.contains("UnitSpatialIndex_nextInCell["));
-        assertTrue("spatial-index hot loop must read cached X directly:\n" + body,
+        assertFalse("spatial-index loop must not retain global cached-X lookups:\n" + body,
             body.contains("UnitSpatialIndex_lastX["));
-        assertTrue("spatial-index hot loop must read cached Y directly:\n" + body,
+        assertFalse("spatial-index loop must not retain global cached-Y lookups:\n" + body,
             body.contains("UnitSpatialIndex_lastY["));
         assertFalse("typed array reads must not retain assurance calls:\n" + body,
             body.contains("__wurst_ensure"));
@@ -2184,6 +2203,114 @@ public void optimizedUnitSpatialIndexInnerLoopUsesRawLuaOperations() {
             body.contains("__wurst_intDiv("));
     }
 
+    @Test
+    public void hotLoopStorageTablesAreLocalizedOnLua() {
+        String compiled = compileOptimizedLua(
+            "hotLoopStorageTablesAreLocalizedOnLua",
+            "package Test",
+            "int array values",
+            "class Counter",
+            "    int total",
+            "@noinline function hotLoop(Counter counter, int limit)",
+            "    var i = 0",
+            "    while i < limit",
+            "        values[i] = values[i] + counter.total",
+            "        counter.total = values[i + 1] + counter.total",
+            "        i++",
+            "@noinline function coldPath(Counter counter) returns int",
+            "    return values[0] + counter.total",
+            "init",
+            "    let counter = new Counter()",
+            "    hotLoop(counter, 16)",
+            "    coldPath(counter)"
+        );
+
+        String hotBody = topLevelFunctionBodyWithPrefix(compiled, "hotLoop");
+        java.util.regex.Matcher arrayAlias = java.util.regex.Pattern
+            .compile("local (\\w+) = Test_values").matcher(hotBody);
+        assertTrue("the loop's generated array table must be cached in a local:\n" + hotBody,
+            arrayAlias.find());
+        String arrayAliasName = arrayAlias.group(1);
+        assertTrue("the localized array must be used in the loop:\n" + hotBody,
+            hotBody.contains(arrayAliasName + "["));
+        assertFalse("loop indexes must not keep resolving the array through a global:\n" + hotBody,
+            hotBody.contains("Test_values["));
+
+        java.util.regex.Matcher fieldAlias = java.util.regex.Pattern
+            .compile("local (\\w+) = Counter_total_storage").matcher(hotBody);
+        assertTrue("the loop's generated field-storage table must be cached in a local:\n" + hotBody,
+            fieldAlias.find());
+        String fieldAliasName = fieldAlias.group(1);
+        assertTrue("the localized field storage must be used in the loop:\n" + hotBody,
+            hotBody.contains(fieldAliasName + "["));
+        assertFalse("loop field accesses must not keep resolving storage through a global:\n" + hotBody,
+            hotBody.contains("Counter_total_storage["));
+
+        String coldBody = topLevelFunctionBodyWithPrefix(compiled, "coldPath");
+        assertFalse("cold array access must not consume a local alias:\n" + coldBody,
+            coldBody.contains("local "));
+        assertTrue("cold array access must keep the generated global table:\n" + coldBody,
+            coldBody.contains("Test_values["));
+        assertTrue("cold field access must keep the generated global storage:\n" + coldBody,
+            coldBody.contains("Counter_total_storage["));
+    }
+
+    @Test
+    public void localizedHotStorageTablesPreserveLuaBehavior() throws IOException {
+        test().testLua(true).executeProg().lines(
+            "package Test",
+            "native testSuccess()",
+            "int array values",
+            "class Counter",
+            "    int total",
+            "function hotLoop(Counter counter, int limit)",
+            "    var i = 0",
+            "    while i < limit",
+            "        values[i] = values[i] + counter.total",
+            "        counter.total = values[i + 1] + counter.total",
+            "        i++",
+            "init",
+            "    values[0] = 1",
+            "    let counter = new Counter()",
+            "    counter.total = 2",
+            "    hotLoop(counter, 3)",
+            "    if values[0] == 3 and values[1] == 2 and counter.total == 2",
+            "        testSuccess()"
+        );
+    }
+
+    @Test
+    public void hotStorageLocalizationRespectsLuaRegisterHeadroom() {
+        List<String> lines = new ArrayList<>();
+        lines.add("package Test");
+        lines.add("native consume(int value)");
+        lines.add("int array values");
+        lines.add("@noinline function crowded(int limit)");
+        lines.add("    var sum = 0");
+        for (int i = 0; i < 187; i++) {
+            lines.add("    let v" + i + " = " + i);
+            lines.add("    sum += v" + i);
+        }
+        lines.add("    var i = 0");
+        lines.add("    while i < limit");
+        lines.add("        sum += values[i]");
+        lines.add("        i++");
+        lines.add("    consume(sum)");
+        lines.add("init");
+        lines.add("    crowded(1)");
+
+        String compiled = compileLuaWithRunArgs(
+            "hotStorageLocalizationRespectsLuaRegisterHeadroom",
+            new RunArgs().with("-lua"), lines.toArray(new String[0]));
+        String body = topLevelFunctionBodyWithPrefix(compiled, "crowded");
+        assertFalse("localization must retain headroom below Lua's hard local limit:\n" + body,
+            body.contains("= Test_values"));
+        assertTrue("a skipped alias must leave the generated array access intact:\n" + body,
+            body.contains("Test_values[i]"));
+        assertFalse("the optimization must not force the locals-table fallback:\n" + body,
+            body.contains("__wurst_locals"));
+    }
+
     /**
      * On Lua a vararg function used to keep its {@code ...} parameter and pack it into a table on
      * every call, and the inliner refused it. With a static argument count at the call site the