Skip to content

Optimize generated Lua table access in hot loops #1297

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Frotty merged 2 commits into from
Sep 4, 2026
+246 −7
Merged

Optimize generated Lua table access in hot loops #1297

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8fe5040
Optimize generated Lua table access in hot loops
Frotty Sep 4, 2026
181109f
Update Lua inliner storage assertion
Frotty Sep 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
114 changes: 113 additions & 1 deletion ...tscript/src/main/java/de/peeeq/wurstscript/translation/lua/translation/LuaTranslator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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<>();
Expand Down Expand Up @@ -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;
}
};

Expand Down Expand Up @@ -382,6 +388,7 @@ public LuaCompilationUnit translate() {

createBootstrapFunction();
cleanStatements();
localizeHotStorageTables();
enforceLuaLocalLimits();

return luaModel;
Expand Down Expand Up @@ -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<>();
Expand Down Expand Up @@ -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<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();
Expand Down Expand Up @@ -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())));
Expand Down
139 changes: 133 additions & 6 deletions de.peeeq.wurstscript/src/test/java/tests/wurstscript/tests/LuaBackendAuditTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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_"));
Expand Down Expand Up @@ -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"));
Expand All @@ -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
Expand Down
Loading