feat: specifying ClassNode is now optional

3 files changed, 96 insertions, 60 deletions

diff --git a/src/main/java/ftbsc/lll/processor/LilleroProcessor.java b/src/main/java/ftbsc/lll/processor/LilleroProcessor.java
index 467fa88..fdddd9a 100644
--- a/src/main/java/ftbsc/lll/processor/LilleroProcessor.java
+++ b/src/main/java/ftbsc/lll/processor/LilleroProcessor.java
@@ -7,6 +7,7 @@ import ftbsc.lll.exceptions.InvalidResourceException;
 import ftbsc.lll.exceptions.OrphanElementException;
 import ftbsc.lll.processor.annotations.*;
 import ftbsc.lll.processor.tools.containers.ClassContainer;
+import ftbsc.lll.processor.tools.containers.InjectorInfo;
 import ftbsc.lll.processor.tools.containers.MethodContainer;
 import ftbsc.lll.processor.tools.obfuscation.ObfuscationMapper;
 import ftbsc.lll.proxies.ProxyType;
@@ -172,9 +173,12 @@ public class LilleroProcessor extends AbstractProcessor {
 			List<? extends TypeMirror> params = ((ExecutableType) e.asType()).getParameterTypes();
 			return e.getAnnotation(Injector.class) != null
 				&& e.getAnnotation(Target.class) == null
-				&& params.size() == 2
-				&& this.processingEnv.getTypeUtils().isSameType(params.get(0), classNodeType)
-				&& this.processingEnv.getTypeUtils().isSameType(params.get(1), methodNodeType);
+				&& (
+					(params.size() == 2
+						&& this.processingEnv.getTypeUtils().isSameType(params.get(0), classNodeType)
+						&& this.processingEnv.getTypeUtils().isSameType(params.get(1), methodNodeType)
+					) || (params.size() == 1 && this.processingEnv.getTypeUtils().isSameType(params.get(0), methodNodeType))
+			);
 		})) return true;
 		else {
 			this.processingEnv.getMessager().printMessage(Diagnostic.Kind.WARNING,
@@ -301,7 +305,7 @@ public class LilleroProcessor extends AbstractProcessor {
 						matchedInjectors.add(injector);
 						toGenerate.put(
 							String.format("%sInjector%d", cl.getSimpleName(), iterationNumber),
-							new InjectorInfo(injector, tg, targetAnn)
+							new InjectorInfo(injector, tg, targetAnn, this.processingEnv, this.mapper)
 						);
 						iterationNumber++; //increment is only used by injectors
 					} else {
@@ -331,21 +335,6 @@ public class LilleroProcessor extends AbstractProcessor {
 
 		//iterate over the map and generate the classes
 		for(String injName : toGenerate.keySet()) {
-			MethodSpec inject = MethodSpec.methodBuilder("inject")
-				.addModifiers(Modifier.PUBLIC)
-				.returns(void.class)
-				.addAnnotation(Override.class)
-				.addParameter(ParameterSpec.builder(
-					TypeName.get(processingEnv
-						.getElementUtils()
-						.getTypeElement("org.objectweb.asm.tree.ClassNode").asType()), "clazz").build())
-				.addParameter(ParameterSpec.builder(
-					TypeName.get(processingEnv
-						.getElementUtils()
-						.getTypeElement("org.objectweb.asm.tree.MethodNode").asType()), "main").build())
-				.addStatement(String.format("super.%s(clazz, main)", toGenerate.get(injName).injector.getSimpleName()), TypeName.get(cl.asType()))
-				.build();
-
 			MethodContainer target = toGenerate.get(injName).target;
 			TypeSpec injectorClass = TypeSpec.classBuilder(injName)
 				.addModifiers(Modifier.PUBLIC)
@@ -358,7 +347,7 @@ public class LilleroProcessor extends AbstractProcessor {
 				.addMethod(buildStringReturnMethod("methodName", obfuscateInjectorMetadata ? target.nameObf : target.name))
 				.addMethod(buildStringReturnMethod("methodDesc", obfuscateInjectorMetadata ? target.descriptorObf : target.descriptor))
 				.addMethods(generateDummies(targets))
-				.addMethod(inject)
+				.addMethod(generateInjector(toGenerate.get(injName), this.processingEnv))
 				.build();
 
 			JavaFile javaFile = JavaFile.builder(packageName, injectorClass).build();
@@ -393,43 +382,4 @@ public class LilleroProcessor extends AbstractProcessor {
 			throw new RuntimeException(e);
 		}
 	}
-
-	/**
-	 * Container for information about a class that is to be generated.
-	 * Only used internally.
-	 */
-	private class InjectorInfo {
-		/**
-		 * The {@link ExecutableElement} corresponding to the injector method.
-		 */
-		public final ExecutableElement injector;
-
-		/**
-		 * The {@link ExecutableElement} corresponding to the target method stub.
-		 */
-		public final ExecutableElement targetStub;
-
-		/**
-		 * The reason for the injection.
-		 */
-		public final String reason;
-
-		/**
-		 * The {@link MethodContainer} corresponding to the target method.
-		 */
-		private final MethodContainer target;
-
-		/**
-		 * Public constructor.
-		 * @param injector the injector {@link ExecutableElement}
-		 * @param targetStub the target {@link ExecutableElement}
-		 * @param targetAnn the relevant {@link Target} annotation
-		 */
-		public InjectorInfo(ExecutableElement injector, ExecutableElement targetStub, Target targetAnn) {
-			this.injector = injector;
-			this.targetStub = targetStub;
-			this.reason = injector.getAnnotation(Injector.class).reason();
-			this.target = MethodContainer.from(targetStub, targetAnn, null, processingEnv, mapper);
-		}
-	}
 }
diff --git a/src/main/java/ftbsc/lll/processor/tools/JavaPoetUtils.java b/src/main/java/ftbsc/lll/processor/tools/JavaPoetUtils.java
index 67d6cec..62d82e3 100644
--- a/src/main/java/ftbsc/lll/processor/tools/JavaPoetUtils.java
+++ b/src/main/java/ftbsc/lll/processor/tools/JavaPoetUtils.java
@@ -5,6 +5,7 @@ import ftbsc.lll.processor.annotations.Find;
 import ftbsc.lll.processor.annotations.Target;
 import ftbsc.lll.processor.tools.containers.ClassContainer;
 import ftbsc.lll.processor.tools.containers.FieldContainer;
+import ftbsc.lll.processor.tools.containers.InjectorInfo;
 import ftbsc.lll.processor.tools.containers.MethodContainer;
 import ftbsc.lll.processor.tools.obfuscation.ObfuscationMapper;
 import ftbsc.lll.proxies.ProxyType;
@@ -122,7 +123,7 @@ public class JavaPoetUtils {
 	/**
 	 * Generates a {@link HashSet} of dummy overrides given a {@link Collection} stubs.
  	 * @param dummies the stubs
-	 * @return the generated {@link HashSet}
+	 * @return a {@link HashSet} containing the generated {@link MethodSpec}s
 	 * @since 0.5.0
 	 */
 	public static HashSet<MethodSpec> generateDummies(Collection<ExecutableElement> dummies) {
@@ -135,4 +136,41 @@ public class JavaPoetUtils {
 				);
 		return specs;
 	}
+
+	/**
+	 * Generates the wrapper around a certain injector.
+	 * @param inj the {@link InjectorInfo} carrying the information about the target injector
+	 * @param env the {@link ProcessingEnvironment} to perform the operation in
+	 * @return the generated {@link MethodSpec} for the injector
+	 * @since 0.6.0
+	 */
+	public static MethodSpec generateInjector(InjectorInfo inj, ProcessingEnvironment env) {
+		MethodSpec.Builder injectBuilder = MethodSpec.methodBuilder("inject")
+			.addModifiers(Modifier.PUBLIC)
+			.returns(void.class)
+			.addAnnotation(Override.class);
+
+		int argumentCount = inj.injector.getParameters().size();
+
+		if(argumentCount == 2) {
+			injectBuilder
+				.addParameter(ParameterSpec.builder(
+						TypeName.get(env
+							.getElementUtils()
+							.getTypeElement("org.objectweb.asm.tree.ClassNode").asType()), "clazz")
+					.build());
+		}
+
+		injectBuilder
+			.addParameter(ParameterSpec.builder(
+				TypeName.get(env
+					.getElementUtils()
+					.getTypeElement("org.objectweb.asm.tree.MethodNode").asType()), "main")
+				.build());
+
+		if(argumentCount == 2) injectBuilder.addStatement("super.$L(clazz, main)", inj.injector.getSimpleName());
+		else injectBuilder.addStatement("super.$L(main)", inj.injector.getSimpleName());
+
+		return injectBuilder.build();
+	}
 }
diff --git a/src/main/java/ftbsc/lll/processor/tools/containers/InjectorInfo.java b/src/main/java/ftbsc/lll/processor/tools/containers/InjectorInfo.java
new file mode 100644
index 0000000..e2d7ec8
--- /dev/null
+++ b/src/main/java/ftbsc/lll/processor/tools/containers/InjectorInfo.java
@@ -0,0 +1,48 @@
+package ftbsc.lll.processor.tools.containers;
+
+import ftbsc.lll.processor.annotations.Injector;
+import ftbsc.lll.processor.annotations.Target;
+import ftbsc.lll.processor.tools.obfuscation.ObfuscationMapper;
+
+import javax.annotation.processing.ProcessingEnvironment;
+import javax.lang.model.element.ExecutableElement;
+
+/**
+ * Container for information about a class that is to be generated.
+ */
+public class InjectorInfo {
+	/**
+	 * The {@link ExecutableElement} corresponding to the injector method.
+	 */
+	public final ExecutableElement injector;
+
+	/**
+	 * The {@link ExecutableElement} corresponding to the target method stub.
+	 */
+	public final ExecutableElement targetStub;
+
+	/**
+	 * The reason for the injection.
+	 */
+	public final String reason;
+
+	/**
+	 * The {@link MethodContainer} corresponding to the target method.
+	 */
+	public final MethodContainer target;
+
+	/**
+	 * Public constructor.
+	 * @param injector the injector {@link ExecutableElement}
+	 * @param targetStub the target {@link ExecutableElement}
+	 * @param targetAnn the relevant {@link Target} annotation
+	 * @param env the {@link ProcessingEnvironment} to be used to locate the class
+	 * @param mapper the {@link ObfuscationMapper} to be used, may be null
+	 */
+	public InjectorInfo(ExecutableElement injector, ExecutableElement targetStub, Target targetAnn, ProcessingEnvironment env, ObfuscationMapper mapper) {
+		this.injector = injector;
+		this.targetStub = targetStub;
+		this.reason = injector.getAnnotation(Injector.class).reason();
+		this.target = MethodContainer.from(targetStub, targetAnn, null, env, mapper);
+	}
+}
\ No newline at end of file