Lua hot reload: callbacks told apart by their order reload, @onHotReload hooks, smaller payloads

de.peeeq.wurstscript/src/main/java/de/peeeq/wurstio/languageserver/requests/LuaHotReloadSession.java

@@ -44,6 +44,8 @@ private static final class State {
         List<String> packageInits = new ArrayList<>();
         /** The initializers of globals a version the map ran declared as variables, which code may have written. */
         List<String> variables = new ArrayList<>();
+        /** Every shape sent of each closure told apart by its order. */
+        Map<String, List<String>> groupShapes = new LinkedHashMap<>();
     }
 
     /**
@@ -150,6 +152,7 @@ private static final class Pending {
         Map<String, String> initializers = new LinkedHashMap<>();
         List<String> packageInits = new ArrayList<>();
         List<String> variables = new ArrayList<>();
+        Map<String, String> groupShapes = new LinkedHashMap<>();
     }
 
     /** A payload and the file it was written to. */
@@ -172,7 +175,7 @@ public static Sent send(Path buildDir, Path customMapData, LuaHotReload.Manifest
         LuaHotReload.Payload payload = LuaHotReload.payload(state.baseline, current, new LuaHotReload.History(
             new LinkedHashSet<>(state.reloaded), orEmpty(state.capturedDefinitions), orEmpty(state.slots),
             orEmpty(state.initializers), new LinkedHashSet<>(orEmpty(state.packageInits)),
-            new LinkedHashSet<>(orEmpty(state.variables))), state.version + 1);
+            new LinkedHashSet<>(orEmpty(state.variables)), orEmpty(state.groupShapes)), state.version + 1);
         Pending pending = new Pending();
         pending.version = payload.version();
         pending.functions = payload.functions();
@@ -181,6 +184,7 @@ public static Sent send(Path buildDir, Path customMapData, LuaHotReload.Manifest
         pending.initializers = new LinkedHashMap<>(payload.initializers());
         pending.packageInits = new ArrayList<>(payload.packageInits());
         pending.variables = new ArrayList<>(payload.variables());
+        pending.groupShapes = new LinkedHashMap<>(payload.groupShapes());
         writeAtomically(pendingFile(buildDir), GSON.toJson(pending));
         Path delivered = deliver(customMapData, payload);
         record(state, pending);
@@ -232,6 +236,14 @@ private static void record(State state, Pending sent) {
         Set<String> variables = new LinkedHashSet<>(orEmpty(state.variables));
         variables.addAll(orEmpty(sent.variables));
         state.variables = new ArrayList<>(variables);
+        // The map runs one of the shapes sent, depending on which payloads it installed: keep them all.
+        Map<String, List<String>> groupShapes = new LinkedHashMap<>(orEmpty(state.groupShapes));
+        for (Map.Entry<String, String> entry : orEmpty(sent.groupShapes).entrySet()) {
+            Set<String> shapes = new LinkedHashSet<>(groupShapes.getOrDefault(entry.getKey(), List.of()));
+            shapes.add(entry.getValue());
+            groupShapes.put(entry.getKey(), List.copyOf(shapes));
+        }
+        state.groupShapes = groupShapes;
         state.version = sent.version;
     }
 

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

@@ -0,0 +1,36 @@
+package de.peeeq.wurstscript;
+
+import de.peeeq.wurstscript.ast.Element;
+import de.peeeq.wurstscript.ast.FuncDef;
+import de.peeeq.wurstscript.ast.StructureDef;
+import de.peeeq.wurstscript.attributes.GenericsHelper;
+
+/**
+ * Functions annotated {@code @onHotReload}: in a {@code -luaHotReload} build each payload calls them, with their bodies
+ * as the payload has them, once it installed everything else, so a map can run its setup code again
+ * (docs/LUA_HOT_RELOAD.md, "Reload hooks"). In any other build the annotation does nothing.
+ */
+public final class HotReloadHooks {
+    public static final String ANNOTATION = "onHotReload";
+
+    private HotReloadHooks() {
+    }
+
+    public static boolean isHook(Element trace) {
+        return trace instanceof FuncDef f && f.attrHasAnnotation(ANNOTATION);
+    }
+
+    /** Whether a function, or a class or module it is declared in, has type parameters, which a payload cannot pass. */
+    public static boolean isGeneric(FuncDef f) {
+        if (!f.getTypeParameters().isEmpty()) {
+            return true;
+        }
+        for (StructureDef s = f.attrNearestStructureDef(); s != null;
+             s = s.getParent() == null ? null : s.getParent().attrNearestStructureDef()) {
+            if (!GenericsHelper.typeParameters(s).isEmpty()) {
+                return true;
+            }
+        }
+        return false;
+    }
+}

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

@@ -140,6 +140,9 @@ private static int remove(ImTranslator trans, boolean followTypeArguments, boole
         for (ImFunction pinned : trans.pinnedFunctions()) {
             reachability.add(pinned);
         }
+        for (ImFunction hook : trans.hotReloadHooks()) {
+            reachability.add(hook);
+        }
         for (ImFunction root : extraRoots) {
             reachability.add(root);
         }

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

@@ -5,6 +5,7 @@
 import com.google.common.collect.ImmutableList.Builder;
 import de.peeeq.datastructures.Partitions;
 import de.peeeq.datastructures.TransitiveClosure;
+import de.peeeq.wurstscript.HotReloadHooks;
 import de.peeeq.wurstscript.RunArgs;
 import de.peeeq.wurstscript.WLogger;
 import de.peeeq.wurstscript.WurstOperator;
@@ -403,6 +404,24 @@ public List<ImFunction> pinnedFunctions() {
         return result;
     }
 
+    /**
+     * In a {@code -luaHotReload} build, the functions annotated {@code @onHotReload}: each payload calls them, so they are
+     * entry points as preserved functions are, even when no Wurst code calls them. Empty in any other build.
+     */
+    public List<ImFunction> hotReloadHooks() {
+        if (!runArgs.isLuaHotReload()) {
+            return List.of();
+        }
+        List<ImFunction> hooks = new ArrayList<>();
+        for (ImFunction f : ImHelper.calculateFunctionsOfProg(imProg)) {
+            // The function made for the annotated declaration, not another with its trace (a stack-trace bridge).
+            if (HotReloadHooks.isHook(f.getTrace()) && functionMap.get(f.getTrace()) == f) {
+                hooks.add(f);
+            }
+        }
+        return hooks;
+    }
+
     /**
      * The globals this translator created for passes which run later and reach them through its handle, not through
      * an access in a body: the variables which manage the instances of a class, which the class elimination builds
@@ -2225,6 +2244,9 @@ private void calculateCallRelationsAndVariables(boolean includeUsedVariables) {
                 calculateCallRelations(function, includeUsedVariables);
             }
         }
+        for (ImFunction hook : hotReloadHooks()) {
+            calculateCallRelations(hook, includeUsedVariables);
+        }
 
         // Mark externally visible globals as read so they survive garbage collection.
         for (ImVar global : imProg.getGlobals()) {