Give vararg calls a fixed-arity copy on Lua

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

@@ -897,6 +897,14 @@ public LuaCompilationUnit transformProgToLua() {
                 timeTaker.endPhase();
             }
         }
+        // Same position as on Jass: after stack traces, before lowering and inlining. Calls with a
+        // static argument count go to fixed-arity copies, so the emitted Lua packs no table and the
+        // copies can inline; originals stay for dispatch, function references and calls above the bound.
+        beginPhase(4, "eliminate varargs");
+        new VarargEliminator(imProg, true).run();
+        imTranslator.assertProperties();
+        timeTaker.endPhase();
+
         ImTranslator imTranslator2 = getImTranslator();
         ImOptimizer optimizer = new ImOptimizer(timeTaker, imTranslator2);
 

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

@@ -21,32 +21,144 @@
 public class VarargEliminator {
 
     private static final int JASS_MAX_PARAMETERS = ImHelper.JASS_MAX_PARAMETERS;
+    /**
+     * Largest number of emitted parameters a fixed-arity copy may have on Lua, counted after tuple
+     * flattening: one four-field tuple argument is four parameters, so an arity that looks modest in
+     * source can exceed what the target accepts. Lua caps a function at 200 locals including its
+     * parameters, and the locals-table fallback cannot spill parameters, so this leaves room for the
+     * body's own locals. A call above it keeps the original `...` function, which is always still
+     * present on that target.
+     */
+    public static final int LUA_MAX_SPECIALISED_VARARG_PARAMETERS = 64;
     private final ImProg prog;
+    /**
+     * On Lua classes are still present when this runs, so a vararg function can also be reached
+     * through a method dispatch or a function reference. Originals are therefore kept, only direct
+     * calls are redirected, and unreferenced originals are left to garbage removal.
+     */
+    private final boolean luaTarget;
     // original + number of args --> new function
     private final Table<ImFunction, Integer, ImFunction> varargFuncs = HashBasedTable.create();
 
     public VarargEliminator(ImProg prog) {
+        this(prog, false);
+    }
+
+    public VarargEliminator(ImProg prog, boolean luaTarget) {
         this.prog = prog;
+        this.luaTarget = luaTarget;
     }
 
     public void run() {
-        // create new vararg functions
-        for (ImFunctionCall c : collectVarargCalls()) {
-            if (c.getFunc().hasFlag(IS_VARARG)) {
-                generateVarargFunc(c);
+        // Create new vararg functions. Repeated to a fixpoint: a generated copy can contain a call
+        // to a vararg function at an arity nothing has needed yet, which is what a recursive vararg
+        // function calling itself with a different argument count produces.
+        boolean generated = true;
+        while (generated) {
+            generated = false;
+            for (ImFunctionCall c : collectVarargCalls()) {
+                if (c.getFunc().hasFlag(IS_VARARG) && shouldSpecialise(c)
+                    && !varargFuncs.contains(c.getFunc(), c.getArguments().size())) {
+                    generateVarargFunc(c);
+                    generated = true;
+                }
+            }
+            if (luaTarget) {
+                // The Lua backend already turns a method call with exactly one possible
+                // implementation into a direct call of that implementation. Doing the same here for
+                // vararg methods is what lets ArrayList.add and friends get a fixed-arity copy at
+                // all: on this target the call is still an ImMethodCall when varargs are eliminated.
+                for (ImMethodCall c : collectMonomorphicVarargMethodCalls()) {
+                    ImFunction implementation = c.getMethod().getImplementation();
+                    List<ImExpr> arguments = receiverAndArguments(c);
+                    if (shouldSpecialise(arguments)
+                        && !varargFuncs.contains(implementation, arguments.size())) {
+                        generateVarargFunc(implementation, arguments, c);
+                        generated = true;
+                    }
+                }
             }
         }
 
-        // remove original vararg functions:
-        prog.getFunctions().removeIf(f -> f.hasFlag(IS_VARARG));
+        if (!luaTarget) {
+            // remove original vararg functions:
+            prog.getFunctions().removeIf(f -> f.hasFlag(IS_VARARG));
+        }
 
         // rewrite calls to use new functions:
         // (need to collect vararg calls again, because first phase can create copies of calls)
         for (ImFunctionCall call : collectVarargCalls()) {
-            redirectCall(call, varargFuncs.get(call.getFunc(), call.getArguments().size()));
+            ImFunction newFunc = varargFuncs.get(call.getFunc(), call.getArguments().size());
+            if (newFunc != null) {
+                redirectCall(call, newFunc);
+            }
+        }
+        if (luaTarget) {
+            for (ImMethodCall call : collectMonomorphicVarargMethodCalls()) {
+                ImFunction implementation = call.getMethod().getImplementation();
+                ImFunction newFunc = varargFuncs.get(implementation, 1 + call.getArguments().size());
+                if (newFunc != null) {
+                    redirectMethodCall(call, newFunc);
+                }
+            }
         }
     }
 
+    /** A method call which can only ever reach one implementation, and that implementation is vararg. */
+    private Collection<ImMethodCall> collectMonomorphicVarargMethodCalls() {
+        final Collection<ImMethodCall> calls = new ArrayList<>();
+        prog.accept(new ImProg.DefaultVisitor() {
+            @Override
+            public void visit(ImMethodCall c) {
+                super.visit(c);
+                ImMethod method = c.getMethod();
+                if (method != null && !method.getIsAbstract() && method.getImplementation() != null
+                    && method.getSubMethods().isEmpty() && method.getImplementation().hasFlag(IS_VARARG)) {
+                    calls.add(c);
+                }
+            }
+        });
+        return calls;
+    }
+
+    /** The implementation's argument list: the receiver is its first parameter. */
+    private static List<ImExpr> receiverAndArguments(ImMethodCall call) {
+        List<ImExpr> arguments = new ArrayList<>(1 + call.getArguments().size());
+        arguments.add(call.getReceiver());
+        arguments.addAll(call.getArguments());
+        return arguments;
+    }
+
+    private void redirectMethodCall(ImMethodCall call, ImFunction newFunc) {
+        ImExprs args = JassIm.ImExprs(call.getReceiver().copy());
+        args.addAll(call.getArguments().removeAll());
+        call.replaceBy(JassIm.ImFunctionCall(call.getTrace(), newFunc, JassIm.ImTypeArguments(), args,
+            call.getTuplesEliminated(), CallType.NORMAL));
+    }
+
+    /** Whether a call gets a fixed-arity copy. Always on Jass; on Lua only within the parameter bound. */
+    private boolean shouldSpecialise(ImFunctionCall call) {
+        return shouldSpecialise(call.getArguments());
+    }
+
+    /**
+     * Counted after tuple flattening, because that is what the emitted parameter list costs: twenty
+     * four-field tuples are eighty parameters, not twenty.
+     */
+    private boolean shouldSpecialise(List<ImExpr> arguments) {
+        if (!luaTarget) {
+            return true;
+        }
+        int parameters = 0;
+        for (ImExpr argument : arguments) {
+            parameters += ImHelper.flattenedJassArity(argument.attrTyp());
+            if (parameters > LUA_MAX_SPECIALISED_VARARG_PARAMETERS) {
+                return false;
+            }
+        }
+        return true;
+    }
+
     @NotNull
     private Collection<ImFunctionCall> collectVarargCalls() {
         // Collect all calls to vararg functions
@@ -70,13 +182,17 @@ public void visit(ImFunctionCall c) {
      * for the function call.
      */
     private void generateVarargFunc(ImFunctionCall sourceCall) {
-        ImFunction func = sourceCall.getFunc();
-        int numberOfParams = sourceCall.getArguments().size();
-        int jassParameterCount = sourceCall.getArguments().stream()
+        generateVarargFunc(sourceCall.getFunc(), sourceCall.getArguments(), sourceCall);
+    }
+
+    /** {@code arguments} are in the callee's parameter order, so for a method they start with the receiver. */
+    private void generateVarargFunc(ImFunction func, List<ImExpr> arguments, Element trace) {
+        int numberOfParams = arguments.size();
+        int jassParameterCount = arguments.stream()
             .mapToInt(argument -> ImHelper.flattenedJassArity(argument.attrTyp()))
             .sum();
-        if (jassParameterCount > JASS_MAX_PARAMETERS) {
-            throw new CompileError(sourceCall, "Vararg call would generate " + jassParameterCount
+        if (!luaTarget && jassParameterCount > JASS_MAX_PARAMETERS) {
+            throw new CompileError(trace, "Vararg call would generate " + jassParameterCount
                 + " Jass parameters; the maximum is " + JASS_MAX_PARAMETERS
                 + ". Use multiple calls (for example with the cascade operator) or pass a collection instead.");
         }
@@ -91,6 +207,19 @@ private void generateVarargFunc(ImFunctionCall sourceCall) {
 
         // Create new function
         ImFunction newFunc = ReferenceRewritingCopy.copy(func);
+        // The copy retargets the function's own references, so a recursive call inside it now names
+        // the copy. That is wrong whenever the recursion uses a different argument count: the call
+        // must go back to naming the vararg original, so the rewrite below maps it to a copy of its
+        // own arity like any other call.
+        newFunc.accept(new Element.DefaultVisitor() {
+            @Override
+            public void visit(ImFunctionCall call) {
+                super.visit(call);
+                if (call.getFunc() == newFunc) {
+                    call.setFunc(func);
+                }
+            }
+        });
         newFunc.setName(func.getName() + "_" + argumentSize);
         // replace vararg with special parameters:
         ImVar varargParam = newFunc.getParameters().remove(newFunc.getParameters().size() - 1);
@@ -131,7 +260,9 @@ public void visit(ImVarargLoop imLoop) {
             params.addAll(list);
 
             // generate function for this new call
-            generateVarargFunc(call);
+            if (shouldSpecialise(call)) {
+                generateVarargFunc(call);
+            }
         }