Fix mapping not effectively traversing up inheritance trees

Author: Col-E

recaf-core/src/main/java/software/coley/recaf/services/mapping/MappingApplier.java

@@ -191,6 +191,12 @@ private Mappings enrich(@Nonnull Mappings mappings) {
 			MappingsAdapter adapter = new MappingsAdapter(fieldDifferentiation, varDifferentiation);
 			adapter.importIntermediate(intermediateMappings);
 			mappings = adapter;
+
+			// Consider the method family of: "Circle.area() --> Shape.area()"
+			// - If you provide a mapping for just "Shape.area()" a parent-lookup from "Circle.area()" can see the mapping by checking parents.
+			// - If you provide a mapping for just "Circle.area()" then "Shape.area()" needs to look at children.
+			// This infilling will flesh out the mappings so that if you provide either end of the family, the other end will be present in the mappings.
+			adapter.infillMethodFamilies(intermediateMappings, inheritanceGraph);
 		}
 
 		// Check if mappings can be enriched with type look-ups

recaf-core/src/main/java/software/coley/recaf/services/mapping/MappingsAdapter.java

@@ -3,15 +3,28 @@
 import jakarta.annotation.Nonnull;
 import jakarta.annotation.Nullable;
 import software.coley.recaf.info.ClassInfo;
+import software.coley.recaf.info.member.MethodMember;
 import software.coley.recaf.path.ClassPathNode;
 import software.coley.recaf.services.inheritance.InheritanceGraph;
 import software.coley.recaf.services.inheritance.InheritanceVertex;
-import software.coley.recaf.services.mapping.data.*;
+import software.coley.recaf.services.mapping.data.ClassMapping;
+import software.coley.recaf.services.mapping.data.ClassMappingKey;
+import software.coley.recaf.services.mapping.data.FieldMapping;
+import software.coley.recaf.services.mapping.data.FieldMappingKey;
+import software.coley.recaf.services.mapping.data.MappingKey;
+import software.coley.recaf.services.mapping.data.MethodMapping;
+import software.coley.recaf.services.mapping.data.MethodMappingKey;
+import software.coley.recaf.services.mapping.data.VariableMapping;
+import software.coley.recaf.services.mapping.data.VariableMappingKey;
 import software.coley.recaf.workspace.model.Workspace;
 
+import java.util.Collection;
+import java.util.Comparator;
 import java.util.HashMap;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.TreeMap;
 import java.util.function.Function;
 
@@ -20,6 +33,7 @@
  * <b>Enhancements</b>
  * <ol>
  * <li>Import mapping entries from a {@link IntermediateMappings} instance.</li>
+ * <li>Enhance mapping entries with method family information from {@link InheritanceGraph}, see {@link #infillMethodFamilies(IntermediateMappings, InheritanceGraph)}</li>
  * <li>Enhance field/method lookups with inheritance info from {@link InheritanceGraph}, see {@link #enableHierarchyLookup(InheritanceGraph)}.</li>
  * <li>Enhance inner/outer class mapping edge cases via {@link #enableClassLookup(Workspace)}.</li>
  * <li>Adapt keys in cases where fields/vars do not have type info associated with them <i>(for formats that suck)</i>.</li>
@@ -66,11 +80,9 @@ public void importIntermediate(@Nonnull IntermediateMappings mappings) {
 				String newName = fieldMapping.getNewName();
 				if (!oldName.equals(newName)) {
 					if (doesSupportFieldTypeDifferentiation()) {
-						addField(fieldMapping.getOwnerName(), oldName,
-								fieldMapping.getDesc(), newName);
+						addField(fieldMapping.getOwnerName(), oldName, fieldMapping.getDesc(), newName);
 					} else {
-						addField(fieldMapping.getOwnerName(), oldName,
-								newName);
+						addField(fieldMapping.getOwnerName(), oldName, newName);
 					}
 				}
 			}
@@ -79,10 +91,8 @@ public void importIntermediate(@Nonnull IntermediateMappings mappings) {
 				String oldMethodDesc = methodMapping.getDesc();
 				String newMethodName = methodMapping.getNewName();
 				if (!oldMethodName.equals(newMethodName))
-					addMethod(methodMapping.getOwnerName(), oldMethodName,
-							oldMethodDesc, newMethodName);
-				for (VariableMapping variableMapping :
-						mappings.getMethodVariableMappings(className, oldMethodName, oldMethodDesc)) {
+					addMethod(methodMapping.getOwnerName(), oldMethodName, oldMethodDesc, newMethodName);
+				for (VariableMapping variableMapping : mappings.getMethodVariableMappings(className, oldMethodName, oldMethodDesc)) {
 					addVariable(className, oldMethodName, oldMethodDesc,
 							variableMapping.getOldName(), variableMapping.getDesc(), variableMapping.getIndex(),
 							variableMapping.getNewName());
@@ -91,6 +101,94 @@ public void importIntermediate(@Nonnull IntermediateMappings mappings) {
 		}
 	}
 
+	/**
+	 * Enriches the provided mappings by normalizing method mappings across method families.
+	 *
+	 * @param inputMappings
+	 * 		Original mappings to enrich.
+	 * @param inheritanceGraph
+	 * 		Inheritance graph to use for looking up method families.
+	 */
+	public void infillMethodFamilies(@Nonnull IntermediateMappings inputMappings, @Nonnull InheritanceGraph inheritanceGraph) {
+		// Collect all method mappings in the input.
+		List<MethodMapping> explicitMappings = inputMappings.getMethods().values().stream()
+				.flatMap(Collection::stream)
+				.sorted(Comparator
+						.comparing(MethodMapping::getOwnerName)
+						.thenComparing(MethodMapping::getOldName)
+						.thenComparing(MethodMapping::getDesc)
+						.thenComparing(MethodMapping::getNewName))
+				.toList();
+
+		// Compute which methods are mapped to what, and detect any conflicts in the provided mappings.
+		Map<MethodMappingKey, MethodMapping> origins = new HashMap<>();
+		for (MethodMapping explicitMapping : explicitMappings) {
+			MethodMappingKey key = new MethodMappingKey(explicitMapping.getOwnerName(), explicitMapping.getOldName(), explicitMapping.getDesc());
+			MethodMapping previous = origins.putIfAbsent(key, explicitMapping);
+			if (previous != null && !previous.getNewName().equals(explicitMapping.getNewName()))
+				throw newMethodConflict(key, key.getOwner(), previous, explicitMapping);
+		}
+
+		// Iterate over the method mappings, and for each mapping,
+		// check if there are any methods in the same override family that would be affected by the same mapping.
+		// If so, add the same mapping for those methods as well.
+		for (MethodMapping explicitMapping : explicitMappings) {
+			// Skip constructors and static initializers.
+			String methodName = explicitMapping.getOldName();
+			if (!methodName.isEmpty() && methodName.charAt(0) == '<')
+				continue;
+
+			// Skip enriching mappings on library classes. Can't map those anyways.
+			String ownerName = explicitMapping.getOwnerName();
+			InheritanceVertex ownerVertex = inheritanceGraph.getVertex(ownerName);
+			if (ownerVertex == null || ownerVertex.isLibraryVertex())
+				continue;
+
+			// Skip enriching mappings on private or static methods, as they cannot be inherited.
+			ClassInfo ownerInfo = ownerVertex.getValue();
+			MethodMember ownerMethod = ownerInfo.getDeclaredMethod(methodName, explicitMapping.getDesc());
+			if (ownerMethod == null || ownerMethod.hasPrivateModifier() || ownerMethod.hasStaticModifier())
+				continue;
+
+			// Iterate over classes in the inheritance family and apply the same mapping to any method with the same signature that is inherited by the owner method.
+			List<InheritanceVertex> sortedFamily = inheritanceGraph.getVertexFamily(ownerName, false).stream()
+					.sorted(Comparator.comparing(InheritanceVertex::getName))
+					.toList();
+			for (InheritanceVertex familyVertex : sortedFamily) {
+				// Again, skipping library classes.
+				if (familyVertex.isLibraryVertex())
+					continue;
+
+				// Check if the method is inherited by this family member.
+				// If it isn't, then we don't want to apply the mapping to it.
+				ClassInfo familyInfo = familyVertex.getValue();
+				MethodMember familyMethod = familyInfo.getDeclaredMethod(methodName, explicitMapping.getDesc());
+				if (familyMethod == null || !MappingsAdapter.isInheritedMethod(ownerInfo, ownerMethod, familyInfo, familyMethod))
+					continue;
+
+				// The method is inherited by this family member. Enrich the adapter.
+				String familyOwnerName = familyInfo.getName();
+				MethodMappingKey familyKey = new MethodMappingKey(familyOwnerName, methodName, explicitMapping.getDesc());
+				String existing = getMappedMethodName(familyOwnerName, methodName, explicitMapping.getDesc());
+				if (existing == null) {
+					// Add the method mapping to this family member, and mark the original mapping as the origin for this mapping.
+					addMethod(familyOwnerName, methodName, explicitMapping.getDesc(), explicitMapping.getNewName());
+					origins.put(familyKey, explicitMapping);
+				} else if (!existing.equals(explicitMapping.getNewName())) {
+					// Shouldn't happen due to earlier check, but just in case... check for conflicting mappings.
+					MethodMapping previous = origins.get(familyKey);
+					if (previous == null)
+						previous = new MethodMapping(familyOwnerName, methodName, explicitMapping.getDesc(), existing);
+					throw newMethodConflict(familyKey, familyOwnerName, previous, explicitMapping);
+				} else {
+					// This family member already has the same mapping, so we can just mark the
+					// original mapping as the origin for this family member's mapping.
+					origins.putIfAbsent(familyKey, explicitMapping);
+				}
+			}
+		}
+	}
+
 	@Nullable
 	@Override
 	public String getMappedClassName(@Nonnull String internalName) {
@@ -106,9 +204,8 @@ public String getMappedClassName(@Nonnull String internalName) {
 					if (outerName != null && outerName.length() < name.length()) {
 						String inner = name.substring(outerName.length() + 1);
 						String outerMapped = getMappedClassName(outerName);
-						if (outerMapped != null) {
+						if (outerMapped != null)
 							mapped = outerMapped + "$" + inner;
-						}
 					}
 				}
 			} else if (isInner(internalName)) {
@@ -117,9 +214,8 @@ public String getMappedClassName(@Nonnull String internalName) {
 				String inner = internalName.substring(split + 1);
 				String outer = internalName.substring(0, split);
 				String outerMapped = getMappedClassName(outer);
-				if (outerMapped != null) {
+				if (outerMapped != null)
 					mapped = outerMapped + "$" + inner;
-				}
 			}
 		}
 		return mapped;
@@ -130,9 +226,8 @@ public String getMappedClassName(@Nonnull String internalName) {
 	public String getMappedFieldName(@Nonnull String ownerName, @Nonnull String fieldName, @Nonnull String fieldDesc) {
 		MappingKey key = getFieldKey(ownerName, fieldName, fieldDesc);
 		String mapped = mappings.get(key);
-		if (mapped == null && inheritanceGraph != null) {
+		if (mapped == null && inheritanceGraph != null)
 			mapped = findInParent(ownerName, parent -> getFieldKey(parent, fieldName, fieldDesc));
-		}
 		return mapped;
 	}
 
@@ -141,9 +236,8 @@ public String getMappedFieldName(@Nonnull String ownerName, @Nonnull String fiel
 	public String getMappedMethodName(@Nonnull String ownerName, @Nonnull String methodName, @Nonnull String methodDesc) {
 		MappingKey key = getMethodKey(ownerName, methodName, methodDesc);
 		String mapped = mappings.get(key);
-		if (mapped == null && inheritanceGraph != null) {
-			mapped = findInParent(ownerName, parent -> getMethodKey(parent, methodName, methodDesc));
-		}
+		if (mapped == null && inheritanceGraph != null)
+			mapped = findMethodInParent(ownerName, methodName, methodDesc);
 		return mapped;
 	}
 
@@ -197,9 +291,8 @@ public boolean doesSupportVariableTypeDifferentiation() {
 	 *
 	 * @return The first mapping match in a parent class found by the lookup function.
 	 */
+	@Nullable
 	private String findInParent(String owner, Function<String, ? extends MappingKey> lookup) {
-		// TODO: Diamond class hierarchy doesn't work with this.
-		//  We need to visit the whole family tree.
 		InheritanceVertex vertex = inheritanceGraph.getVertex(owner);
 		if (vertex == null)
 			return null;
@@ -208,13 +301,86 @@ private String findInParent(String owner, Function<String, ? extends MappingKey>
 			vertex = iterator.next();
 			MappingKey key = lookup.apply(vertex.getName());
 			String result = mappings.get(key);
-			if (result != null) {
+			if (result != null)
 				return result;
-			}
 		}
 		return null;
 	}
 
+	/**
+	 * @param ownerName
+	 * 		Internal name of the class <i>"defining"</i> the member.
+	 * 		<i>(Location in reference may not be where the member is actually defined, hence this lookup)</i>
+	 * @param methodName
+	 * 		Name of the method.
+	 * @param methodDesc
+	 * 		Descriptor of the method.
+	 *
+	 * @return The first mapping match in a parent class with the same name/desc <i>(Some access restrictions apply)</i>.
+	 */
+	@Nullable
+	private String findMethodInParent(@Nonnull String ownerName, @Nonnull String methodName, @Nonnull String methodDesc) {
+		// Skip if not in workspace/inheritance graph.
+		InheritanceVertex ownerVertex = inheritanceGraph.getVertex(ownerName);
+		if (ownerVertex == null)
+			return null;
+
+		// Skip if the method is defined in the owner class and is private/static, as it would not be inherited.
+		ClassInfo ownerInfo = ownerVertex.getValue();
+		MethodMember ownerMethod = ownerInfo.getDeclaredMethod(methodName, methodDesc);
+		if (ownerMethod != null && (ownerMethod.hasPrivateModifier() || ownerMethod.hasStaticModifier()))
+			return null;
+
+		// Iterate through parents and find the first method match that is inherited by the owner class.
+		Iterator<InheritanceVertex> iterator = ownerVertex.allParents().iterator();
+		while (iterator.hasNext()) {
+			InheritanceVertex parentVertex = iterator.next();
+			ClassInfo parentInfo = parentVertex.getValue();
+
+			// If the parent class doesn't have a method with the same name/desc, skip it.
+			MethodMember parentMethod = parentInfo.getDeclaredMethod(methodName, methodDesc);
+			if (parentMethod == null || !isInheritedMethod(ownerInfo, ownerMethod, parentInfo, parentMethod))
+				continue;
+
+			// Otherwise, we have a method that is inherited by the owner class, so we can check for a mapping on it.
+			String mapped = mappings.get(getMethodKey(parentVertex.getName(), methodName, methodDesc));
+			if (mapped != null)
+				return mapped;
+		}
+
+		// No mapping found in any parent class.
+		return null;
+	}
+
+	/**
+	 * @param ownerInfo
+	 * 		Class defining the target method.
+	 * @param ownerMethod
+	 * 		Target method in the owner class, or {@code null} if the method is not declared in the owner class.
+	 * @param parentInfo
+	 * 		Class defining the parent method.
+	 * @param parentMethod
+	 * 		Method in the parent class with the same name/desc as the target method.
+	 *
+	 * @return {@code true} when the parent method is inherited/implemented by the owner class.
+	 * {@code false} when the parent method cannot be inherited/implemented by the owner class due to access restrictions.
+	 */
+	protected static boolean isInheritedMethod(@Nonnull ClassInfo ownerInfo, @Nullable MethodMember ownerMethod,
+	                                           @Nonnull ClassInfo parentInfo, @Nonnull MethodMember parentMethod) {
+		// Cannot inherit private/static methods.
+		if (parentMethod.hasPrivateModifier() || parentMethod.hasStaticModifier())
+			return false;
+
+		// Package-private inheritance only works if the owners are in the same package.
+		String ownerPackage = ownerInfo.getPackageName();
+		String parentPackage = parentInfo.getPackageName();
+		if (ownerMethod == null)
+			return !parentMethod.hasPackagePrivateModifier()
+					|| Objects.equals(ownerPackage, parentPackage);
+		return (!ownerMethod.hasPackagePrivateModifier() && !parentMethod.hasPackagePrivateModifier())
+				|| Objects.equals(ownerPackage, parentPackage);
+	}
+
 	/**
 	 * @param internalName
 	 * 		Some class name.
@@ -408,4 +574,33 @@ protected MappingKey getVariableKey(@Nonnull String className, @Nonnull String m
 	                                    @Nullable String name, @Nullable String desc, int index) {
 		return new VariableMappingKey(className, methodName, methodDesc, name, desc);
 	}
+
+	/**
+	 * Build the mapping conflict exception. Occurs when you have a mapping setup like:
+	 * <pre>
+	 * {@code
+	 * Shape.area() -> foo
+	 * Circle.area() -> bar
+	 * }</pre>
+	 * If {@code Circle} implements {@code Shape}, we cannot have two names for the same 'family'.
+	 *
+	 * @param key
+	 * 		Method info.
+	 * @param affectedOwner
+	 * 		Class that is affected by the conflict.
+	 * @param previous
+	 * 		Existing mapping that is causing the conflict.
+	 * @param current
+	 * 		New method mapping that is causing the conflict.
+	 *
+	 * @return Exception detailing the method mapping conflict.
+	 */
+	@Nonnull
+	private static IllegalStateException newMethodConflict(@Nonnull MethodMappingKey key, @Nonnull String affectedOwner,
+	                                                       @Nonnull MethodMapping previous, @Nonnull MethodMapping current) {
+		return new IllegalStateException("Conflicting method mapping for family '" +
+				key.getName() + key.getDesc() + "' affecting '" + affectedOwner + "': '" +
+				previous.getOwnerName() + "' -> '" + previous.getNewName() + "' conflicts with '" +
+				current.getOwnerName() + "' -> '" + current.getNewName() + "'");
+	}
 }