Lower KeyedTable membership to a native Lua table

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

@@ -0,0 +1,79 @@
+package de.peeeq.wurstscript.translation.imtranslation;
+
+import de.peeeq.wurstscript.CompilerIntrinsics;
+import de.peeeq.wurstscript.ast.FuncDef;
+import de.peeeq.wurstscript.jassIm.ImFunction;
+import de.peeeq.wurstscript.jassIm.ImVoid;
+import de.peeeq.wurstscript.types.TypesHelper;
+
+/**
+ * Recognises the compiler-owned {@code KeyedTable} membership operations.
+ *
+ * <p>Jass has no hashing, so every keyed structure in the library bottoms out in the hashtable
+ * natives. Those take a (parent, child) pair and are emitted on Lua as a nested table plus a nil
+ * check, so a membership test costs a call and two indexes on a runtime that is already a hash
+ * table - a Jass compromise carried into Lua, which AGENTS.md section 7 asks us not to do. On Lua
+ * each keyed table is its own table and the element is the key, so each operation is one index.
+ *
+ * <p>Matching is by <b>declaration</b>, not by name: a function qualifies only if it is annotated
+ * {@code @compilerintrinsic} and its IM signature is exactly the one the lowering assumes.
+ * Name-only matching would silently replace the body of any user function that happened to share
+ * the name, including one with different types, which section 7 rules out - semantic identity must
+ * come from the declaration and its structural signature, never from string comparison.
+ *
+ * <p><b>Iteration is deliberately absent.</b> Enumerating a Lua table needs {@code pairs()}, whose
+ * order follows internal hash layout and so differs between clients, which desyncs a lockstep
+ * game. Anything that must be iterated needs a separately maintained insertion-ordered array -
+ * what SparseSet's dense half provides. This primitive is membership only.
+ */
+public final class LuaKeyedTable {
+
+    /** Source-level names of the operations. Necessary to identify them, never sufficient. */
+    private static final String CREATE = "keyedTableCreate";
+    private static final String ADD = "keyedTableAdd";
+    private static final String CONTAINS = "keyedTableContains";
+    private static final String REMOVE = "keyedTableRemove";
+
+    /** Stub names whose Lua bodies live in {@code LuaNatives}. */
+    public static final String NATIVE_CREATE = "__wurst_keyedTableCreate";
+    public static final String NATIVE_ADD = "__wurst_keyedTableAdd";
+    public static final String NATIVE_CONTAINS = "__wurst_keyedTableContains";
+    public static final String NATIVE_REMOVE = "__wurst_keyedTableRemove";
+
+    private LuaKeyedTable() {
+    }
+
+    /**
+     * The {@code __wurst_} stub {@code f} should be lowered to on Lua, or null if {@code f} is not
+     * a compiler-owned KeyedTable operation.
+     */
+    public static String nativeStubFor(ImFunction f) {
+        if (!(f.attrTrace() instanceof FuncDef fd)
+            || !fd.attrHasAnnotation(CompilerIntrinsics.ANNOTATION)) {
+            return null;
+        }
+        int params = f.getParameters().size();
+        return switch (fd.getName()) {
+            case CREATE -> params == 0 && TypesHelper.isIntType(f.getReturnType())
+                ? NATIVE_CREATE : null;
+            case ADD -> isKeyedPair(f) && f.getReturnType() instanceof ImVoid
+                ? NATIVE_ADD : null;
+            case REMOVE -> isKeyedPair(f) && f.getReturnType() instanceof ImVoid
+                ? NATIVE_REMOVE : null;
+            case CONTAINS -> isKeyedPair(f) && TypesHelper.isBoolType(f.getReturnType())
+                ? NATIVE_CONTAINS : null;
+            default -> null;
+        };
+    }
+
+    /**
+     * A (table, key) parameter pair. Both are {@code int} at source level: on Jass everything is an
+     * integer anyway, and on Lua {@code castTo int} is the identity for class types, so the value
+     * arriving here is already a usable table key with reference identity.
+     */
+    private static boolean isKeyedPair(ImFunction f) {
+        return f.getParameters().size() == 2
+            && TypesHelper.isIntType(f.getParameters().get(0).getType())
+            && TypesHelper.isIntType(f.getParameters().get(1).getType());
+    }
+}

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

@@ -146,6 +146,25 @@ public static void transform(ImProg prog, ImTranslator translator) {
             public void visit(ImFunctionCall call) {
                 super.visit(call);
                 ImFunction f = call.getFunc();
+                // KeyedTable membership becomes a table keyed directly by the element. Done here,
+                // before optimization, rather than at emission: the inliner runs in between, and a
+                // call inlined before an emission-time rewrite would keep the hashtable body while
+                // a surviving one got the Lua table - mixing an integer class id with a table index
+                // for the same value. Replacing the call makes every site agree.
+                String keyedStub = LuaKeyedTable.nativeStubFor(f);
+                if (keyedStub != null) {
+                    ImFunction replacement = specialNativeStubs.computeIfAbsent(keyedStub,
+                        name -> createNativeStub(name, f));
+                    if (!deferredAdditions.contains(replacement)) {
+                        deferredAdditions.add(replacement);
+                    }
+                    call.replaceBy(JassIm.ImFunctionCall(
+                        call.attrTrace(), replacement,
+                        JassIm.ImTypeArguments(),
+                        call.getArguments().copy(),
+                        false, CallType.NORMAL));
+                    return;
+                }
                 if (ENABLE_SELECTIVE_GET_HANDLE_ID_SHIMMING && isCompatGetHandleIdFunction(f)) {
                     if (shouldRewriteGetHandleId(call)) {
                         ImFunction replacement = specialNativeStubs.computeIfAbsent("__wurst_GetHandleId",

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

@@ -271,6 +271,28 @@ public class LuaNatives {
             f.getBody().add(LuaAst.LuaLiteral("__wurst_enumDestructable_override = prev"));
         });
 
+        // KeyedTable: membership keyed directly by the element, one index per operation.
+        // Lowered to these stubs by LuaKeyedTable/LuaNativeLowering before the inliner runs, so
+        // every call site agrees on the representation. No iteration is offered: pairs() order
+        // differs between clients and desyncs a lockstep game.
+        addNative("__wurst_keyedTableCreate", f ->
+            f.getBody().add(LuaAst.LuaLiteral("return {}")));
+        addNative("__wurst_keyedTableAdd", f -> {
+            f.getParams().add(LuaAst.LuaVariable("t", LuaAst.LuaNoExpr()));
+            f.getParams().add(LuaAst.LuaVariable("k", LuaAst.LuaNoExpr()));
+            f.getBody().add(LuaAst.LuaLiteral("t[k] = true"));
+        });
+        addNative("__wurst_keyedTableContains", f -> {
+            f.getParams().add(LuaAst.LuaVariable("t", LuaAst.LuaNoExpr()));
+            f.getParams().add(LuaAst.LuaVariable("k", LuaAst.LuaNoExpr()));
+            f.getBody().add(LuaAst.LuaLiteral("return t[k] ~= nil"));
+        });
+        addNative("__wurst_keyedTableRemove", f -> {
+            f.getParams().add(LuaAst.LuaVariable("t", LuaAst.LuaNoExpr()));
+            f.getParams().add(LuaAst.LuaVariable("k", LuaAst.LuaNoExpr()));
+            f.getBody().add(LuaAst.LuaLiteral("t[k] = nil"));
+        });
+
         addNative(Arrays.asList("InitHashtable", "__wurst_InitHashtable"), f ->
             f.getBody().add(LuaAst.LuaLiteral("return { __wurst_ht_int = {}, __wurst_ht_bool = {}, __wurst_ht_real = {}, __wurst_ht_str = {}, __wurst_ht_handle = {} }")));
 

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

@@ -1868,6 +1868,172 @@ public void genericOverrideChainBindsGlobalStateSlotToMostSpecificImplInLua() th
         assertDoesNotContainRegex(compiled, "GlobalCheckState\\." + dispatchedSlot + "\\s*=\\s*NoOpState_NoOpState_update");
     }
 
+    /**
+     * KeyedTable source shared by the tests below: the Jass path built on the library's hashtable
+     * wrapper, which the Lua backend replaces with a table keyed directly by the element.
+     */
+    private static String[] keyedTableSource(String... usage) {
+        java.util.List<String> lines = new java.util.ArrayList<>(java.util.Arrays.asList(
+            "package KeyedTable",
+            "import Table",
+            "@compilerintrinsic public function keyedTableCreate() returns int",
+            "    return (new Table()) castTo int",
+            "@compilerintrinsic public function keyedTableAdd(int tbl, int key)",
+            "    (tbl castTo Table).saveBoolean(key, true)",
+            "@compilerintrinsic public function keyedTableContains(int tbl, int key) returns boolean",
+            "    return (tbl castTo Table).loadBoolean(key)",
+            "@compilerintrinsic public function keyedTableRemove(int tbl, int key)",
+            "    (tbl castTo Table).removeBoolean(key)",
+            "endpackage"));
+        lines.addAll(java.util.Arrays.asList(usage));
+        return lines.toArray(new String[0]);
+    }
+
+    /** On Lua each keyed table is its own table and membership is a single index. */
+    @Test
+    public void keyedTableLowersToASingleLuaIndex() throws IOException {
+        test().testLua(true).withStdLib().lines(keyedTableSource(
+            "package Test",
+            "import KeyedTable",
+            "init",
+            "    let t = keyedTableCreate()",
+            "    keyedTableAdd(t, 7)",
+            "    print(keyedTableContains(t, 7).toString())",
+            "endpackage"));
+
+        String compiled = Files.toString(
+            new File("test-output/lua/LuaTranslationTests_keyedTableLowersToASingleLuaIndex.lua"), Charsets.UTF_8);
+
+        String add = getFunctionBody(compiled, "__wurst_keyedTableAdd");
+        String contains = getFunctionBody(compiled, "__wurst_keyedTableContains");
+        String create = getFunctionBody(compiled, "__wurst_keyedTableCreate");
+
+        assertTrue("add should be a single table store, was: " + add, add.contains("] = true"));
+        assertTrue("contains should be a single index, was: " + contains, contains.contains("] ~= nil"));
+        assertTrue("create should allocate a bare table, was: " + create, create.contains("return {}"));
+
+        // The whole point: no hashtable machinery on this path.
+        assertFalse("add must not go through the hashtable natives: " + add, add.contains("SaveBoolean"));
+        assertFalse("contains must not go through the hashtable natives: " + contains, contains.contains("LoadBoolean"));
+    }
+
+    /**
+     * Membership must mean the same thing on both backends: the Jass path runs the hashtable
+     * body, the Lua path runs a bare table, and the observable behaviour has to agree.
+     */
+    @Test
+    public void keyedTableMembershipAgreesOnBothBackends() {
+        test().testLua(true).executeProg(true).withStdLib().lines(keyedTableSource(
+            "package Test",
+            "import KeyedTable",
+            "init",
+            "    let t = keyedTableCreate()",
+            "    if keyedTableContains(t, 7)",
+            "        testFail(\"empty table reported 7 as present\")",
+            "    keyedTableAdd(t, 7)",
+            "    if not keyedTableContains(t, 7)",
+            "        testFail(\"7 was added but reported absent\")",
+            "    if keyedTableContains(t, 8)",
+            "        testFail(\"8 was never added but reported present\")",
+            "    keyedTableRemove(t, 7)",
+            "    if keyedTableContains(t, 7)",
+            "        testFail(\"7 was removed but reported present\")",
+            "    testSuccess()",
+            "endpackage"));
+    }
+
+    /**
+     * The lowering happens in IM before the inliner, so inlining cannot leave one call site on the
+     * hashtable body while another gets the Lua table - which would mix an integer class id with a
+     * table index for the same value and fail at runtime.
+     */
+    @Test
+    public void keyedTableMembershipSurvivesInlining() throws IOException {
+        test().testLua(true).executeProg(true).inline().withStdLib().lines(keyedTableSource(
+            "package Test",
+            "import KeyedTable",
+            "init",
+            "    let t = keyedTableCreate()",
+            "    keyedTableAdd(t, 7)",
+            "    if not keyedTableContains(t, 7)",
+            "        testFail(\"7 was added but reported absent under -inline\")",
+            "    keyedTableRemove(t, 7)",
+            "    if keyedTableContains(t, 7)",
+            "        testFail(\"7 was removed but reported present under -inline\")",
+            "    testSuccess()",
+            "endpackage"));
+
+        // Whether a given site actually gets inlined depends on the inliner's local-register
+        // budget, so assert the rule rather than the symptom: every call was replaced in IM, so
+        // the Table-backed originals are unreachable and collected. If any site had kept the old
+        // body, that function would still be here.
+        String compiled = Files.toString(
+            new File("test-output/lua/LuaTranslationTests_keyedTableMembershipSurvivesInlining.lua"),
+            Charsets.UTF_8);
+        assertFalse("no call site may keep the Table-backed keyedTableAdd",
+            compiled.contains("function keyedTableAdd("));
+        assertFalse("no call site may keep the Table-backed keyedTableContains",
+            compiled.contains("function keyedTableContains("));
+        assertTrue("membership must go through the lowered stub",
+            compiled.contains("__wurst_keyedTableContains"));
+    }
+
+    /**
+     * The lowering is opt-in by declaration. A user package that merely shares these names keeps
+     * its own body, so nothing silently changes meaning under it.
+     */
+    @Test
+    public void keyedTableWithoutTheAnnotationIsLeftAlone() throws IOException {
+        test().testLua(true).withStdLib().lines(
+            "package KeyedTable",
+            "import Table",
+            "public function keyedTableCreate() returns int",
+            "    return (new Table()) castTo int",
+            "public function keyedTableAdd(int tbl, int key)",
+            "    (tbl castTo Table).saveBoolean(key, true)",
+            "public function keyedTableContains(int tbl, int key) returns boolean",
+            "    return (tbl castTo Table).loadBoolean(key)",
+            "endpackage",
+            "package Test",
+            "import KeyedTable",
+            "init",
+            "    let t = keyedTableCreate()",
+            "    keyedTableAdd(t, 7)",
+            "    print(keyedTableContains(t, 7).toString())",
+            "endpackage");
+
+        String compiled = Files.toString(
+            new File("test-output/lua/LuaTranslationTests_keyedTableWithoutTheAnnotationIsLeftAlone.lua"),
+            Charsets.UTF_8);
+
+        assertFalse("an unannotated package must not be lowered to the keyed-table stubs",
+            compiled.contains("__wurst_keyedTableAdd"));
+        assertTrue("it should keep its own Table-backed body",
+            getFunctionBody(compiled, "keyedTableAdd").contains("Table_Table_saveBoolean"));
+    }
+
+    /**
+     * pairs() iteration order differs between clients and desyncs a lockstep game, so no emitted
+     * Lua may contain it. Cheap to assert and worth keeping regardless of this feature.
+     */
+    @Test
+    public void keyedTableEmitsNoTableIteration() throws IOException {
+        test().testLua(true).withStdLib().lines(keyedTableSource(
+            "package Test",
+            "import KeyedTable",
+            "init",
+            "    let t = keyedTableCreate()",
+            "    keyedTableAdd(t, 3)",
+            "    keyedTableRemove(t, 3)",
+            "endpackage"));
+
+        String compiled = Files.toString(
+            new File("test-output/lua/LuaTranslationTests_keyedTableEmitsNoTableIteration.lua"), Charsets.UTF_8);
+
+        assertFalse("emitted Lua must never iterate a table with pairs()", compiled.contains("pairs("));
+        assertFalse("emitted Lua must never iterate a table with next()", compiled.contains(" next("));
+    }
+
     private CU[] genericOverrideReproUnits() {
         return new CU[]{
             compilationUnit("fsmLib.wurst",

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

@@ -117,6 +117,7 @@ class TestConfig {
         private boolean stopOnFirstError = true;
         private boolean runCompiletimeFunctions;
         private boolean optimize;
+        private boolean inline;
         private boolean testLua = false;
         private boolean luaOnly = false;
         private boolean uncheckedDispatch = false;
@@ -161,6 +162,12 @@ public TestConfig executeTests(boolean b) {
             return this;
         }
 
+        /** Enables the IM inliner (-inline), which is a separate option from -opt. */
+        TestConfig inline() {
+            this.inline = true;
+            return this;
+        }
+
         TestConfig optimize() {
             this.optimize = true;
             return this;
@@ -281,6 +288,9 @@ private CompilationResult testScript() {
             if (optimize) {
                 runArgs = runArgs.with("-opt");
             }
+            if (inline) {
+                runArgs = runArgs.with("-inline");
+            }
             if (legacyJassTypeChecks) {
                 runArgs.setLegacyJassTypeChecks(true);
             }