Reserve preserved Lua class names safely

Author: Frotty

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

@@ -331,8 +331,8 @@ protected String uniqueName(String rawName) {
     }
 
     public LuaCompilationUnit translate() {
-        assertNoDanglingFunctionReferences(prog);
         collectPredefinedNames();
+        assertNoDanglingFunctionReferences(prog);
 
         normalizeFieldNames();
 
@@ -503,6 +503,15 @@ private void collectPredefinedNames() {
             }
         }
 
+        for (ImClass clazz : prog.getClasses()) {
+            for (ImFunction function : clazz.getFunctions()) {
+                if (NamePreservation.isPreserved(function)) {
+                    LuaFunction luaFunction = luaFunc.getFor(function);
+                    usedNames.add(luaFunction.getName());
+                }
+            }
+        }
+
         for (ImVar global : prog.getGlobals()) {
             if (global.getIsBJ()) {
                 setNameFromTrace(global);

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/validation/NamePreservation.java

@@ -7,7 +7,6 @@
 import de.peeeq.wurstscript.types.WurstType;
 import de.peeeq.wurstscript.types.WurstTypeArray;
 import de.peeeq.wurstscript.types.WurstTypeTuple;
-import org.eclipse.jdt.annotation.Nullable;
 
 import java.util.ArrayList;
 import java.util.LinkedHashMap;
@@ -18,6 +17,7 @@
 public final class NamePreservation {
 
     public static final String ANNOTATION = "@preserveName";
+    private static final String SYNTHETIC_MARKER = "__wurst_trve_preserve_name";
 
     private NamePreservation() {
     }
@@ -42,14 +42,14 @@ public static void preserve(ImFunction function) {
      * name-based side table. The marker remains attached to the AST definition and is copied to
      * the corresponding IM variable through its trace.
      */
-    public static @Nullable Annotation preserve(GlobalVarDef variable) {
+    public static void preserve(GlobalVarDef variable) {
         if (variable.hasAnnotation(ANNOTATION)) {
-            return null;
+            return;
         }
         Annotation marker = Ast.Annotation(variable.getSource(),
-            Ast.Identifier(variable.getSource(), ANNOTATION.substring(1)), Ast.Arguments());
+            Ast.Identifier(variable.getSource(), ANNOTATION.substring(1)),
+            Ast.Arguments(Ast.ExprStringVal(variable.getSource(), SYNTHETIC_MARKER)));
         variable.getModifiers().add(marker);
-        return marker;
     }
 
     /**
@@ -71,6 +71,21 @@ public void visit(GlobalVarDef variable) {
         return result;
     }
 
+    /** Removes markers synthesized for TRVE during an earlier validation run. */
+    public static void clearSyntheticMarkers(WurstModel model) {
+        model.accept(new Element.DefaultVisitor() {
+            @Override
+            public void visit(GlobalVarDef variable) {
+                super.visit(variable);
+                variable.getModifiers().removeIf(modifier -> modifier instanceof Annotation annotation
+                    && annotation.getAnnotationType().equalsIgnoreCase(ANNOTATION)
+                    && annotation.getArgs().size() == 1
+                    && annotation.getArgs().get(0) instanceof ExprStringVal value
+                    && value.getValS().equals(SYNTHETIC_MARKER));
+            }
+        });
+    }
+
     private static void addTupleComponentNames(RuntimeNameIndex index, String name, WurstType type,
                                                GlobalVarDef variable) {
         if (type instanceof WurstTypeArray array) {
@@ -88,26 +103,15 @@ private static void addTupleComponentNames(RuntimeNameIndex index, String name,
 
     public static final class RuntimeNameIndex {
         private final Map<String, List<GlobalVarDef>> globalsByName = new LinkedHashMap<>();
-        private final Map<GlobalVarDef, Annotation> syntheticMarkers = new LinkedHashMap<>();
 
         private void add(String name, GlobalVarDef variable) {
             globalsByName.computeIfAbsent(name, ignored -> new ArrayList<>()).add(variable);
         }
 
         public void preserve(String runtimeName) {
             for (GlobalVarDef variable : globalsByName.getOrDefault(runtimeName, List.of())) {
-                Annotation marker = NamePreservation.preserve(variable);
-                if (marker != null) {
-                    syntheticMarkers.put(variable, marker);
-                }
-            }
-        }
-
-        public void clearSyntheticMarkers() {
-            for (Map.Entry<GlobalVarDef, Annotation> entry : syntheticMarkers.entrySet()) {
-                entry.getKey().getModifiers().remove(entry.getValue());
+                NamePreservation.preserve(variable);
             }
-            syntheticMarkers.clear();
         }
     }
 

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstscript/validation/WurstValidator.java

@@ -90,9 +90,7 @@ public void validate(Collection<CompilationUnit> toCheck) {
             guaranteedClassFieldInitCache.clear();
             moduleFieldCopiesCache.clear();
             moduleFieldCopiesIndexed = false;
-            if (runtimeNameIndex != null) {
-                runtimeNameIndex.clearSyntheticMarkers();
-            }
+            NamePreservation.clearSyntheticMarkers(prog);
             runtimeNameIndex = NamePreservation.indexGlobals(prog);
 
             lightValidation(toCheck);

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

@@ -587,6 +587,27 @@ public void preserveNameAnnotationKeepsClassFunctionNameInLua() throws IOExcepti
             "Expected a preserved class function to keep its emitted name.\n" + output);
     }
 
+    @Test
+    public void preservedClassFunctionReservesItsLuaNameBeforeClassVariables() throws IOException {
+        test().testLua(true).luaOnly(true).executeProg(false).lines(
+            "package test",
+            "    class Foo",
+            "        @preserveName function bar()",
+            "            skip",
+            "    class Foo_bar",
+            "    init",
+            "        new Foo_bar",
+            "endpackage");
+
+        String output = Files.toString(
+            new File("./test-output/lua/OptimizerTests_preservedClassFunctionReservesItsLuaNameBeforeClassVariables.lua"),
+            Charsets.UTF_8);
+        assertTrue(output.contains("function Foo_bar("),
+            "Expected the preserved class function to keep its Lua name.\n" + output);
+        assertTrue(output.contains("Foo_bar1 = ({})"),
+            "Expected the colliding class variable to be uniqued around the preserved function.\n" + output);
+    }
+
     @Test
     public void test_tempVarRemover() throws IOException {
         test().lines(