fix lua inlining by avoiding breaking function references (#1168)

Author: Frotty

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/RunArgs.java

@@ -132,7 +132,6 @@ public RunArgs(String... args) {
 
         addOptionWithArg("functionSplitLimit", "The maximum number of operations in a function before it is split by the function splitter (used for compiletime functions)",
             s -> functionSplitLimit = Integer.parseInt(s, 10));
-
         optionPrettyPrint = addOption("prettyPrint", "Pretty print the input file, or all sub-directory if the given path is: '...'");
 
         nextArg:
@@ -149,6 +148,20 @@ public RunArgs(String... args) {
                         o.isSet = true;
                         continue nextArg;
                     } else if ((o.argHandler != null && isDoubleArg(a, o))) {
+                        String value = a.substring(a.indexOf(" ") + 1).trim();
+                        if (value.isEmpty()) {
+                            throw new RuntimeException("Missing value for option: -" + o.name);
+                        }
+                        o.argHandler.accept(value);
+                        o.isSet = true;
+                        continue nextArg;
+                    } else if (o.argHandler != null && isEqualsArg(a, o)) {
+                        String value = a.substring(a.indexOf("=") + 1).trim();
+                        if (value.isEmpty()) {
+                            throw new RuntimeException("Missing value for option: -" + o.name);
+                        }
+                        o.argHandler.accept(value);
+                        o.isSet = true;
                         continue nextArg;
                     }
                 }
@@ -170,6 +183,10 @@ private boolean isDoubleArg(String arg, RunOption option) {
         return (arg.contains(" ") && ("-" + option.name).equals(arg.substring(0, arg.indexOf(" "))));
     }
 
+    private boolean isEqualsArg(String arg, RunOption option) {
+        return arg.startsWith("-" + option.name + "=");
+    }
+
     private RunOption addOption(String name, String descr) {
         RunOption opt = new RunOption(name, descr);
         options.add(opt);

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/translation/imoptimizer/GlobalsInliner.java

@@ -16,7 +16,6 @@
 import java.util.stream.Collectors;
 
 public class GlobalsInliner implements OptimizerPass {
-
     public int optimize(ImTranslator trans) {
         int obsoleteCount = 0;
         ImProg prog = trans.getImProg();
@@ -44,7 +43,7 @@ public int optimize(ImTranslator trans) {
                 ImVarWrite obs = null;
                 for (ImVarWrite write : v.attrWrites()) {
                     ImFunction func = write.getNearestFunc();
-                    if (isInInit(func)) {
+                    if (isInInitGlobals(func)) {
                         right = write.getRight();
                         obs = write;
                         break;
@@ -67,19 +66,21 @@ public int optimize(ImTranslator trans) {
                 List<ImVarWrite> initWrites = new ArrayList<>();
                 for (ImVarWrite imVarWrite : v.attrWrites()) {
                     ImFunction nearestFunc = imVarWrite.getNearestFunc();
-                    if (isInInit(nearestFunc)) {
+                    if (isInInitGlobals(nearestFunc)) {
                         initWrites.add(imVarWrite);
                     }
                 }
                 if (initWrites.size() == 1) {
                     if(v.getType() instanceof ImSimpleType) {
-                        ImExpr write = v.attrWrites().iterator().next().getRight();
+                        ImVarWrite initWrite = initWrites.get(0);
+                        ImExpr write = initWrite.getRight();
                         try {
                             ImExpr defaultValue = ImHelper.defaultValueForType((ImSimpleType) v.getType());
                             boolean isDefault = defaultValue.structuralEquals(write);
                             if (isDefault) {
-                                // Assignment is default value and can be removed
-                                v.attrWrites().iterator().next().replaceBy(ImHelper.nullExpr());
+                                // Only remove the init write when it assigns the default value.
+                                // Never touch non-init writes here.
+                                initWrite.replaceBy(ImHelper.nullExpr());
                             }
                         } catch (Exception e) {
                             throw new CompileError(write.attrTrace().attrErrorPos(),
@@ -140,9 +141,8 @@ public String getName() {
     }
 
 
-    private static boolean isInInit(ImFunction func) {
-        return func != null && (func.getName().startsWith("init_") || func.getName().equals("main") || func.getName().startsWith("InitTrig_")
-                || func.getName().equals("initGlobals"));
+    private static boolean isInInitGlobals(ImFunction func) {
+        return func != null && func.getName().equals("initGlobals");
     }
 
 }

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/translation/imoptimizer/ImInliner.java

@@ -28,6 +28,7 @@ public class ImInliner {
     private final Map<ImFunction, Integer> callCounts = Maps.newLinkedHashMap();
     private final Map<ImFunction, Integer> funcSizes = Maps.newLinkedHashMap();
     private final Set<ImFunction> done = Sets.newLinkedHashSet();
+    private final Map<ImFunction, Boolean> containsFuncRefCache = Maps.newLinkedHashMap();
     private final double inlineTreshold = 50;
 
     static {
@@ -49,8 +50,7 @@ public void doInlining() {
     }
 
     private void inlineFunctions() {
-
-        for (ImFunction f : ImHelper.calculateFunctionsOfProg(prog)) {
+        for (ImFunction f : sortedFunctions(ImHelper.calculateFunctionsOfProg(prog))) {
             inlineFunctions(f);
         }
     }
@@ -61,7 +61,7 @@ private void inlineFunctions(ImFunction f) {
         }
         done.add(f);
         // first inline functions called from this function
-        for (ImFunction called : translator.getCalledFunctions().get(f)) {
+        for (ImFunction called : sortedFunctions(translator.getCalledFunctions().get(f))) {
             inlineFunctions(called);
         }
         boolean[] changed = new boolean[]{false};
@@ -73,10 +73,10 @@ private ImFunction inlineFunctions(ImFunction f, Element parent, int parentI, El
         if (e instanceof ImFunctionCall) {
             ImFunctionCall call = (ImFunctionCall) e;
             ImFunction called = call.getFunc();
-            boolean canInline = f != called && shouldInline(call, called);
+            boolean canInline = f != called && shouldInline(f, call, called);
             if (LOG_INLINER) {
                 String msg = "[INLINER] caller=" + f.getName() + " callee=" + called.getName() + " decision=" + (canInline ? "inline" : "keep") +
-                    (canInline ? "" : " reason=" + skipReason(call, called));
+                    (canInline ? "" : " reason=" + skipReason(f, call, called));
                 WLogger.info(msg);
                 System.out.println(msg);
             }
@@ -109,13 +109,19 @@ private ImFunction inlineFunctions(ImFunction f, Element parent, int parentI, El
         return null;
     }
 
-    private String skipReason(ImFunctionCall call, ImFunction f) {
+    private String skipReason(ImFunction caller, ImFunctionCall call, ImFunction f) {
         if (f.isNative()) {
             return "native";
         }
         if (call.getCallType() == CallType.EXECUTE) {
             return "execute_call";
         }
+        if (translator.isLuaTarget() && !maxOneReturn(f)) {
+            return "lua_multi_return_inline_disabled";
+        }
+        if (translator.isLuaTarget() && containsFuncRef(f)) {
+            return "lua_callback_funcref_barrier";
+        }
         if (!inlinableFunctions.contains(f)) {
             return "not_in_inlinable_set";
         }
@@ -276,12 +282,18 @@ private ImStmt rewriteStmtForEarlyReturn(ImStmt s, ImVar doneVar, ImVar retVar)
     }
 
     private void rateInlinableFunctions() {
-        for (Map.Entry<ImFunction, ImFunction> edge : translator.getCalledFunctions().entries()) {
+        List<Map.Entry<ImFunction, ImFunction>> edges = new ArrayList<>(translator.getCalledFunctions().entries());
+        edges.sort((a, b) -> {
+            int c = functionSortKey(a.getKey()).compareTo(functionSortKey(b.getKey()));
+            if (c != 0) return c;
+            return functionSortKey(a.getValue()).compareTo(functionSortKey(b.getValue()));
+        });
+        for (Map.Entry<ImFunction, ImFunction> edge : edges) {
             // For bloat control we need how often a function is used (incoming edges),
             // not how many calls it performs itself (outgoing edges).
             incCallCount(edge.getValue());
         }
-        for (ImFunction f : inlinableFunctions) {
+        for (ImFunction f : sortedFunctions(inlinableFunctions)) {
             int size = estimateSize(f);
             funcSizes.put(f, size);
         }
@@ -322,10 +334,25 @@ private int getFuncSize(ImFunction f) {
         }
     }
 
-    private boolean shouldInline(ImFunctionCall call, ImFunction f) {
+    private boolean shouldInline(ImFunction caller, ImFunctionCall call, ImFunction f) {
         if (f.isNative() || call.getCallType() == CallType.EXECUTE) {
             return false;
         }
+        if (translator.isLuaTarget() && !maxOneReturn(f)) {
+            // Conservative safety: Lua inliner multi-return rewriting is not yet fully robust
+            // across all lowered patterns. Keep call semantics intact for now.
+            return false;
+        }
+        if (translator.isLuaTarget() && containsFuncRef(f)) {
+            // Functions that build callback refs are lowered with Lua-specific wrappers/xpcall.
+            // Keeping them as standalone calls avoids callback context/vararg scope breakage.
+            return false;
+        }
+        if (isLuaTypeCastingCompatFunction(f)) {
+            // In Lua these compat wrappers are rewritten to object index helpers.
+            // If they are inlined beforehand, old TypeCasting bodies leak through.
+            return false;
+        }
 
         double threshold = inlineTreshold;
         for (ImExpr arg : call.getArguments()) {
@@ -362,6 +389,31 @@ private boolean containsCallTo(ImFunction f, Element e) {
         return false;
     }
 
+    private boolean containsFuncRef(ImFunction f) {
+        if (f == null) {
+            return false;
+        }
+        Boolean cached = containsFuncRefCache.get(f);
+        if (cached != null) {
+            return cached;
+        }
+        boolean result = containsFuncRef(f.getBody());
+        containsFuncRefCache.put(f, result);
+        return result;
+    }
+
+    private boolean containsFuncRef(Element e) {
+        if (e instanceof ImFuncRef) {
+            return true;
+        }
+        for (int i = 0; i < e.size(); i++) {
+            if (containsFuncRef(e.get(i))) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     private int estimateSize(ImFunction f) {
         int[] r = new int[]{0};
         estimateSize(f.getBody(), r);
@@ -390,24 +442,46 @@ private int getCallCount(ImFunction f) {
     }
 
     private void collectInlinableFunctions() {
-        for (ImFunction f : ImHelper.calculateFunctionsOfProg(prog)) {
+        for (ImFunction f : sortedFunctions(ImHelper.calculateFunctionsOfProg(prog))) {
             if (isInlineCandidate(f)) {
                 inlinableFunctions.add(f);
             }
         }
         // Some call targets can survive in the call graph but not in prog/classes lists.
-        for (ImFunction f : translator.getCalledFunctions().values()) {
+        for (ImFunction f : sortedFunctions(translator.getCalledFunctions().values())) {
             if (isInlineCandidate(f)) {
                 inlinableFunctions.add(f);
             }
         }
     }
 
+    private List<ImFunction> sortedFunctions(Collection<ImFunction> functions) {
+        List<ImFunction> r = new ArrayList<>(functions);
+        r.sort(Comparator.comparing(this::functionSortKey));
+        return r;
+    }
+
+    private String functionSortKey(ImFunction f) {
+        if (f == null) {
+            return "";
+        }
+        StringBuilder sb = new StringBuilder();
+        sb.append(f.getName()).append("|");
+        sb.append(f.getReturnType()).append("|");
+        for (ImVar p : f.getParameters()) {
+            sb.append(p.getType()).append(",");
+        }
+        return sb.toString();
+    }
+
     private boolean isInlineCandidate(ImFunction f) {
         if (f.hasFlag(FunctionFlagEnum.IS_COMPILETIME_NATIVE) || f.hasFlag(FunctionFlagEnum.IS_NATIVE)) {
             // do not inline natives
             return false;
         }
+        if (isLuaTypeCastingCompatFunction(f)) {
+            return false;
+        }
         if (f == translator.getGlobalInitFunc()) {
             return false;
         }
@@ -419,6 +493,20 @@ private boolean isInlineCandidate(ImFunction f) {
         return true;
     }
 
+    private boolean isLuaTypeCastingCompatFunction(ImFunction f) {
+        if (!translator.isLuaTarget() || f == null) {
+            return false;
+        }
+        de.peeeq.wurstscript.ast.Element trace = f.attrTrace();
+        if (trace instanceof de.peeeq.wurstscript.ast.FuncDef fd
+            && fd.attrNearestPackage() instanceof de.peeeq.wurstscript.ast.WPackage p
+            && "TypeCasting".equals(p.getName())) {
+            String name = fd.getName();
+            return name.endsWith("FromIndex") || name.endsWith("ToIndex");
+        }
+        return false;
+    }
+
     private boolean maxOneReturn(ImFunction f) {
         return maxOneReturn(f.getBody());
     }