Give vararg calls a fixed-arity copy on Lua (#1286)

Author: Frotty

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

@@ -13,6 +13,7 @@
 import java.util.stream.Collectors;
 
 import static de.peeeq.wurstscript.translation.imtranslation.FunctionFlagEnum.IS_VARARG;
+import static de.peeeq.wurstscript.translation.imtranslation.FunctionFlagEnum.PRESERVE_NAME;
 
 /**
  * Takes a program and eliminates vararg functions, replacing them with
@@ -21,30 +22,178 @@
 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) && !forwardsAVarargParameter(c.getArguments())
+                    && !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) && !forwardsAVarargParameter(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 && !forwardsAVarargParameter(call.getArguments())) {
+                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 && !forwardsAVarargParameter(receiverAndArguments(call))) {
+                    redirectMethodCall(call, newFunc);
+                }
+            }
+        }
+    }
+
+
+    /**
+     * Whether a call passes a vararg placeholder straight through, which is what the generated
+     * `new_C` wrapper of a vararg constructor does with its own parameter. The placeholder is a
+     * single node standing for however many arguments the caller actually passed, so the call's node
+     * count is not an arity: specialising by it would produce a fixed-arity callee and drop every
+     * argument after the first.
+     *
+     * <p>Only reachable on Lua. The forwarding call lives in the body of a vararg original, and a
+     * copy has its placeholder expanded into real parameters before anything looks at it again, so
+     * this matches only originals - which Jass removes and Lua retains.
+     *
+     * <p>Both the generation and the rewrite loop consult this. Skipping generation alone would not
+     * be enough: another call could have produced a copy at the same node count, and the rewrite
+     * would then redirect the forwarding call to it.
+     */
+    private static boolean forwardsAVarargParameter(List<ImExpr> arguments) {
+        for (ImExpr argument : arguments) {
+            if (argument instanceof ImVarAccess access && isVarargPlaceholder(access.getVar())) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    /** The trailing parameter of a function still marked vararg, as opposed to a local or a copy's. */
+    private static boolean isVarargPlaceholder(ImVar variable) {
+        if (variable.getParent() == null
+            || !(variable.getParent().getParent() instanceof ImFunction function)
+            || !function.hasFlag(IS_VARARG)) {
+            return false;
+        }
+        List<ImVar> parameters = function.getParameters();
+        return !parameters.isEmpty() && parameters.get(parameters.size() - 1) == variable;
+    }
+    /** 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(call.getTypeArguments().removeAll()), 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
@@ -70,13 +219,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 +244,31 @@ private void generateVarargFunc(ImFunctionCall sourceCall) {
 
         // Create new function
         ImFunction newFunc = ReferenceRewritingCopy.copy(func);
+        // ReferenceRewritingCopy retargets the function's own references - both call and reference
+        // nodes - so inside the copy they now name the copy. That is wrong for either kind. A
+        // recursive 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; a self reference must name the original too,
+        // because it is invoked at an arity this pass never sees.
+        newFunc.accept(new Element.DefaultVisitor() {
+            @Override
+            public void visit(ImFunctionCall call) {
+                super.visit(call);
+                if (call.getFunc() == newFunc) {
+                    call.setFunc(func);
+                }
+            }
+
+            @Override
+            public void visit(ImFuncRef ref) {
+                super.visit(ref);
+                // Lua only: nothing redirects a reference afterwards, so it keeps naming whatever it
+                // is set to here, and only this target retains the original. On Jass the original is
+                // removed below and pointing at it would leave the reference dangling.
+                if (luaTarget && ref.getFunc() == newFunc) {
+                    ref.setFunc(func);
+                }
+            }
+        });
         newFunc.setName(func.getName() + "_" + argumentSize);
         // replace vararg with special parameters:
         ImVar varargParam = newFunc.getParameters().remove(newFunc.getParameters().size() - 1);
@@ -131,16 +309,23 @@ public void visit(ImVarargLoop imLoop) {
             params.addAll(list);
 
             // generate function for this new call
-            generateVarargFunc(call);
+            if (shouldSpecialise(call)) {
+                generateVarargFunc(call);
+            }
         }
 
 
-        // Remove vararg flag
+        // Drop the vararg flag, and on Lua the name preservation with it. A preserved name is part
+        // of the map's Warcraft-facing API and belongs to the retained original, which is what
+        // external code calls at an arity this pass never sees. Since a copy shares the original's
+        // trace, and LuaTranslator.collectPredefinedNames() resets every preserved function to its
+        // trace's source name, an inherited flag would emit both under one name.
         List<FunctionFlag> list = new ArrayList<>();
         for (FunctionFlag flag : newFunc.getFlags()) {
-            if (flag != IS_VARARG) {
-                list.add(flag);
+            if (flag == IS_VARARG || (luaTarget && flag == PRESERVE_NAME)) {
+                continue;
             }
+            list.add(flag);
         }
         newFunc.setFlags(list);
         // Add new function to prog
@@ -166,7 +351,13 @@ public void visit(ImVarAccess va) {
 
     private void redirectCall(ImFunctionCall call, ImFunction newFunc) {
         // Redirect call to new function
-        ImFunctionCall newCall = JassIm.ImFunctionCall(call.getTrace(), newFunc, JassIm.ImTypeArguments(), JassIm.ImExprs(call.getArguments().removeAll()), call.getTuplesEliminated(), call.getCallType());
+        // Carry the type arguments over rather than assuming there are none. Jass erases generics
+        // long before this pass, so an empty list was always right there; on Lua the erasure happens
+        // elsewhere and this list is empty in practice too, but rebuilding the call should not be
+        // the step that decides that.
+        ImFunctionCall newCall = JassIm.ImFunctionCall(call.getTrace(), newFunc,
+            JassIm.ImTypeArguments(call.getTypeArguments().removeAll()),
+            JassIm.ImExprs(call.getArguments().removeAll()), call.getTuplesEliminated(), call.getCallType());
         call.replaceBy(newCall);
     }