fix: fixed mappings bug

3 files changed, 91 insertions, 26 deletions

diff --git a/src/main/java/ftbsc/lll/processor/LilleroProcessor.java b/src/main/java/ftbsc/lll/processor/LilleroProcessor.java
index 177e12f..491ce68 100644
--- a/src/main/java/ftbsc/lll/processor/LilleroProcessor.java
+++ b/src/main/java/ftbsc/lll/processor/LilleroProcessor.java
@@ -15,7 +15,6 @@ import javax.annotation.processing.*;
 import javax.lang.model.SourceVersion;
 import javax.lang.model.element.*;
 import javax.lang.model.type.ExecutableType;
-import javax.lang.model.type.MirroredTypesException;
 import javax.lang.model.type.TypeMirror;
 import javax.tools.Diagnostic;
 import javax.tools.FileObject;
@@ -176,8 +175,8 @@ public class LilleroProcessor extends AbstractProcessor {
 	private static String findClassName(Patch patchAnn, FindMethod methodAnn, ObfuscationMapper mapper) {
 		String fullyQualifiedName =
 			methodAnn == null || methodAnn.parent() == Object.class
-				? getClassFullyQualifiedName(patchAnn, p -> patchAnn.value())
-				: getClassFullyQualifiedName(methodAnn, m -> methodAnn.parent());
+				? getClassFullyQualifiedName(patchAnn, Patch::value)
+				: getClassFullyQualifiedName(methodAnn, FindMethod::parent);
 		return findClassName(fullyQualifiedName, mapper);
 	}
 
@@ -200,8 +199,8 @@ public class LilleroProcessor extends AbstractProcessor {
 	 * @return the internal class name
 	 * @since 0.3.0
 	 */
-	private static String findMemberName(String parentFQN, String memberName, ObfuscationMapper mapper) {
-		return mapper == null ? memberName : mapper.obfuscateMember(parentFQN, memberName);
+	private static String findMemberName(String parentFQN, String memberName, String methodDescriptor, ObfuscationMapper mapper) {
+		return mapper == null ? memberName : mapper.obfuscateMember(parentFQN, memberName, methodDescriptor);
 	}
 
 	/**
@@ -213,14 +212,18 @@ public class LilleroProcessor extends AbstractProcessor {
 	 * @return the internal class name
 	 * @since 0.3.0
 	 */
-	private static String findMethodName(String parentFQN, FindMethod methodAnn, ExecutableElement stub, ObfuscationMapper mapper) {
+	private String findMethodName(String parentFQN, FindMethod methodAnn, ExecutableElement stub, ObfuscationMapper mapper) {
 		String methodName = methodAnn == null ? stub.getSimpleName().toString() : methodAnn.name();
+		String methodDescriptor;
+		if(methodAnn == null)
+			methodDescriptor = descriptorFromExecutableElement(stub);
+		else methodDescriptor = methodDescriptorFromParams(methodAnn, FindMethod::params, processingEnv.getElementUtils());
 		try {
-			methodName = findMemberName(parentFQN, methodName, mapper);
+			methodName = findMemberName(parentFQN, methodName, methodDescriptor, mapper);
 		} catch(MappingNotFoundException e) {
 			//not found: try again with the name of the annotated method
 			if(methodAnn == null) {
-				methodName = findMemberName(parentFQN, stub.getSimpleName().toString(), mapper);
+				methodName = findMemberName(parentFQN, stub.getSimpleName().toString(), methodDescriptor, mapper);
 			} else throw e;
 		}
 		return methodName;
@@ -235,7 +238,7 @@ public class LilleroProcessor extends AbstractProcessor {
 	 * @return the internal class name
 	 * @since 0.3.0
 	 */
-	private static String findMethodName(Patch patchAnn, FindMethod methodAnn, ExecutableElement stub, ObfuscationMapper mapper) {
+	private String findMethodName(Patch patchAnn, FindMethod methodAnn, ExecutableElement stub, ObfuscationMapper mapper) {
 		return findMethodName(findClassName(patchAnn, methodAnn, mapper), methodAnn, stub, mapper);
 	}
 
@@ -325,8 +328,8 @@ public class LilleroProcessor extends AbstractProcessor {
 		FindField fieldAnn = stub.getAnnotation(FindField.class);
 		String parentName;
 		if(fieldAnn.parent().equals(Object.class))
-			parentName = getClassFullyQualifiedName(patchAnn, p -> patchAnn.value());
-		else parentName = getClassFullyQualifiedName(fieldAnn, f -> fieldAnn.parent());
+			parentName = getClassFullyQualifiedName(patchAnn, Patch::value);
+		else parentName = getClassFullyQualifiedName(fieldAnn, FindField::parent);
 		parentName = findClassName(parentName, mapper);
 		String name = fieldAnn.name().equals("")
 			? stub.getSimpleName().toString()
@@ -408,15 +411,7 @@ public class LilleroProcessor extends AbstractProcessor {
 				if(injectionCandidates.size() == 1)
 					injectionTarget = injectionCandidates.get(0);
 
-				List<TypeMirror> params = new ArrayList<>();
-				try {
-					params.addAll(Arrays.stream(injectorAnn.params())
-							.map(Class::getCanonicalName)
-							.map(fqn -> processingEnv.getElementUtils().getTypeElement(fqn).asType())
-							.collect(Collectors.toList()));
-				} catch(MirroredTypesException e) {
-					params.addAll(e.getTypeMirrors());
-				}
+				List<TypeMirror> params = classArrayFromAnnotation(injectorAnn, Injector::params, processingEnv.getElementUtils());
 
 				if(params.size() != 0) {
 					StringBuilder descr = new StringBuilder("(");
diff --git a/src/main/java/ftbsc/lll/processor/tools/ASTUtils.java b/src/main/java/ftbsc/lll/processor/tools/ASTUtils.java
index 8361367..faac3dc 100644
--- a/src/main/java/ftbsc/lll/processor/tools/ASTUtils.java
+++ b/src/main/java/ftbsc/lll/processor/tools/ASTUtils.java
@@ -6,12 +6,17 @@ import com.squareup.javapoet.MethodSpec;
 import com.squareup.javapoet.TypeName;
 import ftbsc.lll.tools.DescriptorBuilder;
 
+import javax.annotation.processing.ProcessingEnvironment;
 import javax.lang.model.element.ExecutableElement;
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.TypeElement;
 import javax.lang.model.type.MirroredTypeException;
+import javax.lang.model.type.MirroredTypesException;
 import javax.lang.model.type.TypeMirror;
+import javax.lang.model.util.Elements;
 import java.lang.annotation.Annotation;
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -135,19 +140,53 @@ public class ASTUtils {
 	}
 
 	/**
-	 * Safely converts a {@link Class} to its fully qualified name. See
-	 * <a href="https://area-51.blog/2009/02/13/getting-class-values-from-annotations-in-an-annotationprocessor">this blogpost</a>
-	 * for more information.
+	 * Safely extracts a {@link Class} from an annotation and gets its fully qualified name.
 	 * @param ann the annotation containing the class
 	 * @param fun the annotation function returning the class
 	 * @return the fully qualified name of the given class
 	 * @since 0.3.0
 	 */
-	public static <T extends Annotation> String getClassFullyQualifiedName(T ann, Function<Annotation, Class<?>> fun) {
+	public static <T extends Annotation> String getClassFullyQualifiedName(T ann, Function<T, Class<?>> fun) {
 		try {
 			return fun.apply(ann).getCanonicalName();
 		} catch(MirroredTypeException e) {
 			return e.getTypeMirror().toString();
 		}
 	}
+
+	/**
+	 * Safely extracts a {@link Class} array from an annotation.
+	 * @param ann the annotation containing the class
+	 * @param fun the annotation function returning the class
+	 * @param elementUtils the element utils corresponding to the {@link ProcessingEnvironment}
+	 * @return a list of {@link TypeMirror}s representing the classes
+	 * @since 0.3.0
+	 */
+	public static <T extends Annotation> List<TypeMirror> classArrayFromAnnotation(T ann, Function<T, Class<?>[]> fun, Elements elementUtils) {
+		List<TypeMirror> params = new ArrayList<>();
+		try {
+			params.addAll(Arrays.stream(fun.apply(ann))
+				.map(Class::getCanonicalName)
+				.map(fqn -> elementUtils.getTypeElement(fqn).asType())
+				.collect(Collectors.toList()));
+		} catch(MirroredTypesException e) {
+			params.addAll(e.getTypeMirrors());
+		}
+		return params;
+	}
+
+	/**
+	 * Builds a (partial, not including the return type) method descriptor from its parameters
+	 * @param ann the annotation containing the class
+	 * @param fun the annotation function returning the class
+	 * @return the method descriptor
+	 */
+	public static <T extends Annotation> String methodDescriptorFromParams(T ann, Function<T, Class<?>[]> fun, Elements elementUtils) {
+		List<TypeMirror> mirrors = classArrayFromAnnotation(ann, fun, elementUtils);
+		StringBuilder sb = new StringBuilder("(");
+		for(TypeMirror t : mirrors)
+			sb.append(descriptorFromType(t));
+		sb.append(")");
+		return sb.toString();
+	}
 }
diff --git a/src/main/java/ftbsc/lll/processor/tools/obfuscation/ObfuscationMapper.java b/src/main/java/ftbsc/lll/processor/tools/obfuscation/ObfuscationMapper.java
index 12ae23f..5e54775 100644
--- a/src/main/java/ftbsc/lll/processor/tools/obfuscation/ObfuscationMapper.java
+++ b/src/main/java/ftbsc/lll/processor/tools/obfuscation/ObfuscationMapper.java
@@ -1,10 +1,13 @@
 package ftbsc.lll.processor.tools.obfuscation;
 
+import ftbsc.lll.exceptions.AmbiguousDefinitionException;
 import ftbsc.lll.exceptions.MappingNotFoundException;
 
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 /**
@@ -76,14 +79,15 @@ public class ObfuscationMapper {
 	 * with a space, because that's how it is in .tsrg files.
 	 * @param parentName the unobfuscated internal name of the parent class
 	 * @param memberName the field name or method signature
+	 * @param methodDescriptor the optional descriptor of the member, may be null or partial
 	 * @return the obfuscated name of the given member
 	 * @throws MappingNotFoundException if no mapping is found
 	 */
-	public String obfuscateMember(String parentName, String memberName) {
+	public String obfuscateMember(String parentName, String memberName, String methodDescriptor) {
 		ObfuscationData data = mapper.get(parentName.replace('.', '/'));
 		if(data == null)
 			throw new MappingNotFoundException(parentName + "::" + memberName);
-		String member = data.members.get(memberName);
+		String member = data.get(memberName, methodDescriptor);
 		if(member == null)
 			throw new MappingNotFoundException(parentName + "::" + memberName);
 		return member;
@@ -166,5 +170,32 @@ public class ObfuscationMapper {
 			else if (split.length == 3) //method
 				members.put(split[0] + " " + split[1], split[2]);
 		}
+
+		/**
+		 * Gets an obfuscated member given the method name and a method descriptor,
+		 * which may be partial (i.e. not include return type) or null if the member
+		 * is not a method.
+		 * @param memberName member name
+		 * @param methodDescriptor the method descriptor, or null if it's not a method
+		 * @return the requested obfuscated name, or null if nothing was found
+		 * @throws AmbiguousDefinitionException if not enough data was given to uniquely identify a mapping
+		 */
+		public String get(String memberName, String methodDescriptor) {
+			if(methodDescriptor == null)
+				return members.get(memberName);
+			List<String> candidates = members.values().stream().filter(m -> m.startsWith(memberName)).collect(Collectors.toList());
+			if(candidates.size() == 1)
+				return members.get(candidates.get(0));
+			String signature = memberName + " " + methodDescriptor;
+			candidates = candidates.stream().filter(m -> m.startsWith(signature)).collect(Collectors.toList());
+			switch(candidates.size()) {
+				case 0:
+					return null;
+				case 1:
+					return members.get(candidates.get(0));
+				default:
+					throw new AmbiguousDefinitionException("Mapper could not uniquely identify method " + unobf + "::" + memberName);
+			}
+		}
 	}
 }