diff --git a/de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/translation/lua/translation/LuaTranslator.java b/de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/translation/lua/translation/LuaTranslator.java index 29188d783..eba923b3a 100644 --- a/de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/translation/lua/translation/LuaTranslator.java +++ b/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 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 uniqueNameCounters = new HashMap<>(); private final Set usedNames = LuaReservedNames.all(); + private final Set localizableStorageTables = Collections.newSetFromMap(new IdentityHashMap<>()); + private LuaFunction bootstrapFunction; private final Set emittedDispatchSlots = new HashSet<>(); private final Map> emittedDispatchSlotsByClass = new IdentityHashMap<>(); private final Map>> emittedDispatchSlotGroupsByClass = new IdentityHashMap<>(); @@ -222,8 +226,10 @@ public LuaVariable initFor(ImClass a) { GetAForB luaFieldStorage = new GetAForB() { @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 stmts = new ArrayList<>(); @@ -957,6 +965,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 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 selected = new ArrayList<>(loopAccessCounts.keySet()); + selected.sort(Comparator.comparingInt(loopAccessCounts::get).reversed() + .thenComparing(LuaVariable::getName)); + if (selected.size() > aliasCount) { + selected = new ArrayList<>(selected.subList(0, aliasCount)); + } + + Map 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 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 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 scopeLocals = collectFunctionScopeLocals(body); int localCount = params.size() + scopeLocals.size(); @@ -2039,6 +2148,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()))); diff --git a/de.peeeq.wurstscript/src/test/java/tests/wurstscript/tests/LuaBackendAuditTests.java b/de.peeeq.wurstscript/src/test/java/tests/wurstscript/tests/LuaBackendAuditTests.java index 5d579530e..5724aaa1e 100644 --- a/de.peeeq.wurstscript/src/test/java/tests/wurstscript/tests/LuaBackendAuditTests.java +++ b/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 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