Inline monomorphic methods in optimized Lua

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

@@ -917,6 +917,13 @@ public LuaCompilationUnit transformProgToLua() {
         // inliner
         stage = 5;
         if (runArgs.isInline()) {
+            // Expose hot loop calls which cannot dispatch anywhere else to the ordinary inliner.
+            // Calls outside loops keep their established method/slot representation.
+            beginPhase(5, "lower monomorphic Lua method calls");
+            LuaMethodCallLowering.transform(imProg);
+            imTranslator.assertProperties();
+            timeTaker.endPhase();
+
             beginPhase(5, "inlining");
             optimizer.doInlining();
             imTranslator2.assertProperties();

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

@@ -0,0 +1,67 @@
+package de.peeeq.wurstscript.translation.imtranslation;
+
+import de.peeeq.wurstscript.jassIm.*;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Lowers loop-local method calls with exactly one possible implementation to ordinary function
+ * calls on Lua.
+ *
+ * <p>The Lua emitter has always used the same direct-call fast path. Performing the lowering before
+ * optimization exposes hot calls to the ordinary inliner without guessing about receiver types or
+ * generated method names. Calls outside loops, and calls which can participate in virtual dispatch,
+ * remain untouched to avoid broad code-shape churn for a speculative gain.
+ */
+public final class LuaMethodCallLowering {
+
+    private LuaMethodCallLowering() {
+    }
+
+    public static int transform(ImProg prog) {
+        List<ImMethodCall> calls = new ArrayList<>();
+        prog.accept(new ImProg.DefaultVisitor() {
+            @Override
+            public void visit(ImMethodCall call) {
+                super.visit(call);
+                if (isInsideLoop(call) && canLowerDirectly(call.getMethod())) {
+                    calls.add(call);
+                }
+            }
+        });
+
+        for (ImMethodCall call : calls) {
+            lower(call);
+        }
+        return calls.size();
+    }
+
+    private static boolean isInsideLoop(ImMethodCall call) {
+        Element owner = call.getParent();
+        while (owner != null && !(owner instanceof ImFunction)) {
+            if (owner instanceof ImLoop || owner instanceof ImVarargLoop) {
+                return true;
+            }
+            owner = owner.getParent();
+        }
+        return false;
+    }
+
+    public static boolean canLowerDirectly(ImMethod method) {
+        return method != null
+            && !method.getIsAbstract()
+            && method.getImplementation() != null
+            && method.getSubMethods().isEmpty();
+    }
+
+    private static void lower(ImMethodCall call) {
+        ImExpr receiver = call.getReceiver();
+        receiver.setParent(null);
+        ImExprs arguments = JassIm.ImExprs(receiver);
+        arguments.addAll(call.getArguments().removeAll());
+        call.replaceBy(JassIm.ImFunctionCall(call.getTrace(), call.getMethod().getImplementation(),
+            JassIm.ImTypeArguments(call.getTypeArguments().removeAll()), arguments,
+            call.getTuplesEliminated(), CallType.NORMAL));
+    }
+}

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

@@ -148,8 +148,8 @@ private Collection<ImMethodCall> collectMonomorphicVarargMethodCalls() {
             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)) {
+                if (LuaMethodCallLowering.canLowerDirectly(method)
+                    && method.getImplementation().hasFlag(IS_VARARG)) {
                     calls.add(c);
                 }
             }

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/translation/lua/translation/ExprTranslation.java

@@ -5,6 +5,7 @@
 import de.peeeq.wurstscript.jassIm.*;
 import de.peeeq.wurstscript.luaAst.*;
 import de.peeeq.wurstscript.translation.imtranslation.ImTranslator;
+import de.peeeq.wurstscript.translation.imtranslation.LuaMethodCallLowering;
 import de.peeeq.wurstscript.types.TypesHelper;
 
 import java.util.Optional;
@@ -187,9 +188,7 @@ public static LuaExpr translate(ImMemberAccess e, LuaTranslator tr) {
 
     public static LuaExpr translate(ImMethodCall e, LuaTranslator tr) {
         ImMethod method = e.getMethod();
-        if (!method.getIsAbstract()
-            && method.getImplementation() != null
-            && method.getSubMethods().isEmpty()) {
+        if (LuaMethodCallLowering.canLowerDirectly(method)) {
             LuaExprlist args = LuaAst.LuaExprlist();
             args.add(e.getReceiver().translateToLua(tr));
             for (ImExpr arg : e.getArguments()) {

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

@@ -2016,6 +2016,118 @@ public void tinyPopularLuaHelpersInlineWithoutAnnotations() {
             compiled.contains("arithmetic("));
     }
 
+    @Test
+    public void tinyMonomorphicMethodsInlineOnLua() {
+        String compiled = compileOptimizedLua(
+            "tinyMonomorphicMethodsInlineOnLua",
+            "package Test",
+            "native consume(int value)",
+            "class Accumulator",
+            "    int offset",
+            "    construct(int offset)",
+            "        this.offset = offset",
+            "    function add(int value) returns int",
+            "        return value + offset",
+            "abstract class Operation",
+            "    abstract function apply(int value) returns int",
+            "class DoubleOperation extends Operation",
+            "    override function apply(int value) returns int",
+            "        return value * 2",
+            "@noinline function hotLoop(Accumulator accumulator)",
+            "    var i = 0",
+            "    while i < 16",
+            "        consume(accumulator.add(i))",
+            "        i++",
+            "@noinline function dynamicCall(Operation operation)",
+            "    consume(operation.apply(3))",
+            "init",
+            "    hotLoop(new Accumulator(2))",
+            "    dynamicCall(new DoubleOperation())"
+        );
+
+        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["));
+        String dynamicBody = topLevelFunctionBodyWithPrefix(compiled, "dynamicCall");
+        assertTrue("a genuinely virtual method call must retain dispatch:\n" + dynamicBody,
+            dynamicBody.contains("dispatch_"));
+    }
+
+    @Test
+    public void monomorphicMethodInliningEvaluatesReceiverOnce() {
+        test().testLua(true).luaOnly(false).optimize().executeProg().lines(
+            "package Test",
+            "native testSuccess()",
+            "int receiverEvaluations = 0",
+            "class Accumulator",
+            "    int offset",
+            "    construct(int offset)",
+            "        this.offset = offset",
+            "    function add(int value) returns int",
+            "        return value + offset",
+            "function makeAccumulator() returns Accumulator",
+            "    receiverEvaluations++",
+            "    return new Accumulator(2)",
+            "function evaluate() returns int",
+            "    var result = 0",
+            "    var i = 0",
+            "    while i < 1",
+            "        result = makeAccumulator().add(5)",
+            "        i++",
+            "    return result",
+            "init",
+            "    if evaluate() == 7 and receiverEvaluations == 1",
+            "        testSuccess()"
+        );
+    }
+
+    @Test
+    public void monomorphicMethodInliningKeepsCallbackBoundary() {
+        String compiled = compileOptimizedLua(
+            "monomorphicMethodInliningKeepsCallbackBoundary",
+            "package Test",
+            "native consume(code callback)",
+            "function callback()",
+            "class Registrar",
+            "    function install()",
+            "        consume(function callback)",
+            "@noinline function hotPath(Registrar registrar)",
+            "    var i = 0",
+            "    while i < 1",
+            "        registrar.install()",
+            "        i++",
+            "init",
+            "    hotPath(new Registrar())"
+        );
+
+        assertFunctionBodyContains(compiled, "hotPath", "Registrar_install(", true);
+    }
+
+    @Test
+    public void monomorphicMethodInliningKeepsLocalPlayerBoundary() {
+        String compiled = compileOptimizedLua(
+            "monomorphicMethodInliningKeepsLocalPlayerBoundary",
+            "type player extends handle",
+            "package Test",
+            "@extern native GetLocalPlayer() returns player",
+            "native consume(bool value)",
+            "class Probe",
+            "    function isLocal() returns bool",
+            "        return GetLocalPlayer() != null",
+            "@noinline function hotPath(Probe probe)",
+            "    var i = 0",
+            "    while i < 1",
+            "        consume(probe.isLocal())",
+            "        i++",
+            "init",
+            "    hotPath(new Probe())"
+        );
+
+        assertFunctionBodyContains(compiled, "hotPath", "Probe_isLocal(", true);
+    }
+
     /**
      * Measured after Lua native lowering: unit_getX = 59 IM nodes,
      * unit_getAbilityLevel = 63, real_floor = 35, and __wurst_intDiv = 31. Each helper is called
@@ -2050,14 +2162,19 @@ public void optimizedUnitSpatialIndexInnerLoopUsesRawLuaOperations() {
             "    query(vec2(0., 0.))"
         );
 
-        // The spatial-index helpers each have one call site, so the optimizer folds the whole
-        // query chain into our retained entry point. Inspect that surviving hot-loop owner.
+        // 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");
-        assertTrue("query loop must read the next-link array directly:\n" + body,
+        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,
             body.contains("UnitSpatialIndex_nextInCell["));
-        assertTrue("query loop must read cached X directly:\n" + body,
+        assertTrue("spatial-index hot loop must read cached X directly:\n" + body,
             body.contains("UnitSpatialIndex_lastX["));
-        assertTrue("query loop must read cached Y directly:\n" + body,
+        assertTrue("spatial-index hot loop must read cached Y directly:\n" + body,
             body.contains("UnitSpatialIndex_lastY["));
         assertFalse("typed array reads must not retain assurance calls:\n" + body,
             body.contains("__wurst_ensure"));