Fix #49

We are checking if there are methods with the same srg name that are already added, if it is, then we would not fill in these names
This is especially troublesome with methods annotated with @DontObfuscate (e.g. m_129629_), with environments like yarn where methods with the same srg name may not inherit the same names due to parameter mappings and inheritance

This requires further testing!

Signed-off-by: shedaniel <daniel@shedaniel.me>

1 files changed, 129 insertions, 75 deletions

diff --git a/src/main/java/net/fabricmc/loom/util/srg/SrgMerger.java b/src/main/java/net/fabricmc/loom/util/srg/SrgMerger.java
index 9765fa25..ede2dbad 100644
--- a/src/main/java/net/fabricmc/loom/util/srg/SrgMerger.java
+++ b/src/main/java/net/fabricmc/loom/util/srg/SrgMerger.java
@@ -27,19 +27,23 @@ package net.fabricmc.loom.util.srg;
 import java.io.BufferedReader;
 import java.io.IOException;
 import java.io.StringReader;
+import java.io.UncheckedIOException;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import com.google.common.base.MoreObjects;
+import com.google.common.base.Stopwatch;
 import org.apache.commons.io.IOUtils;
 import org.gradle.api.logging.Logger;
 import org.jetbrains.annotations.Nullable;
@@ -60,6 +64,49 @@ import net.fabricmc.mappingio.tree.MemoryMappingTree;
  * @author Juuz
  */
 public final class SrgMerger {
+	private final Logger logger;
+	private final Map<String, List<MappingTreeView.MemberMappingView>> addRegardlessSrg = new HashMap<>();
+	private final MemoryMappingTree srg;
+	private final MemoryMappingTree src;
+	private final MemoryMappingTree output;
+	private final FlatMappingVisitor flatOutput;
+	private final List<Runnable> postProcesses = new ArrayList<>();
+	private final boolean lenient;
+	private final Set<String> methodSrgNames = new HashSet<>();
+
+	public SrgMerger(Logger logger, Path srg, @Nullable Supplier<Path> mojmap, Path tiny, boolean lenient) throws IOException {
+		this.logger = logger;
+		this.srg = readSrg(srg, mojmap);
+		this.src = new MemoryMappingTree();
+		this.output = new MemoryMappingTree();
+		this.flatOutput = new RegularAsFlatMappingVisitor(output);
+		this.lenient = lenient;
+
+		MappingReader.read(tiny, this.src);
+
+		if (!"official".equals(this.src.getSrcNamespace())) {
+			throw new MappingException("Mapping file " + tiny + " does not have the 'official' namespace as the default!");
+		}
+
+		this.output.visitNamespaces(this.src.getSrcNamespace(), Stream.concat(Stream.of("srg"), this.src.getDstNamespaces().stream()).collect(Collectors.toList()));
+	}
+
+	public MemoryMappingTree merge() throws IOException {
+		for (MappingTree.ClassMapping klass : this.srg.getClasses()) {
+			classToTiny(klass);
+		}
+
+		try {
+			for (Runnable process : postProcesses) {
+				process.run();
+			}
+		} catch (UncheckedIOException e) {
+			throw e.getCause();
+		}
+
+		return output;
+	}
+
 	/**
 	 * Merges SRG mappings with a tiny mappings tree through the obf names.
 	 *
@@ -72,7 +119,9 @@ public final class SrgMerger {
 	 * @throws MappingException if the input tiny tree's default namespace is not 'official'
 	 *                          or if an element mentioned in the SRG file does not have tiny mappings
 	 */
-	public static void mergeSrg(Logger logger, @Nullable Supplier<Path> mojmap, Path srg, Path tiny, Path out, boolean lenient) throws IOException, MappingException {
+	public static void mergeSrg(Logger logger, @Nullable Supplier<Path> mojmap, Path srg, Path tiny, Path out, boolean lenient)
+			throws IOException, MappingException {
+		Stopwatch stopwatch = Stopwatch.createStarted();
 		MemoryMappingTree tree = mergeSrg(logger, mojmap, srg, tiny, lenient);
 
 		try (Tiny2Writer writer = new Tiny2Writer(Files.newBufferedWriter(out), false)) {
@@ -80,6 +129,8 @@ public final class SrgMerger {
 		} catch (IOException e) {
 			e.printStackTrace();
 		}
+
+		logger.info(":merged srg mappings in " + stopwatch.stop());
 	}
 
 	/**
@@ -94,37 +145,17 @@ public final class SrgMerger {
 	 * @throws MappingException if the input tiny tree's default namespace is not 'official'
 	 *                          or if an element mentioned in the SRG file does not have tiny mappings
 	 */
-	public static MemoryMappingTree mergeSrg(Logger logger, @Nullable Supplier<Path> mojmap, Path srg, Path tiny, boolean lenient) throws IOException, MappingException {
-		Map<String, List<MappingTreeView.MemberMappingView>> addRegardlessSrgs = new HashMap<>();
-		MemoryMappingTree arr = readSrg(srg, mojmap, addRegardlessSrgs);
-		MemoryMappingTree foss = new MemoryMappingTree();
-
-		try (BufferedReader reader = Files.newBufferedReader(tiny)) {
-			MappingReader.read(reader, foss);
-		}
-
-		if (!"official".equals(foss.getSrcNamespace())) {
-			throw new MappingException("Mapping file " + tiny + " does not have the 'official' namespace as the default!");
-		}
-
-		MemoryMappingTree output = new MemoryMappingTree();
-		output.visitNamespaces(foss.getSrcNamespace(), Stream.concat(Stream.of("srg"), foss.getDstNamespaces().stream()).collect(Collectors.toList()));
-		RegularAsFlatMappingVisitor flatMappingVisitor = new RegularAsFlatMappingVisitor(output);
-
-		for (MappingTree.ClassMapping klass : arr.getClasses()) {
-			classToTiny(logger, addRegardlessSrgs, foss, klass, output, flatMappingVisitor, lenient);
-		}
-
-		return output;
+	public static MemoryMappingTree mergeSrg(Logger logger, @Nullable Supplier<Path> mojmap, Path srg, Path tiny, boolean lenient)
+			throws IOException, MappingException {
+		return new SrgMerger(logger, srg, mojmap, tiny, lenient).merge();
 	}
 
-	private static MemoryMappingTree readSrg(Path srg, @Nullable Supplier<Path> mojmap, Map<String, List<MappingTreeView.MemberMappingView>> addRegardlessSrgs)
-			throws IOException {
+	private MemoryMappingTree readSrg(Path srg, @Nullable Supplier<Path> mojmap) throws IOException {
 		try (BufferedReader reader = Files.newBufferedReader(srg)) {
 			String content = IOUtils.toString(reader);
 
 			if (content.startsWith("tsrg2") && mojmap != null) {
-				addRegardlessSrgs(mojmap, addRegardlessSrgs);
+				addRegardlessSrgs(mojmap);
 			}
 
 			MemoryMappingTree tsrg = new MemoryMappingTree();
@@ -133,7 +164,7 @@ public final class SrgMerger {
 		}
 	}
 
-	private static void addRegardlessSrgs(Supplier<Path> mojmap, Map<String, List<MappingTreeView.MemberMappingView>> addRegardlessSrgs) throws IOException {
+	private void addRegardlessSrgs(Supplier<Path> mojmap) throws IOException {
 		MemoryMappingTree mojmapTree = readTsrg2ToTinyTree(mojmap.get());
 
 		for (MappingTree.ClassMapping classDef : mojmapTree.getClasses()) {
@@ -141,13 +172,13 @@ public final class SrgMerger {
 				String name = methodDef.getSrcName();
 
 				if (name.indexOf('<') != 0 && name.equals(methodDef.getDstName(0))) {
-					addRegardlessSrgs.computeIfAbsent(classDef.getSrcName(), $ -> new ArrayList<>()).add(methodDef);
+					addRegardlessSrg.computeIfAbsent(classDef.getSrcName(), $ -> new ArrayList<>()).add(methodDef);
 				}
 			}
 
 			for (MappingTree.FieldMapping fieldDef : classDef.getFields()) {
 				if (fieldDef.getSrcName().equals(fieldDef.getDstName(0))) {
-					addRegardlessSrgs.computeIfAbsent(classDef.getSrcName(), $ -> new ArrayList<>()).add(fieldDef);
+					addRegardlessSrg.computeIfAbsent(classDef.getSrcName(), $ -> new ArrayList<>()).add(fieldDef);
 				}
 			}
 		}
@@ -163,11 +194,10 @@ public final class SrgMerger {
 		return tree;
 	}
 
-	private static void classToTiny(Logger logger, Map<String, List<MappingTreeView.MemberMappingView>> addRegardlessSrgs, MappingTree foss, MappingTree.ClassMapping klass, MappingTree output, FlatMappingVisitor flatOutput, boolean lenient)
-			throws IOException {
+	private void classToTiny(MappingTree.ClassMapping klass) throws IOException {
 		String obf = klass.getSrcName();
 		String srg = klass.getDstName(0);
-		MappingTree.ClassMapping classDef = foss.getClass(obf);
+		MappingTree.ClassMapping classDef = this.src.getClass(obf);
 
 		if (classDef == null) {
 			if (lenient) {
@@ -194,46 +224,83 @@ public final class SrgMerger {
 					m -> m.getName("official").equals(method.getSrcName()) && m.getDesc("official").equals(method.getSrcDesc())
 			).orElse(null);
 
-			if (def == null) {
-				if (tryMatchRegardlessSrgsMethod(addRegardlessSrgs, obf, output, flatOutput, method)) continue;
+			String methodSrgName = method.getDstName(0);
 
-				if (!lenient) {
-					throw new MappingException("Missing method: " + method.getSrcName() + " (srg: " + method.getDstName(0) + ")");
+			if (def == null) {
+				if (lenient) {
+					// TODO Big Hack!
+					// We are checking if there are methods with the same srg name that are already added, if it is, then we would not fill in these names
+					// This is especially troublesome with methods annotated with @DontObfuscate (e.g. m_129629_)
+					// with environments like yarn where methods with the same srg name may not inherit the same names due to parameter mappings and inheritance
+					// This requires further testing!
+					postProcesses.add(() -> {
+						if (!methodSrgNames.contains(methodSrgName)) {
+							List<String> methodNames = CollectionUtil.map(
+									output.getDstNamespaces(),
+									namespace -> "srg".equals(namespace) ? methodSrgName : method.getSrcName()
+							);
+
+							try {
+								flatOutput.visitMethod(obf, method.getSrcName(), method.getSrcDesc(), methodNames.toArray(new String[0]));
+							} catch (IOException e) {
+								throw new UncheckedIOException(e);
+							}
+						}
+					});
+				} else {
+					throw new MappingException("Missing method: " + method.getSrcName() + " (srg: " + methodSrgName + ")");
 				}
 
 				continue;
 			}
 
-			methodToTiny(obf, method, null, def, output, flatOutput);
-		}
+			methodToTiny(obf, method, methodSrgName, def);
 
-		// TODO: This second iteration seems a bit wasteful.
-		//  Is it possible to just iterate this and leave SRG out?
-		for (MappingTree.MethodMapping def : classDef.getMethods()) {
-			// If obf = some other name: some special name that srg might not contain.
-			// This includes constructors and overridden JDK methods.
-			if (!def.getSrcName().equals(def.getDstName(0))) {
-				continue;
+			if (methodSrgName.startsWith("func_") || methodSrgName.startsWith("m_")) {
+				methodSrgNames.add(methodSrgName);
 			}
+		}
 
-			MappingTree.MethodMapping method = CollectionUtil.find(
-					klass.getMethods(),
-					m -> m.getSrcName().equals(def.getName("official")) && m.getSrcDesc().equals(def.getDesc("official"))
-			).orElse(null);
+		postProcesses.add(() -> {
+			// TODO: This second iteration seems a bit wasteful.
+			//  Is it possible to just iterate this and leave SRG out?
+			for (MappingTree.MethodMapping def : classDef.getMethods()) {
+				// If obf = some other name: some special name that srg might not contain.
+				// This includes constructors and overridden JDK methods.
+				if (!def.getSrcName().equals(def.getDstName(0))) {
+					continue;
+				}
 
-			if (method == null) {
-				methodToTiny(obf, null, def.getSrcName(), def, output, flatOutput);
+				MappingTree.MethodMapping method = CollectionUtil.find(
+						klass.getMethods(),
+						m -> m.getSrcName().equals(def.getName("official")) && m.getSrcDesc().equals(def.getDesc("official"))
+				).orElse(null);
+
+				if (method == null) {
+					try {
+						methodToTiny(obf, null, def.getSrcName(), def);
+					} catch (IOException e) {
+						throw new UncheckedIOException(e);
+					}
+				}
 			}
-		}
+		});
 
 		for (MappingTree.FieldMapping field : klass.getFields()) {
 			MappingTree.FieldMapping def = CollectionUtil.find(
 					classDef.getFields(),
 					f -> f.getName("official").equals(field.getSrcName())
-			).orElse(nullOrThrow(lenient, () -> new MappingException("Missing field: " + field.getSrcName() + " (srg: " + field.getDstName(0) + ")")));
+			).orElse(nullOrThrow(() -> new MappingException("Missing field: " + field.getSrcName() + " (srg: " + field.getDstName(0) + ")")));
 
 			if (def == null) {
-				if (tryMatchRegardlessSrgsField(addRegardlessSrgs, obf, output, flatOutput, field)) continue;
+				if (tryMatchRegardlessSrgsField(obf, field)) {
+					List<String> fieldNames = CollectionUtil.map(
+							output.getDstNamespaces(),
+							namespace -> "srg".equals(namespace) ? field.getDstName(0) : field.getSrcName()
+					);
+
+					flatOutput.visitField(obf, field.getSrcName(), field.getSrcDesc(), fieldNames.toArray(new String[0]));
+				}
 
 				continue;
 			}
@@ -251,8 +318,9 @@ public final class SrgMerger {
 		}
 	}
 
-	private static void methodToTiny(String obfClassName, @Nullable MappingTree.MethodMapping srgMethod, @Nullable String srgMethodName, MappingTree.MethodMapping actualMethod, MappingTree output, FlatMappingVisitor flatOutput) throws IOException {
-		if (srgMethod != null) {
+	private void methodToTiny(String obfClassName, @Nullable MappingTree.MethodMapping srgMethod, @Nullable String srgMethodName, MappingTree.MethodMapping actualMethod)
+			throws IOException {
+		if (srgMethod != null && srgMethodName != null) {
 			srgMethodName = srgMethod.getDstName(0);
 		}
 
@@ -299,19 +367,12 @@ public final class SrgMerger {
 		}
 	}
 
-	private static boolean tryMatchRegardlessSrgsMethod(Map<String, List<MappingTreeView.MemberMappingView>> addRegardlessSrgs, String obf,
-			MappingTree output, FlatMappingVisitor flatOutput, MappingTree.MethodMapping method) throws IOException {
-		List<MappingTreeView.MemberMappingView> mutableDescriptoredList = addRegardlessSrgs.get(obf);
+	private boolean tryMatchRegardlessSrgsMethod(String obf, MappingTree.MethodMapping method) {
+		List<MappingTreeView.MemberMappingView> mutableDescriptoredList = addRegardlessSrg.get(obf);
 
 		if (!method.getDstName(0).equals(method.getSrcName())) {
 			for (MappingTreeView.MemberMappingView descriptored : MoreObjects.firstNonNull(mutableDescriptoredList, Collections.<MappingTreeView.MemberMappingView>emptyList())) {
 				if (descriptored instanceof MappingTree.MethodMapping && descriptored.getSrcName().equals(method.getSrcName()) && descriptored.getSrcDesc().equals(method.getSrcDesc())) {
-					List<String> methodNames = CollectionUtil.map(
-							output.getDstNamespaces(),
-							namespace -> "srg".equals(namespace) ? method.getDstName(0) : method.getSrcName()
-					);
-
-					flatOutput.visitMethod(obf, method.getSrcName(), method.getSrcDesc(), methodNames.toArray(new String[0]));
 					return true;
 				}
 			}
@@ -320,19 +381,12 @@ public final class SrgMerger {
 		return false;
 	}
 
-	private static boolean tryMatchRegardlessSrgsField(Map<String, List<MappingTreeView.MemberMappingView>> addRegardlessSrgs, String obf,
-			MappingTree output, FlatMappingVisitor flatOutput, MappingTree.FieldMapping field) throws IOException {
-		List<MappingTreeView.MemberMappingView> mutableDescriptoredList = addRegardlessSrgs.get(obf);
+	private boolean tryMatchRegardlessSrgsField(String obf, MappingTree.FieldMapping field) {
+		List<MappingTreeView.MemberMappingView> mutableDescriptoredList = addRegardlessSrg.get(obf);
 
 		if (!field.getDstName(0).equals(field.getSrcName())) {
 			for (MappingTreeView.MemberMappingView descriptored : MoreObjects.firstNonNull(mutableDescriptoredList, Collections.<MappingTreeView.MemberMappingView>emptyList())) {
 				if (descriptored instanceof MappingTree.FieldMapping && descriptored.getSrcName().equals(field.getSrcName())) {
-					List<String> fieldNames = CollectionUtil.map(
-							output.getDstNamespaces(),
-							namespace -> "srg".equals(namespace) ? field.getDstName(0) : field.getSrcName()
-					);
-
-					flatOutput.visitField(obf, field.getSrcName(), field.getSrcDesc(), fieldNames.toArray(new String[0]));
 					return true;
 				}
 			}
@@ -342,7 +396,7 @@ public final class SrgMerger {
 	}
 
 	@Nullable
-	private static <T, X extends Exception> T nullOrThrow(boolean lenient, Supplier<X> exception) throws X {
+	private <T, X extends Exception> T nullOrThrow(Supplier<X> exception) throws X {
 		if (lenient) {
 			return null;
 		} else {