Skip to content

Commit e2dde9b

Browse files
FrottyFrotty
committed
Give vararg calls a fixed-arity copy on Lua
On Lua a vararg function kept its `...` parameter and every call packed the arguments into a table, and the inliner refused vararg functions outright. `max`, `min` and `ArrayList.add` are vararg, so every relink in the spatial index allocated four tables and every element added to a list allocated one. Jass has always eliminated varargs before inlining by generating one copy per call arity. The same pass now runs on Lua at the same position, after stack traces and before lowering, with three differences the target needs. Method calls are handled as well as function calls. Classes still exist when this runs on Lua, so `list.add(x)` is an ImMethodCall; the backend already turns a call with exactly one possible implementation into a direct call, and the eliminator does the same for vararg methods. Without it ArrayList.add would never have been specialised. There is no Jass parameter cap. Instead a call with more than 64 vararg arguments keeps the original, which is always still present on Lua. Originals are kept. A vararg function may still be reached through a polymorphic method call, a function reference, or a call above the bound; unreferenced originals go with garbage removal. Measured on the stdlib probe with release flags: table.pack is gone from the output, max and min are max_2 and min_2, and ArrayList.add no longer exists as a function at all because the one-element copy inlines at every call site to a capacity check, one store and one increment.
1 parent d1bf067 commit e2dde9b

3 files changed

Lines changed: 249 additions & 11 deletions

File tree

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -897,6 +897,14 @@ public LuaCompilationUnit transformProgToLua() {
897897
timeTaker.endPhase();
898898
}
899899
}
900+
// Same position as on Jass: after stack traces, before lowering and inlining. Calls with a
901+
// static argument count go to fixed-arity copies, so the emitted Lua packs no table and the
902+
// copies can inline; originals stay for dispatch, function references and calls above the bound.
903+
beginPhase(4, "eliminate varargs");
904+
new VarargEliminator(imProg, true).run();
905+
imTranslator.assertProperties();
906+
timeTaker.endPhase();
907+
900908
ImTranslator imTranslator2 = getImTranslator();
901909
ImOptimizer optimizer = new ImOptimizer(timeTaker, imTranslator2);
902910

‎de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/translation/imtranslation/VarargEliminator.java‎

Lines changed: 105 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -21,30 +21,119 @@
2121
public class VarargEliminator {
2222

2323
private static final int JASS_MAX_PARAMETERS = ImHelper.JASS_MAX_PARAMETERS;
24+
/**
25+
* Largest vararg arity given a fixed-arity copy on Lua. Lua caps a function at 200 locals
26+
* including parameters; this leaves room for the body's own locals. A call above it keeps the
27+
* original `...` function, which is always still present on that target.
28+
*/
29+
public static final int LUA_MAX_SPECIALISED_VARARG_ARITY = 64;
2430
private final ImProg prog;
31+
/**
32+
* On Lua classes are still present when this runs, so a vararg function can also be reached
33+
* through a method dispatch or a function reference. Originals are therefore kept, only direct
34+
* calls are redirected, and unreferenced originals are left to garbage removal.
35+
*/
36+
private final boolean luaTarget;
2537
// original + number of args --> new function
2638
private final Table<ImFunction, Integer, ImFunction> varargFuncs = HashBasedTable.create();
2739

2840
public VarargEliminator(ImProg prog) {
41+
this(prog, false);
42+
}
43+
44+
public VarargEliminator(ImProg prog, boolean luaTarget) {
2945
this.prog = prog;
46+
this.luaTarget = luaTarget;
3047
}
3148

3249
public void run() {
3350
// create new vararg functions
3451
for (ImFunctionCall c : collectVarargCalls()) {
35-
if (c.getFunc().hasFlag(IS_VARARG)) {
52+
if (c.getFunc().hasFlag(IS_VARARG) && shouldSpecialise(c)) {
3653
generateVarargFunc(c);
3754
}
3855
}
56+
if (luaTarget) {
57+
// The Lua backend already turns a method call with exactly one possible implementation
58+
// into a direct call of that implementation. Doing the same here for vararg methods is
59+
// what lets ArrayList.add and friends get a fixed-arity copy at all: on this target the
60+
// call is still an ImMethodCall when varargs are eliminated.
61+
for (ImMethodCall c : collectMonomorphicVarargMethodCalls()) {
62+
ImFunction implementation = c.getMethod().getImplementation();
63+
List<ImExpr> arguments = receiverAndArguments(c);
64+
if (shouldSpecialise(implementation, arguments.size())) {
65+
generateVarargFunc(implementation, arguments, c);
66+
}
67+
}
68+
}
3969

40-
// remove original vararg functions:
41-
prog.getFunctions().removeIf(f -> f.hasFlag(IS_VARARG));
70+
if (!luaTarget) {
71+
// remove original vararg functions:
72+
prog.getFunctions().removeIf(f -> f.hasFlag(IS_VARARG));
73+
}
4274

4375
// rewrite calls to use new functions:
4476
// (need to collect vararg calls again, because first phase can create copies of calls)
4577
for (ImFunctionCall call : collectVarargCalls()) {
46-
redirectCall(call, varargFuncs.get(call.getFunc(), call.getArguments().size()));
78+
ImFunction newFunc = varargFuncs.get(call.getFunc(), call.getArguments().size());
79+
if (newFunc != null) {
80+
redirectCall(call, newFunc);
81+
}
82+
}
83+
if (luaTarget) {
84+
for (ImMethodCall call : collectMonomorphicVarargMethodCalls()) {
85+
ImFunction implementation = call.getMethod().getImplementation();
86+
ImFunction newFunc = varargFuncs.get(implementation, 1 + call.getArguments().size());
87+
if (newFunc != null) {
88+
redirectMethodCall(call, newFunc);
89+
}
90+
}
91+
}
92+
}
93+
94+
/** A method call which can only ever reach one implementation, and that implementation is vararg. */
95+
private Collection<ImMethodCall> collectMonomorphicVarargMethodCalls() {
96+
final Collection<ImMethodCall> calls = new ArrayList<>();
97+
prog.accept(new ImProg.DefaultVisitor() {
98+
@Override
99+
public void visit(ImMethodCall c) {
100+
super.visit(c);
101+
ImMethod method = c.getMethod();
102+
if (method != null && !method.getIsAbstract() && method.getImplementation() != null
103+
&& method.getSubMethods().isEmpty() && method.getImplementation().hasFlag(IS_VARARG)) {
104+
calls.add(c);
105+
}
106+
}
107+
});
108+
return calls;
109+
}
110+
111+
/** The implementation's argument list: the receiver is its first parameter. */
112+
private static List<ImExpr> receiverAndArguments(ImMethodCall call) {
113+
List<ImExpr> arguments = new ArrayList<>(1 + call.getArguments().size());
114+
arguments.add(call.getReceiver());
115+
arguments.addAll(call.getArguments());
116+
return arguments;
117+
}
118+
119+
private void redirectMethodCall(ImMethodCall call, ImFunction newFunc) {
120+
ImExprs args = JassIm.ImExprs(call.getReceiver().copy());
121+
args.addAll(call.getArguments().removeAll());
122+
call.replaceBy(JassIm.ImFunctionCall(call.getTrace(), newFunc, JassIm.ImTypeArguments(), args,
123+
call.getTuplesEliminated(), CallType.NORMAL));
124+
}
125+
126+
/** Whether a call gets a fixed-arity copy. Always on Jass; on Lua only up to the arity bound. */
127+
private boolean shouldSpecialise(ImFunctionCall call) {
128+
return shouldSpecialise(call.getFunc(), call.getArguments().size());
129+
}
130+
131+
private boolean shouldSpecialise(ImFunction func, int totalArguments) {
132+
if (!luaTarget) {
133+
return true;
47134
}
135+
int varargCount = 1 + totalArguments - func.getParameters().size();
136+
return varargCount <= LUA_MAX_SPECIALISED_VARARG_ARITY;
48137
}
49138

50139
@NotNull
@@ -70,13 +159,17 @@ public void visit(ImFunctionCall c) {
70159
* for the function call.
71160
*/
72161
private void generateVarargFunc(ImFunctionCall sourceCall) {
73-
ImFunction func = sourceCall.getFunc();
74-
int numberOfParams = sourceCall.getArguments().size();
75-
int jassParameterCount = sourceCall.getArguments().stream()
162+
generateVarargFunc(sourceCall.getFunc(), sourceCall.getArguments(), sourceCall);
163+
}
164+
165+
/** {@code arguments} are in the callee's parameter order, so for a method they start with the receiver. */
166+
private void generateVarargFunc(ImFunction func, List<ImExpr> arguments, Element trace) {
167+
int numberOfParams = arguments.size();
168+
int jassParameterCount = arguments.stream()
76169
.mapToInt(argument -> ImHelper.flattenedJassArity(argument.attrTyp()))
77170
.sum();
78-
if (jassParameterCount > JASS_MAX_PARAMETERS) {
79-
throw new CompileError(sourceCall, "Vararg call would generate " + jassParameterCount
171+
if (!luaTarget && jassParameterCount > JASS_MAX_PARAMETERS) {
172+
throw new CompileError(trace, "Vararg call would generate " + jassParameterCount
80173
+ " Jass parameters; the maximum is " + JASS_MAX_PARAMETERS
81174
+ ". Use multiple calls (for example with the cascade operator) or pass a collection instead.");
82175
}
@@ -131,7 +224,9 @@ public void visit(ImVarargLoop imLoop) {
131224
params.addAll(list);
132225

133226
// generate function for this new call
134-
generateVarargFunc(call);
227+
if (shouldSpecialise(call)) {
228+
generateVarargFunc(call);
229+
}
135230
}
136231

137232

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

Lines changed: 136 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -271,7 +271,7 @@ public void optimizedTupleVarargLoopUsesAttachedScalarLocals() {
271271
" let bag = new Bag<handles>()",
272272
" bag.add(handles(makeFrame(), makeFrame()))"
273273
);
274-
assertTrue(compiled.contains("table.pack(...)"));
274+
assertFalse("a static-arity vararg call must not pack a table on Lua", compiled.contains("table.pack(...)"));
275275
assertFalse(compiled.contains("tupleCopy"));
276276
}
277277

@@ -1988,6 +1988,141 @@ public void localPlayerEffectfulBooleanOperandSurvivesOptimization() {
19881988
compiled.indexOf("localProbe()", definitionOrCall + 1) >= 0);
19891989
}
19901990

1991+
/**
1992+
* On Lua a vararg function used to keep its {@code ...} parameter and pack it into a table on
1993+
* every call, and the inliner refused it. With a static argument count at the call site the
1994+
* call is redirected to a fixed-arity copy, as Jass has always done, so the pack is gone and
1995+
* the copy inlines like any other small function.
1996+
*/
1997+
@Test
1998+
public void staticArityVarargCallsAreFixedArityOnLua() {
1999+
String compiled = compileOptimizedLua(
2000+
"staticArityVarargCallsAreFixedArityOnLua",
2001+
"package Test",
2002+
"native consume(int i)",
2003+
"int array values",
2004+
"function biggest(vararg int xs) returns int",
2005+
" var best = -2147483648",
2006+
" for x in xs",
2007+
" if x > best",
2008+
" best = x",
2009+
" return best",
2010+
"@noinline function query(int a, int b, int c)",
2011+
" var i = 0",
2012+
" while i < 16",
2013+
" consume(biggest(a, values[i]))",
2014+
" consume(biggest(a, b, c))",
2015+
" i++",
2016+
"init",
2017+
" query(1, 2, 3)"
2018+
);
2019+
assertFalse("no vararg call site may pack a table:\n" + compiled, compiled.contains("table.pack"));
2020+
assertFalse("the vararg original must not survive with a ... parameter:\n" + compiled,
2021+
compiled.contains("function biggest(...)"));
2022+
assertFunctionBodyContains(compiled, "query", "biggest(", false);
2023+
}
2024+
2025+
@Test
2026+
public void fixedArityVarargLoweringKeepsSemanticsOnLua() throws IOException {
2027+
test().testLua(true).executeProg().lines(
2028+
"package Test",
2029+
"native testSuccess()",
2030+
"tuple pair(int x, int y)",
2031+
"function sum(vararg int xs) returns int",
2032+
" var total = 0",
2033+
" for x in xs",
2034+
" total += x",
2035+
" return total",
2036+
"function count(vararg int xs) returns int",
2037+
" var n = 0",
2038+
" for x in xs",
2039+
" n++",
2040+
" return n",
2041+
"function firstOr(vararg int xs) returns int",
2042+
" for x in xs",
2043+
" return x",
2044+
" return -1",
2045+
"function pairs(vararg pair ps) returns int",
2046+
" var result = 0",
2047+
" for p in ps",
2048+
" result = result * 100 + p.x * 10 + p.y",
2049+
" return result",
2050+
"class Bag",
2051+
" int total = 0",
2052+
" function add(vararg int xs)",
2053+
" for x in xs",
2054+
" total += x",
2055+
"init",
2056+
" let bag = new Bag()",
2057+
" bag.add(1)",
2058+
" bag.add(2, 3)",
2059+
" bag.add()",
2060+
" if sum() == 0 and sum(5) == 5 and sum(1, 2, 3, 4) == 10",
2061+
" and count() == 0 and count(9, 9, 9) == 3",
2062+
" and firstOr() == -1 and firstOr(4, 5) == 4",
2063+
" and pairs(pair(1, 2), pair(3, 4)) == 1234",
2064+
" and bag.total == 6",
2065+
" testSuccess()"
2066+
);
2067+
String compiled = compiledLua("fixedArityVarargLoweringKeepsSemanticsOnLua");
2068+
assertFalse("every call above has a static arity, so nothing may pack:\n" + compiled,
2069+
compiled.contains("table.pack"));
2070+
}
2071+
2072+
@Test
2073+
public void varargCallAboveTheLuaArityBoundKeepsThePackedPath() throws IOException {
2074+
StringBuilder args = new StringBuilder();
2075+
int n = 150;
2076+
for (int i = 1; i <= n; i++) {
2077+
if (i > 1) {
2078+
args.append(", ");
2079+
}
2080+
args.append(i);
2081+
}
2082+
test().testLua(true).executeProg().lines(
2083+
"package Test",
2084+
"native testSuccess()",
2085+
"function sum(vararg int xs) returns int",
2086+
" var total = 0",
2087+
" for x in xs",
2088+
" total += x",
2089+
" return total",
2090+
"init",
2091+
" if sum(" + args + ") == " + (n * (n + 1) / 2) + " and sum(1, 2) == 3",
2092+
" testSuccess()"
2093+
);
2094+
String compiled = compiledLua("varargCallAboveTheLuaArityBoundKeepsThePackedPath");
2095+
assertTrue("a call above the bound keeps the vararg original:\n" + compiled,
2096+
compiled.contains("table.pack"));
2097+
}
2098+
2099+
@Test
2100+
public void virtuallyDispatchedVarargMethodKeepsThePackedPath() throws IOException {
2101+
test().testLua(true).executeProg().lines(
2102+
"package Test",
2103+
"native testSuccess()",
2104+
"interface Summer",
2105+
" function sum(vararg int xs) returns int",
2106+
"class Plain implements Summer",
2107+
" override function sum(vararg int xs) returns int",
2108+
" var total = 0",
2109+
" for x in xs",
2110+
" total += x",
2111+
" return total",
2112+
"class Doubling implements Summer",
2113+
" override function sum(vararg int xs) returns int",
2114+
" var total = 0",
2115+
" for x in xs",
2116+
" total += 2 * x",
2117+
" return total",
2118+
"init",
2119+
" Summer a = new Plain()",
2120+
" Summer b = new Doubling()",
2121+
" if a.sum(1, 2, 3) == 6 and b.sum(1, 2, 3) == 12",
2122+
" testSuccess()"
2123+
);
2124+
}
2125+
19912126
@Test
19922127
public void localPlayerTaintFlowsThroughVarargLoopValues() {
19932128
String compiled = compileOptimizedLua(

0 commit comments

Comments
 (0)