RegisterWidget annotation!

refactor some init stuff to have less duplicate code

9 files changed, 323 insertions, 72 deletions

diff --git a/buildSrc/build.gradle b/buildSrc/build.gradle
index e6cf6a1f..c6398c0a 100644
--- a/buildSrc/build.gradle
+++ b/buildSrc/build.gradle
@@ -16,7 +16,7 @@ gradlePlugin {
 		simplePlugin {
 			id = 'de.hysky.skyblocker.annotation-processor'
 			// The plugin entry point could be changed to a different class that then appropriately calls the different processors when there's more than one.
-			implementationClass = 'de.hysky.skyblocker.init.InitProcessor'
+			implementationClass = 'de.hysky.skyblocker.Processor'
 		}
 	}
 }
diff --git a/buildSrc/src/main/java/de/hysky/skyblocker/Processor.java b/buildSrc/src/main/java/de/hysky/skyblocker/Processor.java
new file mode 100644
index 00000000..24560ae3
--- /dev/null
+++ b/buildSrc/src/main/java/de/hysky/skyblocker/Processor.java
@@ -0,0 +1,115 @@
+package de.hysky.skyblocker;
+
+import de.hysky.skyblocker.hud.HudProcessor;
+import de.hysky.skyblocker.init.InitProcessor;
+import org.gradle.api.Plugin;
+import org.gradle.api.Project;
+import org.gradle.api.logging.Logger;
+import org.gradle.api.logging.Logging;
+import org.gradle.api.tasks.compile.JavaCompile;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.tree.ClassNode;
+
+import java.io.File;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.file.FileVisitResult;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.SimpleFileVisitor;
+import java.nio.file.attribute.BasicFileAttributes;
+import java.util.Objects;
+import java.util.function.Consumer;
+
+public class Processor implements Plugin<Project> {
+
+	public static final Logger logger = Logging.getLogger(Processor.class);
+	public static File classesDir;
+
+	@Override
+	public void apply(@NotNull Project project) {
+		// https://docs.gradle.org/current/userguide/task_configuration_avoidance.html
+		// This only configures the `compileJava` task and not other `JavaCompile` tasks such as `compileTestJava`. https://stackoverflow.com/a/77047012
+		project.getTasks().withType(JavaCompile.class).named("compileJava").get().doLast(task -> {
+			JavaCompile javaCompile = (JavaCompile) task;
+			classesDir = javaCompile.getDestinationDirectory().get().getAsFile();
+
+			new InitProcessor().apply(javaCompile);
+			new HudProcessor().apply(javaCompile);
+		});
+	}
+
+	public static void forEachClass(@NotNull File directory, final Consumer<InputStream> consumer) {
+		try {
+			Files.walkFileTree(directory.toPath(), new SimpleFileVisitor<>() {
+				@Override
+				public FileVisitResult visitFile(Path path, BasicFileAttributes attrs) {
+					if (!path.toString().endsWith(".class")) return FileVisitResult.CONTINUE;
+					try (InputStream inputStream = Files.newInputStream(path)) {
+						consumer.accept(inputStream);
+					} catch (IOException e) {
+						logger.error("Failed to run consumer on class {}", path, e);
+
+					}
+
+					return FileVisitResult.CONTINUE;
+				}
+			});
+		} catch (IOException e) {
+			logger.error("Failed to walk classes", e);
+		}
+	}
+
+	public static void forEachClass(final Consumer<InputStream> consumer) {
+		forEachClass(classesDir, consumer);
+	}
+
+	public static @Nullable File findClass(File directory, String className) {
+		if (!className.endsWith(".class")) className += ".class";
+
+		if (!directory.isDirectory()) throw new IllegalArgumentException("Not a directory");
+
+		for (File file : Objects.requireNonNull(directory.listFiles())) {
+			if (file.isDirectory()) {
+				File foundFile = findClass(file, className);
+
+				if (foundFile != null) return foundFile;
+			} else if (file.getName().equals(className)) {
+				return file;
+			}
+		}
+		return null;
+	}
+
+	public static @Nullable File findClass(String className) {
+		return findClass(classesDir, className);
+	}
+
+	/**
+	 * Pretty much [child instanceof superClass]
+	 * <p>
+	 * Classes are full name. Example: de/hysky/skyblocker/SkyblockerMod
+	 * @param child the class to test
+	 * @param superClass super
+	 * @return if child is an instance of superclass
+	 */
+	public static boolean instanceOf(String child, String superClass) {
+		Path start = classesDir.toPath();
+		String sup = child;
+		while (sup != null) {
+			if (sup.equals(superClass)) return true;
+			Path resolve = start.resolve(sup + ".class");
+			try (InputStream stream = Files.newInputStream(resolve)) {
+				ClassReader classReader = new ClassReader(stream);
+				sup = classReader.getSuperName();
+			} catch (IOException e) {
+				logger.error("Failed to read class {}", resolve, e);
+				return false;
+			}
+		}
+		return false;
+	}
+}
diff --git a/buildSrc/src/main/java/de/hysky/skyblocker/hud/HudInjectClassVisitor.java b/buildSrc/src/main/java/de/hysky/skyblocker/hud/HudInjectClassVisitor.java
new file mode 100644
index 00000000..8acd06f0
--- /dev/null
+++ b/buildSrc/src/main/java/de/hysky/skyblocker/hud/HudInjectClassVisitor.java
@@ -0,0 +1,45 @@
+package de.hysky.skyblocker.hud;
+
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.tree.ClassNode;
+import org.objectweb.asm.tree.MethodNode;
+
+import java.util.List;
+
+public class HudInjectClassVisitor extends ClassVisitor {
+
+	private final List<ClassNode> widgetClasses;
+
+	protected HudInjectClassVisitor(ClassVisitor delegate, List<ClassNode> widgetClasses) {
+		super(Opcodes.ASM9, delegate);
+		this.widgetClasses = widgetClasses;
+	}
+
+
+	@Override
+	public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
+		MethodVisitor methodVisitor = super.visitMethod(access, name, descriptor, signature, exceptions);
+
+		if ((access & Opcodes.ACC_PRIVATE) != 0 && (access & Opcodes.ACC_STATIC) != 0 && name.equals("instantiateWidgets") && descriptor.equals("()V")) {
+			MethodNode methodNode = new MethodNode(Opcodes.ASM9, access, name, descriptor, signature, exceptions);
+
+			for (ClassNode widget : widgetClasses) {
+				methodNode.visitTypeInsn(Opcodes.NEW, widget.name);
+				methodNode.visitInsn(Opcodes.DUP);
+				methodNode.visitMethodInsn(Opcodes.INVOKESPECIAL, widget.name, "<init>", "()V", false);
+				methodNode.visitMethodInsn(Opcodes.INVOKESTATIC, "de/hysky/skyblocker/skyblock/tabhud/screenbuilder/ScreenMaster", "addWidgetInstance", "(Lde/hysky/skyblocker/skyblock/tabhud/widget/HudWidget;)V", false);
+			}
+
+
+			// Return from the method
+			methodNode.visitInsn(Opcodes.RETURN);
+
+			// Apply our new method node to the visitor to replace the original one
+			methodNode.accept(methodVisitor);
+		}
+
+		return methodVisitor;
+	}
+}
diff --git a/buildSrc/src/main/java/de/hysky/skyblocker/hud/HudProcessor.java b/buildSrc/src/main/java/de/hysky/skyblocker/hud/HudProcessor.java
new file mode 100644
index 00000000..34c9a5d2
--- /dev/null
+++ b/buildSrc/src/main/java/de/hysky/skyblocker/hud/HudProcessor.java
@@ -0,0 +1,100 @@
+package de.hysky.skyblocker.hud;
+
+import de.hysky.skyblocker.Processor;
+import org.gradle.api.tasks.compile.JavaCompile;
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassWriter;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.tree.AnnotationNode;
+import org.objectweb.asm.tree.ClassNode;
+import org.objectweb.asm.tree.MethodNode;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.*;
+
+public class HudProcessor {
+
+	private final Map<ClassNode, Integer> annotatedClassesConstructors = new HashMap<>();
+
+	public void apply(JavaCompile task) {
+		long start = System.currentTimeMillis();
+
+		Processor.forEachClass(this::visitClass);
+
+		List<ClassNode> constructors = new ArrayList<>(annotatedClassesConstructors.keySet());
+		constructors.sort(Comparator.comparingInt(annotatedClassesConstructors::get));
+
+		inject(constructors);
+
+		System.out.println("Injecting widget instancing took: " + (System.currentTimeMillis() - start) + "ms");
+	}
+
+	private void visitClass(InputStream inputStream) {
+		try {
+			ClassReader classReader = new ClassReader(inputStream);
+			ClassNode classNode = new ClassNode(Opcodes.ASM9);
+			classReader.accept(classNode, 0);
+
+			// Look for the annotation
+			boolean annotationFound = false;
+			int priority = 0;
+			List<AnnotationNode> annotationNodes = new ArrayList<>(classNode.visibleAnnotations == null ? List.of() : classNode.visibleAnnotations);
+			annotationNodes.addAll(classNode.invisibleAnnotations == null ? List.of() : classNode.invisibleAnnotations);
+			for (AnnotationNode annotationNode : annotationNodes) {
+				String desc = annotationNode.desc;
+				if (!desc.equals("Lde/hysky/skyblocker/annotations/RegisterWidget;")) continue;
+
+				annotationFound = true;
+				// null if no parameters are given, defaults don't show up :shrug:
+				if (annotationNode.values != null) {
+					for (int i = 0; i < annotationNode.values.size(); i++) {
+						if ("priority".equals(annotationNode.values.get(i))) {
+							priority = (int) annotationNode.values.get(i + 1);
+						}
+					}
+				}
+				break;
+			}
+			if (!annotationFound) return;
+			if (!Processor.instanceOf(classNode.name, "de/hysky/skyblocker/skyblock/tabhud/widget/HudWidget")) {
+				throw new IllegalArgumentException("Class " + classNode.name + " has @RegisterWidget annotation but does not extend HudWidget");
+			}
+
+			// Look for constructor
+			MethodNode constructor = null;
+			for (MethodNode method : classNode.methods) {
+				if (!method.name.equals("<init>")) continue;
+				if (!method.desc.equals("()V")) continue;
+				constructor = method;
+				break;
+			}
+			if (constructor == null) throw new IllegalStateException("No parameterless constructor found for " + classNode.name);
+
+			annotatedClassesConstructors.put(classNode, priority);
+
+
+
+		} catch (IOException e) {
+			throw new RuntimeException(e);
+		}
+	}
+
+	private void inject(List<ClassNode> constructors) {
+		Path mainClassFile = Objects.requireNonNull(Processor.findClass("ScreenMaster.class"), "ScreenMaster class wasn't found :(").toPath();
+
+		try (InputStream inputStream = Files.newInputStream(mainClassFile)) {
+			ClassReader classReader = new ClassReader(inputStream);
+			ClassWriter classWriter = new ClassWriter(ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
+			classReader.accept(new HudInjectClassVisitor(classWriter, constructors), 0);
+			try (OutputStream outputStream = Files.newOutputStream(mainClassFile)) {
+				outputStream.write(classWriter.toByteArray());
+			}
+		} catch (Exception e) {
+			throw new RuntimeException(e);
+		}
+	}
+}
diff --git a/buildSrc/src/main/java/de/hysky/skyblocker/init/InitProcessor.java b/buildSrc/src/main/java/de/hysky/skyblocker/init/InitProcessor.java
index b3c353a4..3c864bde 100644
--- a/buildSrc/src/main/java/de/hysky/skyblocker/init/InitProcessor.java
+++ b/buildSrc/src/main/java/de/hysky/skyblocker/init/InitProcessor.java
@@ -1,77 +1,54 @@
 package de.hysky.skyblocker.init;
 
-import org.gradle.api.Plugin;
-import org.gradle.api.Project;
+import de.hysky.skyblocker.Processor;
 import org.gradle.api.tasks.compile.JavaCompile;
-import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
 import org.objectweb.asm.*;
 
-import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
-import java.nio.file.FileVisitResult;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.nio.file.SimpleFileVisitor;
-import java.nio.file.attribute.BasicFileAttributes;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 
 @SuppressWarnings("unused")
-public abstract class InitProcessor implements Plugin<Project> {
-	@Override
-	public void apply(Project project) {
-		// https://docs.gradle.org/current/userguide/task_configuration_avoidance.html
-		// This only configures the `compileJava` task and not other `JavaCompile` tasks such as `compileTestJava`. https://stackoverflow.com/a/77047012
-        project.getTasks().withType(JavaCompile.class).named("compileJava").get().doLast(task -> {
-			long start = System.currentTimeMillis();
-			File classesDir = ((JavaCompile) task).getDestinationDirectory().get().getAsFile();
-			Map<MethodReference, Integer> methodSignatures = new HashMap<>();
+public class InitProcessor {
+	public void apply(JavaCompile task) {
+		long start = System.currentTimeMillis();
+		Map<MethodReference, Integer> methodSignatures = new HashMap<>();
 
-			//Find all methods with the @Init annotation
-			findInitMethods(classesDir, methodSignatures);
+		//Find all methods with the @Init annotation
+		findInitMethods(methodSignatures);
 
-			//Sort the methods by their priority. It's also converted to a list because the priority values are useless from here on
-			List<MethodReference> sortedMethodSignatures = methodSignatures.entrySet()
-					.stream()
-					.sorted(Map.Entry.<MethodReference, Integer>comparingByValue().thenComparing(entry -> entry.getKey().className()))
-					.map(Map.Entry::getKey)
-					.toList();
+		//Sort the methods by their priority. It's also converted to a list because the priority values are useless from here on
+		List<MethodReference> sortedMethodSignatures = methodSignatures.entrySet()
+				.stream()
+				.sorted(Map.Entry.<MethodReference, Integer>comparingByValue().thenComparing(entry -> entry.getKey().className()))
+				.map(Map.Entry::getKey)
+				.toList();
 
-			//Inject calls to the @Init annotated methods in the SkyblockerMod class
-			injectInitCalls(classesDir, sortedMethodSignatures);
+		//Inject calls to the @Init annotated methods in the SkyblockerMod class
+		injectInitCalls(sortedMethodSignatures);
 
-			System.out.println("Injecting init methods took: " + (System.currentTimeMillis() - start) + "ms");
-		});
+		System.out.println("Injecting init methods took: " + (System.currentTimeMillis() - start) + "ms");
 	}
 
-	public void findInitMethods(@NotNull File directory, Map<MethodReference, Integer> methodSignatures) {
-		try {
-			Files.walkFileTree(directory.toPath(), new SimpleFileVisitor<>() {
-				@Override
-				public FileVisitResult visitFile(Path path, BasicFileAttributes attrs) {
-                    if (!path.toString().endsWith(".class")) return FileVisitResult.CONTINUE;
-					try (InputStream inputStream = Files.newInputStream(path)) {
-						ClassReader classReader = new ClassReader(inputStream);
-						classReader.accept(new InitReadingClassVisitor(classReader, methodSignatures), ClassReader.SKIP_CODE | ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);
-					} catch (IOException e) {
-						e.printStackTrace();
-					}
-
-					return FileVisitResult.CONTINUE;
-				}
-			});
-		} catch (IOException e) {
-			e.printStackTrace();
-		}
+	public void findInitMethods(Map<MethodReference, Integer> methodSignatures) {
+		Processor.forEachClass(inputStream -> {
+			try {
+				ClassReader classReader = new ClassReader(inputStream);
+				classReader.accept(new InitReadingClassVisitor(classReader, methodSignatures), ClassReader.SKIP_CODE | ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);
+			} catch (IOException e) {
+				throw new RuntimeException(e);
+			}
+		});
 	}
 
-	public void injectInitCalls(File directory, List<MethodReference> methodSignatures) {
-		Path mainClassFile = Objects.requireNonNull(findMainClass(directory), "SkyblockerMod class wasn't found :(").toPath();
+	public void injectInitCalls(List<MethodReference> methodSignatures) {
+		Path mainClassFile = Objects.requireNonNull(Processor.findClass("SkyblockerMod.class"), "SkyblockerMod class wasn't found :(").toPath();
 
 		try (InputStream inputStream = Files.newInputStream(mainClassFile)) {
 			ClassReader classReader = new ClassReader(inputStream);
@@ -85,23 +62,6 @@ public abstract class InitProcessor implements Plugin<Project> {
 		}
 	}
 
-	// Find the main SkyblockerMod class
-	@Nullable
-	public File findMainClass(@NotNull File directory) {
-		if (!directory.isDirectory()) throw new IllegalArgumentException("Not a directory");
-		for (File file : Objects.requireNonNull(directory.listFiles())) {
-			if (file.isDirectory()) {
-				File foundFile = findMainClass(file);
-
-				if (foundFile != null) return foundFile;
-			} else if (file.getName().equals("SkyblockerMod.class")) {
-				return file;
-			}
-		}
-
-		return null;
-	}
-
 	/**
 	 * @param className  the class name (e.g. de/hysky/skyblocker/skyblock/ChestValue)
 	 * @param methodName the method's name (e.g. init)
diff --git a/src/main/java/de/hysky/skyblocker/annotations/RegisterWidget.java b/src/main/java/de/hysky/skyblocker/annotations/RegisterWidget.java
new file mode 100644
index 00000000..219b0c9b
--- /dev/null
+++ b/src/main/java/de/hysky/skyblocker/annotations/RegisterWidget.java
@@ -0,0 +1,15 @@
+package de.hysky.skyblocker.annotations;
+
+import java.lang.annotation.*;
+
+@Documented
+@Retention(RetentionPolicy.CLASS)
+@Target(ElementType.TYPE)
+public @interface RegisterWidget {
+	/**
+	 * The priority of the widget.
+	 * The higher the number, the later the widget will be instantiated.
+	 * Use this to ensure that your widget is instantiated after widget method if it depends on it.
+	 */
+	int priority() default 0;
+}
diff --git a/src/main/java/de/hysky/skyblocker/skyblock/end/EndHudWidget.java b/src/main/java/de/hysky/skyblocker/skyblock/end/EndHudWidget.java
index 8e0a4c30..ee1b95c1 100644
--- a/src/main/java/de/hysky/skyblocker/skyblock/end/EndHudWidget.java
+++ b/src/main/java/de/hysky/skyblocker/skyblock/end/EndHudWidget.java
@@ -1,6 +1,7 @@
 package de.hysky.skyblocker.skyblock.end;
 
 import com.mojang.authlib.properties.PropertyMap;
+import de.hysky.skyblocker.annotations.RegisterWidget;
 import de.hysky.skyblocker.config.SkyblockerConfigManager;
 import de.hysky.skyblocker.skyblock.tabhud.widget.ComponentBasedWidget;
 import de.hysky.skyblocker.skyblock.tabhud.widget.component.IcoTextComponent;
@@ -18,10 +19,11 @@ import java.text.NumberFormat;
 import java.util.Locale;
 import java.util.Optional;
 
+@RegisterWidget
 public class EndHudWidget extends ComponentBasedWidget {
     private static final MutableText TITLE = Text.literal("The End").formatted(Formatting.LIGHT_PURPLE, Formatting.BOLD);
 
-    public static final EndHudWidget INSTANCE = new EndHudWidget(TITLE, Formatting.DARK_PURPLE.getColorValue());
+    public static final EndHudWidget INSTANCE = new EndHudWidget();
     private static final NumberFormat DECIMAL_FORMAT = NumberFormat.getInstance(Locale.US);
     private static final ItemStack ENDERMAN_HEAD = new ItemStack(Items.PLAYER_HEAD);
     private static final ItemStack POPPY = new ItemStack(Items.POPPY);
@@ -34,8 +36,8 @@ public class EndHudWidget extends ComponentBasedWidget {
         POPPY.set(DataComponentTypes.ENCHANTMENT_GLINT_OVERRIDE, true);
     }
 
-    public EndHudWidget(MutableText title, Integer colorValue) {
-        super(title, colorValue, "hud_end");
+    public EndHudWidget() {
+        super(TITLE, Formatting.DARK_PURPLE.getColorValue(), "hud_end");
         this.update();
     }
     @Override
diff --git a/src/main/java/de/hysky/skyblocker/skyblock/garden/FarmingHudWidget.java b/src/main/java/de/hysky/skyblocker/skyblock/garden/FarmingHudWidget.java
index 383e2061..f686b0ae 100644
--- a/src/main/java/de/hysky/skyblocker/skyblock/garden/FarmingHudWidget.java
+++ b/src/main/java/de/hysky/skyblocker/skyblock/garden/FarmingHudWidget.java
@@ -1,5 +1,6 @@
 package de.hysky.skyblocker.skyblock.garden;
 
+import de.hysky.skyblocker.annotations.RegisterWidget;
 import de.hysky.skyblocker.config.SkyblockerConfigManager;
 import de.hysky.skyblocker.skyblock.item.tooltip.info.TooltipInfoType;
 import de.hysky.skyblocker.skyblock.itemlist.ItemRepository;
@@ -20,6 +21,7 @@ import net.minecraft.util.math.MathHelper;
 
 import java.util.Map;
 
+@RegisterWidget
 public class FarmingHudWidget extends ComponentBasedWidget {
     private static final MutableText TITLE = Text.literal("Farming").formatted(Formatting.YELLOW, Formatting.BOLD);
     public static final Map<String, String> FARMING_TOOLS = Map.ofEntries(
diff --git a/src/main/java/de/hysky/skyblocker/skyblock/tabhud/screenbuilder/ScreenMaster.java b/src/main/java/de/hysky/skyblocker/skyblock/tabhud/screenbuilder/ScreenMaster.java
index 06207554..bffa5418 100644
--- a/src/main/java/de/hysky/skyblocker/skyblock/tabhud/screenbuilder/ScreenMaster.java
+++ b/src/main/java/de/hysky/skyblocker/skyblock/tabhud/screenbuilder/ScreenMaster.java
@@ -151,10 +151,13 @@ public class ScreenMaster {
 
     }
 
-    @Init
+	// we probably want this to run pretty early?
+    @Init(priority = -1)
     public static void init() {
         SkyblockEvents.LOCATION_CHANGE.register(location -> ScreenBuilder.positionsNeedsUpdating = true);
 
+		instantiateWidgets();
+
         ClientLifecycleEvents.CLIENT_STARTED.register(client -> {
             try {
                 ClassPath.from(TabHudWidget.class.getClassLoader()).getTopLevelClasses("de.hysky.skyblocker.skyblock.tabhud.widget").iterator().forEachRemaining(classInfo -> {
@@ -187,6 +190,15 @@ public class ScreenMaster {
         ClientLifecycleEvents.CLIENT_STOPPING.register(client -> saveConfig());
     }
 
+
+	private static void instantiateWidgets() {}
+
+	@SuppressWarnings("unused")
+	public static void addWidgetInstance(HudWidget widget) {
+		HudWidget put = widgetInstances.put(widget.getInternalID(), widget);
+		if (put != null) LOGGER.warn("[Skyblocker] Duplicate hud widget found: {}", widget);
+	}
+
     /**
      * @implNote !! The 3 first ones shouldn't be moved, ordinal is used in some places
      */