Primitive MCP Support!

2 files changed, 317 insertions, 3 deletions

diff --git a/src/main/java/net/fabricmc/loom/configuration/providers/mappings/MappingsProvider.java b/src/main/java/net/fabricmc/loom/configuration/providers/mappings/MappingsProvider.java
index 90488bcd..cf9fab37 100644
--- a/src/main/java/net/fabricmc/loom/configuration/providers/mappings/MappingsProvider.java
+++ b/src/main/java/net/fabricmc/loom/configuration/providers/mappings/MappingsProvider.java
@@ -31,6 +31,7 @@ import java.net.URL;
 import java.nio.file.FileSystem;
 import java.nio.file.FileSystems;
 import java.nio.file.Files;
+import java.nio.file.NoSuchFileException;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.nio.file.StandardCopyOption;
@@ -43,6 +44,7 @@ import com.google.common.net.UrlEscapers;
 import org.apache.commons.io.FileUtils;
 import org.apache.tools.ant.util.StringUtils;
 import org.gradle.api.Project;
+import org.gradle.api.artifacts.Configuration;
 import org.zeroturnaround.zip.FileSource;
 import org.zeroturnaround.zip.ZipEntrySource;
 import org.zeroturnaround.zip.ZipUtil;
@@ -54,10 +56,12 @@ import net.fabricmc.loom.configuration.processors.JarProcessorManager;
 import net.fabricmc.loom.configuration.processors.MinecraftProcessedProvider;
 import net.fabricmc.loom.configuration.providers.MinecraftProvider;
 import net.fabricmc.loom.configuration.providers.forge.MinecraftPatchedProvider;
+import net.fabricmc.loom.configuration.providers.forge.SrgProvider;
 import net.fabricmc.loom.configuration.providers.minecraft.MinecraftMappedProvider;
 import net.fabricmc.loom.util.Constants;
 import net.fabricmc.loom.util.DeletingFileVisitor;
 import net.fabricmc.loom.util.DownloadUtil;
+import net.fabricmc.loom.util.srg.MCPReader;
 import net.fabricmc.loom.util.srg.SrgMerger;
 import net.fabricmc.loom.util.srg.SrgNamedWriter;
 import net.fabricmc.mapping.reader.v2.TinyV2Factory;
@@ -66,6 +70,8 @@ import net.fabricmc.stitch.Command;
 import net.fabricmc.stitch.commands.CommandProposeFieldNames;
 import net.fabricmc.stitch.commands.tinyv2.CommandMergeTinyV2;
 import net.fabricmc.stitch.commands.tinyv2.CommandReorderTinyV2;
+import net.fabricmc.stitch.commands.tinyv2.TinyFile;
+import net.fabricmc.stitch.commands.tinyv2.TinyV2Writer;
 
 public class MappingsProvider extends DependencyProvider {
 	public MinecraftMappedProvider mappedProvider;
@@ -125,6 +131,13 @@ public class MappingsProvider extends DependencyProvider {
 
 		boolean isV2;
 
+		if (isMCP(mappingsJar.toPath())) {
+			File old = mappingsJar;
+			mappingsJar = mappingsDir.resolve(StringUtils.removeSuffix(mappingsJar.getName(), ".zip") + "-" + minecraftVersion + ".jar").toFile();
+			FileUtils.copyFile(old, mappingsJar);
+			mappingsName += "-" + minecraftVersion;
+		}
+
 		// Only do this for official yarn, there isn't really a way we can get the mc version for all mappings
 		if (dependency.getDependency().getGroup() != null && dependency.getDependency().getGroup().equals("net.fabricmc") && dependency.getDependency().getName().equals("yarn") && dependency.getDependency().getVersion() != null) {
 			String yarnVersion = dependency.getDependency().getVersion();
@@ -166,7 +179,7 @@ public class MappingsProvider extends DependencyProvider {
 		srgToNamedSrg = mappingsDir.resolve(StringUtils.removeSuffix(mappingsJar.getName(), ".jar") + "-srg-named.srg").toFile();
 
 		if (!tinyMappings.exists() || isRefreshDeps()) {
-			storeMappings(getProject(), minecraftProvider, mappingsJar.toPath());
+			storeMappings(getProject(), minecraftProvider, mappingsJar.toPath(), postPopulationScheduler);
 		}
 
 		if (!tinyMappingsJar.exists() || isRefreshDeps()) {
@@ -224,9 +237,15 @@ public class MappingsProvider extends DependencyProvider {
 		mappedProvider.provide(dependency, postPopulationScheduler);
 	}
 
-	private void storeMappings(Project project, MinecraftProvider minecraftProvider, Path yarnJar) throws IOException {
+	private void storeMappings(Project project, MinecraftProvider minecraftProvider, Path yarnJar, Consumer<Runnable> postPopulationScheduler)
+			throws Exception {
 		project.getLogger().lifecycle(":extracting " + yarnJar.getFileName());
 
+		if (isMCP(yarnJar)) {
+			readAndMergeMCP(yarnJar, postPopulationScheduler);
+			return;
+		}
+
 		try (FileSystem fileSystem = FileSystems.newFileSystem(yarnJar, (ClassLoader) null)) {
 			extractMappings(fileSystem, baseTinyMappings);
 		}
@@ -245,6 +264,34 @@ public class MappingsProvider extends DependencyProvider {
 		}
 	}
 
+	private void readAndMergeMCP(Path mcpJar, Consumer<Runnable> postPopulationScheduler) throws Exception {
+		Path intermediaryTinyPath = getIntermediaryTiny();
+		SrgProvider provider = getExtension().getSrgProvider();
+
+		if (provider == null) {
+			if (!getExtension().shouldGenerateSrgTiny()) {
+				Configuration srg = getProject().getConfigurations().maybeCreate(Constants.Configurations.SRG);
+				srg.setTransitive(false);
+			}
+
+			provider = new SrgProvider(getProject());
+			getProject().getDependencies().add(provider.getTargetConfig(), "de.oceanlabs.mcp:mcp_config:" + minecraftVersion);
+			Configuration configuration = getProject().getConfigurations().getByName(provider.getTargetConfig());
+			provider.provide(DependencyInfo.create(getProject(), configuration.getDependencies().iterator().next(), configuration), postPopulationScheduler);
+		}
+
+		Path srgPath = provider.getSrg().toPath();
+
+		TinyFile file = new MCPReader(intermediaryTinyPath, srgPath).read(mcpJar);
+		TinyV2Writer.write(file, tinyMappings.toPath());
+	}
+
+	private boolean isMCP(Path path) throws IOException {
+		try (FileSystem fs = FileSystems.newFileSystem(path, (ClassLoader) null)) {
+			return Files.exists(fs.getPath("fields.csv")) && Files.exists(fs.getPath("methods.csv"));
+		}
+	}
+
 	private boolean baseMappingsAreV2() throws IOException {
 		try (BufferedReader reader = Files.newBufferedReader(baseTinyMappings)) {
 			TinyV2Factory.readMetadata(reader);
@@ -260,7 +307,7 @@ public class MappingsProvider extends DependencyProvider {
 			try (BufferedReader reader = Files.newBufferedReader(fs.getPath("mappings", "mappings.tiny"))) {
 				TinyV2Factory.readMetadata(reader);
 				return true;
-			} catch (IllegalArgumentException e) {
+			} catch (IllegalArgumentException | NoSuchFileException e) {
 				return false;
 			}
 		}
diff --git a/src/main/java/net/fabricmc/loom/util/srg/MCPReader.java b/src/main/java/net/fabricmc/loom/util/srg/MCPReader.java
new file mode 100644
index 00000000..cae61758
--- /dev/null
+++ b/src/main/java/net/fabricmc/loom/util/srg/MCPReader.java
@@ -0,0 +1,267 @@
+/*
+ * This file is part of fabric-loom, licensed under the MIT License (MIT).
+ *
+ * Copyright (c) 2016, 2017, 2018 FabricMC
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package net.fabricmc.loom.util.srg;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.FileSystem;
+import java.nio.file.FileSystems;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+import au.com.bytecode.opencsv.CSVReader;
+import org.cadixdev.lorenz.MappingSet;
+import org.cadixdev.lorenz.io.srg.tsrg.TSrgReader;
+import org.cadixdev.lorenz.model.ClassMapping;
+import org.cadixdev.lorenz.model.FieldMapping;
+import org.cadixdev.lorenz.model.InnerClassMapping;
+import org.cadixdev.lorenz.model.MethodMapping;
+import org.cadixdev.lorenz.model.TopLevelClassMapping;
+import org.jetbrains.annotations.Nullable;
+
+import net.fabricmc.stitch.commands.tinyv2.TinyClass;
+import net.fabricmc.stitch.commands.tinyv2.TinyField;
+import net.fabricmc.stitch.commands.tinyv2.TinyFile;
+import net.fabricmc.stitch.commands.tinyv2.TinyMethod;
+import net.fabricmc.stitch.commands.tinyv2.TinyV2Reader;
+
+public class MCPReader {
+	private final Path intermediaryTinyPath;
+	private final Path srgTsrgPath;
+
+	public MCPReader(Path intermediaryTinyPath, Path srgTsrgPath) {
+		this.intermediaryTinyPath = intermediaryTinyPath;
+		this.srgTsrgPath = srgTsrgPath;
+	}
+
+	public TinyFile read(Path mcpJar) throws IOException {
+		Map<MemberToken, String> srgTokens = readSrg();
+		TinyFile intermediaryTiny = TinyV2Reader.read(intermediaryTinyPath);
+		Map<String, String> intermediaryToMCPMap = createIntermediaryToMCPMap(intermediaryTiny, srgTokens);
+		injectMcp(mcpJar, intermediaryToMCPMap);
+
+		mergeTokensIntoIntermediary(intermediaryTiny, intermediaryToMCPMap);
+		return intermediaryTiny;
+	}
+
+	private Map<String, String> createIntermediaryToMCPMap(TinyFile tiny, Map<MemberToken, String> officialToMCP) {
+		Map<String, String> map = new HashMap<>();
+
+		for (TinyClass tinyClass : tiny.getClassEntries()) {
+			String classObf = tinyClass.getMapping().get(0);
+			String classIntermediary = tinyClass.getMapping().get(1);
+			MemberToken classTokenObf = MemberToken.ofClass(classObf);
+
+			if (officialToMCP.containsKey(classTokenObf)) {
+				map.put(classIntermediary, officialToMCP.get(classTokenObf));
+			}
+
+			for (TinyField tinyField : tinyClass.getFields()) {
+				String fieldObf = tinyField.getMapping().get(0);
+				String fieldIntermediary = tinyField.getMapping().get(1);
+				MemberToken fieldTokenObf = MemberToken.ofField(classTokenObf, fieldObf);
+
+				if (officialToMCP.containsKey(fieldTokenObf)) {
+					map.put(fieldIntermediary, officialToMCP.get(fieldTokenObf));
+				}
+			}
+
+			for (TinyMethod tinyMethod : tinyClass.getMethods()) {
+				String methodObf = tinyMethod.getMapping().get(0);
+				String methodIntermediary = tinyMethod.getMapping().get(1);
+				MemberToken methodTokenObf = MemberToken.ofMethod(classTokenObf, methodObf, tinyMethod.getMethodDescriptorInFirstNamespace());
+
+				if (officialToMCP.containsKey(methodTokenObf)) {
+					map.put(methodIntermediary, officialToMCP.get(methodTokenObf));
+				}
+			}
+		}
+
+		return map;
+	}
+
+	private void mergeTokensIntoIntermediary(TinyFile tiny, Map<String, String> intermediaryToMCPMap) {
+		stripTinyWithParametersAndLocal(tiny);
+
+		// We will be adding the "named" namespace with MCP
+		tiny.getHeader().getNamespaces().add("named");
+
+		for (TinyClass tinyClass : tiny.getClassEntries()) {
+			String classIntermediary = tinyClass.getMapping().get(1);
+			tinyClass.getMapping().add(intermediaryToMCPMap.getOrDefault(classIntermediary, classIntermediary));
+
+			for (TinyField tinyField : tinyClass.getFields()) {
+				String fieldIntermediary = tinyField.getMapping().get(1);
+				tinyField.getMapping().add(intermediaryToMCPMap.getOrDefault(fieldIntermediary, fieldIntermediary));
+			}
+
+			for (TinyMethod tinyMethod : tinyClass.getMethods()) {
+				String methodIntermediary = tinyMethod.getMapping().get(1);
+				tinyMethod.getMapping().add(intermediaryToMCPMap.getOrDefault(methodIntermediary, methodIntermediary));
+			}
+		}
+	}
+
+	private void stripTinyWithParametersAndLocal(TinyFile tiny) {
+		for (TinyClass tinyClass : tiny.getClassEntries()) {
+			for (TinyMethod tinyMethod : tinyClass.getMethods()) {
+				tinyMethod.getParameters().clear();
+				tinyMethod.getLocalVariables().clear();
+			}
+		}
+	}
+
+	private Map<MemberToken, String> readSrg() throws IOException {
+		Map<MemberToken, String> tokens = new HashMap<>();
+
+		try (TSrgReader reader = new TSrgReader(Files.newBufferedReader(srgTsrgPath, StandardCharsets.UTF_8))) {
+			MappingSet mappingSet = reader.read();
+
+			for (TopLevelClassMapping classMapping : mappingSet.getTopLevelClassMappings()) {
+				appendClass(tokens, classMapping);
+			}
+		}
+
+		return tokens;
+	}
+
+	private void injectMcp(Path mcpJar, Map<String, String> intermediaryToMCPMap) throws IOException {
+		Map<String, List<String>> inverseMap = inverseMap(intermediaryToMCPMap);
+
+		try (FileSystem fs = FileSystems.newFileSystem(mcpJar, null)) {
+			Path fields = fs.getPath("fields.csv");
+			Path methods = fs.getPath("methods.csv");
+
+			try (CSVReader reader = new CSVReader(Files.newBufferedReader(fields, StandardCharsets.UTF_8))) {
+				reader.readNext();
+				String[] line;
+
+				while ((line = reader.readNext()) != null) {
+					List<String> intermediaryField = inverseMap.get(line[0]);
+
+					if (intermediaryField != null) {
+						for (String s : intermediaryField) {
+							intermediaryToMCPMap.put(s, line[1]);
+						}
+					}
+				}
+			}
+
+			try (CSVReader reader = new CSVReader(Files.newBufferedReader(methods, StandardCharsets.UTF_8))) {
+				reader.readNext();
+				String[] line;
+
+				while ((line = reader.readNext()) != null) {
+					List<String> intermediaryMethod = inverseMap.get(line[0]);
+
+					if (intermediaryMethod != null) {
+						for (String s : intermediaryMethod) {
+							intermediaryToMCPMap.put(s, line[1]);
+						}
+					}
+				}
+			}
+		}
+	}
+
+	private Map<String, List<String>> inverseMap(Map<String, String> intermediaryToMCPMap) {
+		Map<String, List<String>> map = new HashMap<>();
+
+		for (Map.Entry<String, String> token : intermediaryToMCPMap.entrySet()) {
+			map.computeIfAbsent(token.getValue(), s -> new ArrayList<>()).add(token.getKey());
+		}
+
+		return map;
+	}
+
+	private void appendClass(Map<MemberToken, String> tokens, ClassMapping<?, ?> classMapping) {
+		MemberToken ofClass = MemberToken.ofClass(classMapping.getFullObfuscatedName());
+		tokens.put(ofClass, classMapping.getFullDeobfuscatedName());
+
+		for (FieldMapping fieldMapping : classMapping.getFieldMappings()) {
+			tokens.put(MemberToken.ofField(ofClass, fieldMapping.getObfuscatedName()), fieldMapping.getDeobfuscatedName());
+		}
+
+		for (MethodMapping methodMapping : classMapping.getMethodMappings()) {
+			tokens.put(MemberToken.ofMethod(ofClass, methodMapping.getObfuscatedName(), methodMapping.getObfuscatedDescriptor()), methodMapping.getDeobfuscatedName());
+		}
+
+		for (InnerClassMapping mapping : classMapping.getInnerClassMappings()) {
+			appendClass(tokens, mapping);
+		}
+	}
+
+	private static class MemberToken {
+		private final TokenType type;
+		@Nullable
+		private final MemberToken owner;
+		private final String name;
+		@Nullable
+		private final String descriptor;
+
+		MemberToken(TokenType type, @Nullable MemberToken owner, String name, @Nullable String descriptor) {
+			this.type = type;
+			this.owner = owner;
+			this.name = name;
+			this.descriptor = descriptor;
+		}
+
+		static MemberToken ofClass(String name) {
+			return new MemberToken(TokenType.CLASS, null, name, null);
+		}
+
+		static MemberToken ofField(MemberToken owner, String name) {
+			return new MemberToken(TokenType.FIELD, owner, name, null);
+		}
+
+		static MemberToken ofMethod(MemberToken owner, String name, String descriptor) {
+			return new MemberToken(TokenType.METHOD, owner, name, descriptor);
+		}
+
+		@Override
+		public boolean equals(Object o) {
+			if (this == o) return true;
+			if (o == null || getClass() != o.getClass()) return false;
+			MemberToken that = (MemberToken) o;
+			return type == that.type && name.equals(that.name) && Objects.equals(descriptor, that.descriptor) && Objects.equals(owner, that.owner);
+		}
+
+		@Override
+		public int hashCode() {
+			return Objects.hash(type, name, descriptor, owner);
+		}
+	}
+
+	private enum TokenType {
+		CLASS,
+		METHOD,
+		FIELD
+	}
+}