Use ASM compile-time class transformation for class init via an `@Init` annotation (#924)

* Add annotation processor for init methods and @Init annotation

* Use ASM for @Init

* Separate the annotation processor to its own plugin file inside buildSrc

* Actually implement priority

* Reverse annotation equality check and method check to warn about misuse of annotation

* Add gradle.properties to buildSrc and move asm version into it

* Reformat buildscripts

Less conflicting for other PRs

* Refactor to use a record over strings

* Rebase onto master and add more documentation

* Remove rebasing artifact

* Apply suggestions from code review

Simplifies the `itf`

Co-authored-by: Kevin <92656833+kevinthegreat1@users.noreply.github.com>

* Use Files class' methods for reading and writing to files

* Apply suggestion

* Then sort by name

* Clean up InitProcessor

* Separate classes into java files

* Fix indent

---------

Co-authored-by: Aaron <51387595+AzureAaron@users.noreply.github.com>
Co-authored-by: Kevin <92656833+kevinthegreat1@users.noreply.github.com>

3 files changed, 224 insertions, 0 deletions

diff --git a/buildSrc/src/main/java/de/hysky/skyblocker/init/InitInjectingClassVisitor.java b/buildSrc/src/main/java/de/hysky/skyblocker/init/InitInjectingClassVisitor.java
new file mode 100644
index 00000000..e1cde6ca
--- /dev/null
+++ b/buildSrc/src/main/java/de/hysky/skyblocker/init/InitInjectingClassVisitor.java
@@ -0,0 +1,41 @@
+package de.hysky.skyblocker.init;
+
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.tree.MethodNode;
+
+import java.util.List;
+
+public class InitInjectingClassVisitor extends ClassVisitor {
+	private final List<InitProcessor.MethodReference> methodSignatures;
+
+	public InitInjectingClassVisitor(ClassVisitor classVisitor, List<InitProcessor.MethodReference> methodSignatures) {
+		super(Opcodes.ASM9, classVisitor);
+		this.methodSignatures = methodSignatures;
+	}
+
+	@Override
+	public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
+		MethodVisitor methodVisitor = super.visitMethod(access, name, descriptor, signature, exceptions);
+
+		// Limit replacing to the init method which is private, static, named init, has no args, and has a void return type
+		if ((access & Opcodes.ACC_PRIVATE) != 0 && (access & Opcodes.ACC_STATIC) != 0 && name.equals("init") && descriptor.equals("()V")) {
+			// Method node that we will overwrite the init method with
+			MethodNode methodNode = new MethodNode(Opcodes.ASM9, access, name, descriptor, signature, exceptions);
+
+			// Inject calls to each found @Init annotated method
+			for (InitProcessor.MethodReference methodCall : methodSignatures) {
+				methodNode.visitMethodInsn(Opcodes.INVOKESTATIC, methodCall.className(), methodCall.methodName(), methodCall.descriptor(), methodCall.itf());
+			}
+
+			// 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/init/InitProcessor.java b/buildSrc/src/main/java/de/hysky/skyblocker/init/InitProcessor.java
new file mode 100644
index 00000000..b3c353a4
--- /dev/null
+++ b/buildSrc/src/main/java/de/hysky/skyblocker/init/InitProcessor.java
@@ -0,0 +1,112 @@
+package de.hysky.skyblocker.init;
+
+import org.gradle.api.Plugin;
+import org.gradle.api.Project;
+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<>();
+
+			//Find all methods with the @Init annotation
+			findInitMethods(classesDir, 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();
+
+			//Inject calls to the @Init annotated methods in the SkyblockerMod class
+			injectInitCalls(classesDir, sortedMethodSignatures);
+
+			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 injectInitCalls(File directory, List<MethodReference> methodSignatures) {
+		Path mainClassFile = Objects.requireNonNull(findMainClass(directory), "SkyblockerMod 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 InitInjectingClassVisitor(classWriter, methodSignatures), 0);
+			try (OutputStream outputStream = Files.newOutputStream(mainClassFile)) {
+				outputStream.write(classWriter.toByteArray());
+			}
+		} catch (Exception e) {
+			e.printStackTrace();
+		}
+	}
+
+	// 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)
+	 * @param descriptor the method's descriptor (only ()V for now)
+	 * @param itf        whether the target class is an {@code interface} or not
+	 */
+	public record MethodReference(String className, String methodName, String descriptor, boolean itf) {}
+}
diff --git a/buildSrc/src/main/java/de/hysky/skyblocker/init/InitReadingClassVisitor.java b/buildSrc/src/main/java/de/hysky/skyblocker/init/InitReadingClassVisitor.java
new file mode 100644
index 00000000..e5c9cf48
--- /dev/null
+++ b/buildSrc/src/main/java/de/hysky/skyblocker/init/InitReadingClassVisitor.java
@@ -0,0 +1,71 @@
+package de.hysky.skyblocker.init;
+
+import org.jetbrains.annotations.NotNull;
+import org.objectweb.asm.*;
+
+import java.util.Map;
+
+public class InitReadingClassVisitor extends ClassVisitor {
+	private final Map<InitProcessor.MethodReference, Integer> methodSignatures;
+	private final ClassReader classReader;
+
+	public InitReadingClassVisitor(ClassReader classReader, Map<InitProcessor.MethodReference, Integer> methodSignatures) {
+		super(Opcodes.ASM9);
+		this.classReader = classReader;
+		this.methodSignatures = methodSignatures;
+	}
+
+	@Override
+	public MethodVisitor visitMethod(int access, String methodName, String descriptor, String signature, String[] exceptions) {
+		return new MethodVisitor(Opcodes.ASM9) {
+			@Override
+			public AnnotationVisitor visitAnnotation(String desc, boolean visible) {
+				//This method visitor checks all methods and only acts upon those with the Init annotation.
+				//This lets us warn the user about invalid init methods and misuse of the annotation
+				if (!desc.equals("Lde/hysky/skyblocker/annotations/Init;")) return super.visitAnnotation(desc, visible);
+
+				//Delegates adding the method call to the map to the InitAnnotationVisitor since we don't have a value to put in the map here
+				return new InitAnnotationVisitor(methodSignatures, getMethodCall());
+			}
+
+			private @NotNull InitProcessor.MethodReference getMethodCall() {
+				String className = classReader.getClassName();
+				String methodCallString = className + "." + methodName;
+				if ((access & Opcodes.ACC_PUBLIC) == 0) throw new IllegalStateException(methodCallString + ": Initializer methods must be public");
+				if ((access & Opcodes.ACC_STATIC) == 0) throw new IllegalStateException(methodCallString + ": Initializer methods must be static");
+				if (!descriptor.equals("()V")) throw new IllegalStateException(methodCallString + ": Initializer methods must have no args and a void return type");
+
+				//Interface static methods need special handling, so we add a special marker for that
+				boolean itf = (classReader.getAccess() & Opcodes.ACC_INTERFACE) != 0;
+
+				return new InitProcessor.MethodReference(className, methodName, descriptor, itf);
+			}
+		};
+	}
+
+	static class InitAnnotationVisitor extends AnnotationVisitor {
+		private final Map<InitProcessor.MethodReference, Integer> methodSignatures;
+		private final InitProcessor.MethodReference methodCall;
+
+		protected InitAnnotationVisitor(Map<InitProcessor.MethodReference, Integer> methodSignatures, InitProcessor.MethodReference methodCall) {
+			super(Opcodes.ASM9);
+			this.methodSignatures = methodSignatures;
+			this.methodCall = methodCall;
+		}
+
+		@Override
+		public void visitEnd() {
+			//Annotations that use the default value for the priority field will not be called by the visit method, so we have to handle them here.
+			methodSignatures.putIfAbsent(methodCall, 0);
+			super.visitEnd();
+		}
+
+		@Override
+		public void visit(String name, Object value) {
+			if (name.equals("priority")) {
+				methodSignatures.put(methodCall, (int) value);
+			}
+			super.visit(name, value);
+		}
+	}
+}