Record what a specialised node was copied from, in one place

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

@@ -159,7 +159,7 @@ public boolean removeGarbage() {
                 int classFieldsBefore = c.getFields().size();
                 changes |= c.getFields().retainAll(c.getFields().stream()
                     .filter(field -> readVars.contains(field)
-                        || readVars.contains(trans.originalOfSpecializedField(field)))
+                        || readVars.contains(trans.canonical(field)))
                     .collect(Collectors.toSet()));
                 int classFieldsAfter = c.getFields().size();
                 totalGlobalsRemoved += classFieldsBefore - classFieldsAfter;

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

@@ -913,7 +913,7 @@ private void moveFunctionsOutOfClass(ImClass c) {
                 ImTypeVar copy = imTypeVar.copy();
                 // One source parameter becomes several nodes here. Recorded so the two can be
                 // recognised as the same parameter without falling back to comparing names.
-                translator.recordCopiedTypeVar(copy, imTypeVar);
+                translator.recordSpecialisation(copy, imTypeVar);
                 newTypeVars.add(copy);
             }
             f.getTypeVariables().addAll(0, newTypeVars);
@@ -1136,6 +1136,7 @@ private ImFunction specializeFunction(ImFunction f, GenericTypes generics) {
         prog.getFunctions().add(newF);
 
         // concrete clone => no type vars
+        translator.recordSpecialisation(newF, f, generics.getTypeArguments());
         recordCopiedTypeVars(f.getTypeVariables(), newF.getTypeVariables());
         newF.getTypeVariables().removeAll();
 
@@ -1223,6 +1224,7 @@ private ImFunction specializeMethodImplementation(ImMethod method, GenericTypes
         specializedFunctions.put(implementation, generics, newImplementation);
         specializedFunctionGenerics.put(newImplementation, generics);
         prog.getFunctions().add(newImplementation);
+        translator.recordSpecialisation(newImplementation, implementation, generics.getTypeArguments());
         recordCopiedTypeVars(implementation.getTypeVariables(), newImplementation.getTypeVariables());
         newImplementation.getTypeVariables().removeAll();
         newImplementation.setName(implementation.getName() + "_specialized");
@@ -1416,14 +1418,14 @@ private static String enclosingFunctionName(Element e) {
      */
     private void recordCopiedTypeVars(List<ImTypeVar> originals, List<ImTypeVar> copies) {
         for (int i = 0; i < originals.size() && i < copies.size(); i++) {
-            translator.recordCopiedTypeVar(copies.get(i), originals.get(i));
+            translator.recordSpecialisation(copies.get(i), originals.get(i));
         }
     }
 
     private int indexOfTypeVar(List<ImTypeVar> typeVars, ImTypeVar target) {
-        ImTypeVar wanted = translator.canonicalTypeVar(target);
+        ImTypeVar wanted = translator.canonical(target);
         for (int i = 0; i < typeVars.size(); i++) {
-            if (translator.canonicalTypeVar(typeVars.get(i)) == wanted) {
+            if (translator.canonical(typeVars.get(i)) == wanted) {
                 return i;
             }
         }
@@ -1500,8 +1502,9 @@ private ImClass specializeClass(ImClass c, GenericTypes generics) {
         // The copy is structural, so field i of the copy is field i of the original. Nothing will
         // refer to the copies, so this is the only record that they are the same fields.
         for (int i = 0; i < c.getFields().size() && i < newC.getFields().size(); i++) {
-            translator.recordSpecializedField(newC.getFields().get(i), c.getFields().get(i));
+            translator.recordSpecialisation(newC.getFields().get(i), c.getFields().get(i), generics.getTypeArguments());
         }
+        translator.recordSpecialisation(newC, c, generics.getTypeArguments());
         specializedClasses.put(c, generics, newC);
         prog.getClasses().add(newC);
         recordCopiedTypeVars(c.getTypeVariables(), newC.getTypeVariables());

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

@@ -43,53 +43,61 @@ public class ImTranslator {
     public static final String $DEBUG_PRINT = "$debugPrint";
 
     /**
-     * The field each field of a specialised class was copied from.
+     * What each specialised node was copied from, and under which type arguments.
      * <p>
-     * Nothing refers to a copy: an access made before specialisation still names the original's
-     * variable. A pass which drops fields nothing reads would drop every copy, leaving an instance of
-     * the specialised class allocated with no fields while the emitted code goes on reading them. A
-     * copy is live exactly when the field it was made from is.
+     * Specialising a generic entity makes a new node rather than recording a relation, so the copy has
+     * no way to say what it stands for. Passes then recover that by other means, and each of those
+     * means has been wrong: a pass which drops what nothing reads dropped every copied field, because
+     * an access made before specialisation still names the original's variable; a lookup which matched
+     * type variables by name took two parameters sharing a name for one, which dispatched a value
+     * through the wrong instance; and a name composed from a copy's mangled name reads the type
+     * argument as though it were a method name.
+     * <p>
+     * This is that relation, in one place: a copy, what it was copied from, and the type arguments the
+     * copy was made for. Identity, liveness and naming can all be answered from it rather than from a
+     * name, which is what those three passes were doing by hand.
      */
-    private final java.util.Map<ImVar, ImVar> specializedFieldOrigins = new java.util.IdentityHashMap<>();
+    public record Specialisation(Element original, List<ImTypeArgument> typeArguments) {
+    }
+
+    private final Map<Element, Specialisation> specialisations = new IdentityHashMap<>();
 
     /**
-     * The type variable each copy was made from.
-     * <p>
-     * Moving a function out of its class copies the class's type variables onto it, deliberately, so
-     * one source parameter is several nodes and identity alone cannot recognise them. Matching on the
-     * name instead makes two parameters which merely share a name look like one, which is a wrong
-     * dispatch rather than a missed one; following the copy back to what it was made from tells them
-     * apart.
+     * @param typeArguments the arguments the copy was made for, empty when a copy carries none of its
+     *                      own - moving a function out of its class copies the class's type variables
+     *                      onto it without specialising anything
      */
-    private final java.util.Map<ImTypeVar, ImTypeVar> typeVarOrigins = new java.util.IdentityHashMap<>();
-
-    public void recordCopiedTypeVar(ImTypeVar copy, ImTypeVar original) {
-        typeVarOrigins.put(copy, original);
+    public void recordSpecialisation(Element copy, Element original, List<ImTypeArgument> typeArguments) {
+        specialisations.put(copy, new Specialisation(original, List.copyOf(typeArguments)));
     }
 
-    /** The type variable {@code tv} was ultimately copied from, or {@code tv} itself. */
-    public ImTypeVar canonicalTypeVar(ImTypeVar tv) {
-        ImTypeVar current = tv;
-        // A copy of a copy is possible, so follow to the root; the map is acyclic by construction
-        // because a copy is always newer than what it was made from.
-        for (int steps = 0; steps < 100; steps++) {
-            ImTypeVar origin = typeVarOrigins.get(current);
-            if (origin == null) {
-                return current;
-            }
-            current = origin;
-        }
-        return current;
+    public void recordSpecialisation(Element copy, Element original) {
+        recordSpecialisation(copy, original, List.of());
     }
 
-    public void recordSpecializedField(ImVar copy, ImVar original) {
-        specializedFieldOrigins.put(copy, original);
+    /** What {@code copy} was made from and for, or null when it is not a copy. */
+    public @Nullable Specialisation specialisationOf(Element copy) {
+        return specialisations.get(copy);
     }
 
-    /** The field {@code copy} was specialised from, or {@code copy} itself if it is not a copy. */
-    public ImVar originalOfSpecializedField(ImVar copy) {
-        ImVar origin = specializedFieldOrigins.get(copy);
-        return origin == null ? copy : origin;
+    /**
+     * The node {@code copy} was ultimately copied from, or {@code copy} itself.
+     * <p>
+     * A copy of a copy is possible, so this follows to the root. The relation is acyclic by
+     * construction because a copy is always newer than what it was made from; the bound is there so a
+     * mistake elsewhere fails loudly rather than hanging.
+     */
+    @SuppressWarnings("unchecked")
+    public <T extends Element> T canonical(T copy) {
+        Element current = copy;
+        for (int steps = 0; steps < 1000; steps++) {
+            Specialisation specialisation = specialisations.get(current);
+            if (specialisation == null) {
+                return (T) current;
+            }
+            current = specialisation.original();
+        }
+        throw new IllegalStateException("specialisation chain does not terminate at " + copy);
     }
 
     private static final de.peeeq.wurstscript.ast.Element emptyTrace = Ast.NoExpr();

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

@@ -422,7 +422,7 @@ private Map<ImVar, Set<String>> collectNamesEachFieldMustAvoid() {
             collectMethodNames(c, reserved, new HashSet<>());
             for (ImVar field : c.getFields()) {
                 namesToAvoid
-                    .computeIfAbsent(imTr.originalOfSpecializedField(field), origin -> new HashSet<>())
+                    .computeIfAbsent(imTr.canonical(field), origin -> new HashSet<>())
                     .addAll(reserved);
             }
         }
@@ -445,7 +445,7 @@ private void normalizeFieldNames(ImClass c, Set<ImClass> processed,
         collectMethodNames(c, reserved, new HashSet<>());
         collectSuperFieldNames(c, reserved, new HashSet<>());
         for (ImVar field : c.getFields()) {
-            ImVar origin = imTr.originalOfSpecializedField(field);
+            ImVar origin = imTr.canonical(field);
             String settled = chosenNames.get(origin);
             if (settled != null) {
                 // The original and its copies are one key, decided the first time any of them is met.

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

@@ -93,7 +93,7 @@ public static void removeGarbage(ImProg prog, ImTranslator translator) {
         // no fields at all while the emitted code goes on reading them.
         for (ImClass c : prog.getClasses()) {
             c.getFields().removeIf(g -> !used.getVars().contains(g)
-                && !used.getVars().contains(translator.originalOfSpecializedField(g)));
+                && !used.getVars().contains(translator.canonical(g)));
             c.getFunctions().removeIf(f -> !used.getFunctions().contains(f));
             c.getMethods().removeIf(m -> !used.getMethods().contains(m));
             for (ImMethod m : c.getMethods()) {

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

@@ -0,0 +1,114 @@
+package tests.wurstscript.tests;
+
+import de.peeeq.wurstscript.jassIm.ImTypeArgument;
+import de.peeeq.wurstscript.jassIm.ImVar;
+import de.peeeq.wurstscript.jassIm.JassIm;
+import de.peeeq.wurstscript.translation.imtranslation.ImTranslator;
+import org.testng.annotations.Test;
+
+import java.util.Collections;
+import java.util.List;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.assertThrows;
+
+/**
+ * The specialisation relation on its own.
+ * <p>
+ * Three passes answer questions from it that they previously answered from names, and each of those
+ * name-based answers was wrong in a way that reached a user: fields dropped as dead, a value
+ * dispatched through the wrong instance, a slot named after a type argument. So the relation is worth
+ * covering directly rather than only through the passes which consume it.
+ */
+public class SpecialisationOriginTest {
+
+    private static ImVar var(String name) {
+        return JassIm.ImVar(de.peeeq.wurstscript.ast.Ast.NoExpr(), JassIm.ImSimpleType("integer"), name, false);
+    }
+
+    private static ImTranslator translator() {
+        // The relation is a plain side table on the translator and touches nothing else in it, so an
+        // instance with no program is enough and keeps this a unit test.
+        return new ImTranslator(de.peeeq.wurstscript.ast.Ast.WurstModel(), false, null);
+    }
+
+    @Test
+    public void aNodeWhichIsNotACopyIsItsOwnOriginal() {
+        ImTranslator translator = translator();
+        ImVar original = var("count");
+
+        assertSame(translator.canonical(original), original);
+        assertNull(translator.specialisationOf(original));
+    }
+
+    @Test
+    public void aCopyLeadsBackToWhatItWasMadeFrom() {
+        ImTranslator translator = translator();
+        ImVar original = var("count");
+        ImVar copy = var("count");
+        translator.recordSpecialisation(copy, original);
+
+        assertSame(translator.canonical(copy), original);
+    }
+
+    /** Specialising a specialisation happens, so the relation has to be followed to its root. */
+    @Test
+    public void aCopyOfACopyLeadsBackToTheRoot() {
+        ImTranslator translator = translator();
+        ImVar original = var("count");
+        ImVar once = var("count");
+        ImVar twice = var("count");
+        translator.recordSpecialisation(once, original);
+        translator.recordSpecialisation(twice, once);
+
+        assertSame(translator.canonical(twice), original);
+        assertSame(translator.canonical(once), original);
+    }
+
+    /**
+     * Two nodes sharing a name are not the same parameter. This is the case the name matching this
+     * relation replaced got wrong, and it crashed the interpreter rather than answering differently.
+     */
+    @Test
+    public void sharingANameIsNotSharingAnOrigin() {
+        ImTranslator translator = translator();
+        ImVar first = var("T");
+        ImVar second = var("T");
+
+        assertEquals(first.getName(), second.getName());
+        assertSame(translator.canonical(first), first);
+        assertSame(translator.canonical(second), second);
+    }
+
+    /** The arguments a copy was made for are kept, which is what a name can be composed from. */
+    @Test
+    public void theTypeArgumentsOfTheCopyAreKept() {
+        ImTranslator translator = translator();
+        ImVar original = var("keys");
+        ImVar copy = var("keys");
+        List<ImTypeArgument> arguments = List.of(
+            JassIm.ImTypeArgument(JassIm.ImSimpleType("integer"), Collections.emptyMap()));
+        translator.recordSpecialisation(copy, original, arguments);
+
+        ImTranslator.Specialisation specialisation = translator.specialisationOf(copy);
+        assertSame(specialisation.original(), original);
+        assertEquals(specialisation.typeArguments().size(), 1);
+    }
+
+    /**
+     * A cycle cannot arise from specialising - a copy is always newer than what it was made from - so
+     * one means a mistake elsewhere, and the relation says so rather than looping.
+     */
+    @Test
+    public void acycleIsReportedRatherThanFollowedForever() {
+        ImTranslator translator = translator();
+        ImVar a = var("a");
+        ImVar b = var("b");
+        translator.recordSpecialisation(a, b);
+        translator.recordSpecialisation(b, a);
+
+        assertThrows(IllegalStateException.class, () -> translator.canonical(a));
+    }
+}