Constant unpicking (#328)

* Start adding constant unpicking

* Update to use unpick cli

* Fix build?

* Fix?

* Fix log spam when unpicking

* Improve unpick tests

8 files changed, 257 insertions, 14 deletions

diff --git a/src/main/java/net/fabricmc/loom/LoomGradleExtension.java b/src/main/java/net/fabricmc/loom/LoomGradleExtension.java
index ea089af0..3234a628 100644
--- a/src/main/java/net/fabricmc/loom/LoomGradleExtension.java
+++ b/src/main/java/net/fabricmc/loom/LoomGradleExtension.java
@@ -354,6 +354,10 @@ public class LoomGradleExtension {
 		return new File(getProjectPersistentCache(), "log4j.xml");
 	}
 
+	public File getUnpickLoggingConfigFile() {
+		return new File(getProjectPersistentCache(), "unpick-logging.properties");
+	}
+
 	public ConfigurableFileCollection getLog4jConfigs() {
 		return log4jConfigs;
 	}
diff --git a/src/main/java/net/fabricmc/loom/configuration/CompileConfiguration.java b/src/main/java/net/fabricmc/loom/configuration/CompileConfiguration.java
index e6e7787e..6552089b 100644
--- a/src/main/java/net/fabricmc/loom/configuration/CompileConfiguration.java
+++ b/src/main/java/net/fabricmc/loom/configuration/CompileConfiguration.java
@@ -64,9 +64,13 @@ public final class CompileConfiguration {
 		Configuration includeConfig = project.getConfigurations().maybeCreate(Constants.Configurations.INCLUDE);
 		includeConfig.setTransitive(false); // Dont get transitive deps
 
+		project.getConfigurations().maybeCreate(Constants.Configurations.MAPPING_CONSTANTS);
+		extendsFrom(JavaPlugin.COMPILE_ONLY_CONFIGURATION_NAME, Constants.Configurations.MAPPING_CONSTANTS, project);
+
 		project.getConfigurations().maybeCreate(Constants.Configurations.MAPPINGS);
 		project.getConfigurations().maybeCreate(Constants.Configurations.MAPPINGS_FINAL);
 		project.getConfigurations().maybeCreate(Constants.Configurations.LOOM_DEVELOPMENT_DEPENDENCIES);
+		project.getConfigurations().maybeCreate(Constants.Configurations.UNPICK_CLASSPATH);
 
 		for (RemappedConfigurationEntry entry : Constants.MOD_COMPILE_ENTRIES) {
 			Configuration compileModsConfig = project.getConfigurations().maybeCreate(entry.getSourceConfiguration());
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 8b19a3da..448082b9 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
@@ -28,6 +28,7 @@ import java.io.BufferedReader;
 import java.io.File;
 import java.io.IOException;
 import java.net.URL;
+import java.nio.charset.StandardCharsets;
 import java.nio.file.FileSystem;
 import java.nio.file.FileSystems;
 import java.nio.file.Files;
@@ -38,6 +39,7 @@ import java.util.function.Consumer;
 
 import com.google.common.base.Preconditions;
 import com.google.common.net.UrlEscapers;
+import com.google.gson.JsonObject;
 import org.apache.commons.io.FileUtils;
 import org.apache.tools.ant.util.StringUtils;
 import org.gradle.api.Project;
@@ -46,6 +48,7 @@ import org.zeroturnaround.zip.ZipEntrySource;
 import org.zeroturnaround.zip.ZipUtil;
 
 import net.fabricmc.loom.LoomGradleExtension;
+import net.fabricmc.loom.LoomGradlePlugin;
 import net.fabricmc.loom.configuration.DependencyProvider;
 import net.fabricmc.loom.configuration.accesswidener.AccessWidenerJarProcessor;
 import net.fabricmc.loom.configuration.processors.JarProcessorManager;
@@ -78,6 +81,9 @@ public class MappingsProvider extends DependencyProvider {
 	// The mappings we use in practice
 	public File tinyMappings;
 	public File tinyMappingsJar;
+	private File unpickDefinitionsFile;
+	private boolean hasUnpickDefinitions;
+	private UnpickMetadata unpickMetadata;
 
 	public MappingsProvider(Project project) {
 		super(project);
@@ -142,16 +148,32 @@ public class MappingsProvider extends DependencyProvider {
 		}
 
 		tinyMappings = mappingsDir.resolve(StringUtils.removeSuffix(mappingsJar.getName(), ".jar") + ".tiny").toFile();
+		unpickDefinitionsFile = mappingsDir.resolve(StringUtils.removeSuffix(mappingsJar.getName(), ".jar") + ".unpick").toFile();
 		tinyMappingsJar = new File(getExtension().getUserCache(), mappingsJar.getName().replace(".jar", "-" + jarClassifier + ".jar"));
 
 		if (!tinyMappings.exists() || isRefreshDeps()) {
 			storeMappings(getProject(), minecraftProvider, mappingsJar.toPath());
+		} else {
+			try (FileSystem fileSystem = FileSystems.newFileSystem(mappingsJar.toPath(), (ClassLoader) null)) {
+				extractUnpickDefinitions(fileSystem, unpickDefinitionsFile.toPath());
+			}
 		}
 
 		if (!tinyMappingsJar.exists() || isRefreshDeps()) {
 			ZipUtil.pack(new ZipEntrySource[] {new FileSource("mappings/mappings.tiny", tinyMappings)}, tinyMappingsJar);
 		}
 
+		if (hasUnpickDefinitions()) {
+			String notation = String.format("%s:%s:%s:constants",
+					dependency.getDependency().getGroup(),
+					dependency.getDependency().getName(),
+					dependency.getDependency().getVersion()
+			);
+
+			getProject().getDependencies().add(Constants.Configurations.MAPPING_CONSTANTS, notation);
+			populateUnpickClasspath();
+		}
+
 		addDependency(tinyMappingsJar, Constants.Configurations.MAPPINGS_FINAL);
 
 		LoomGradleExtension extension = getExtension();
@@ -180,6 +202,7 @@ public class MappingsProvider extends DependencyProvider {
 
 		try (FileSystem fileSystem = FileSystems.newFileSystem(yarnJar, (ClassLoader) null)) {
 			extractMappings(fileSystem, baseTinyMappings);
+			extractUnpickDefinitions(fileSystem, unpickDefinitionsFile.toPath());
 		}
 
 		if (baseMappingsAreV2()) {
@@ -221,6 +244,40 @@ public class MappingsProvider extends DependencyProvider {
 		Files.copy(jar.getPath("mappings/mappings.tiny"), extractTo, StandardCopyOption.REPLACE_EXISTING);
 	}
 
+	private void extractUnpickDefinitions(FileSystem jar, Path extractTo) throws IOException {
+		Path unpickPath = jar.getPath("extras/definitions.unpick");
+		Path unpickMetadataPath = jar.getPath("extras/unpick.json");
+
+		if (!Files.exists(unpickPath) || !Files.exists(unpickMetadataPath)) {
+			return;
+		}
+
+		Files.copy(unpickPath, extractTo, StandardCopyOption.REPLACE_EXISTING);
+
+		unpickMetadata = parseUnpickMetadata(unpickMetadataPath);
+		hasUnpickDefinitions = true;
+	}
+
+	private UnpickMetadata parseUnpickMetadata(Path input) throws IOException {
+		JsonObject jsonObject = LoomGradlePlugin.GSON.fromJson(new String(Files.readAllBytes(input), StandardCharsets.UTF_8), JsonObject.class);
+
+		if (!jsonObject.has("version") || jsonObject.get("version").getAsInt() != 1) {
+			throw new UnsupportedOperationException("Unsupported unpick version");
+		}
+
+		return new UnpickMetadata(
+				jsonObject.get("unpickGroup").getAsString(),
+				jsonObject.get("unpickVersion").getAsString()
+		);
+	}
+
+	private void populateUnpickClasspath() {
+		String unpickCliName = "unpick-cli";
+		getProject().getDependencies().add(Constants.Configurations.UNPICK_CLASSPATH,
+				String.format("%s:%s:%s", unpickMetadata.unpickGroup, unpickCliName, unpickMetadata.unpickVersion)
+		);
+	}
+
 	private void extractIntermediary(Path intermediaryJar, Path intermediaryTiny) throws IOException {
 		getProject().getLogger().info(":extracting " + intermediaryJar.getFileName());
 
@@ -343,4 +400,22 @@ public class MappingsProvider extends DependencyProvider {
 	public String getMappingsKey() {
 		return mappingsName + "." + minecraftVersion.replace(' ', '_').replace('.', '_').replace('-', '_') + "." + mappingsVersion;
 	}
+
+	public File getUnpickDefinitionsFile() {
+		return unpickDefinitionsFile;
+	}
+
+	public boolean hasUnpickDefinitions() {
+		return hasUnpickDefinitions;
+	}
+
+	public static class UnpickMetadata {
+		public final String unpickGroup;
+		public final String unpickVersion;
+
+		public UnpickMetadata(String unpickGroup, String unpickVersion) {
+			this.unpickGroup = unpickGroup;
+			this.unpickVersion = unpickVersion;
+		}
+	}
 }
diff --git a/src/main/java/net/fabricmc/loom/configuration/providers/minecraft/MinecraftMappedProvider.java b/src/main/java/net/fabricmc/loom/configuration/providers/minecraft/MinecraftMappedProvider.java
index 6301050a..0a61f2ba 100644
--- a/src/main/java/net/fabricmc/loom/configuration/providers/minecraft/MinecraftMappedProvider.java
+++ b/src/main/java/net/fabricmc/loom/configuration/providers/minecraft/MinecraftMappedProvider.java
@@ -169,6 +169,10 @@ public class MinecraftMappedProvider extends DependencyProvider {
 		return minecraftMappedJar;
 	}
 
+	public File getUnpickedJar() {
+		return new File(getJarDirectory(getExtension().getUserCache(), "mapped"), "minecraft-" + getJarVersionString("unpicked") + ".jar");
+	}
+
 	@Override
 	public String getTargetConfig() {
 		return Constants.Configurations.MINECRAFT_NAMED;
diff --git a/src/main/java/net/fabricmc/loom/task/GenerateSourcesTask.java b/src/main/java/net/fabricmc/loom/task/GenerateSourcesTask.java
index 7c5862e4..25d6fe47 100644
--- a/src/main/java/net/fabricmc/loom/task/GenerateSourcesTask.java
+++ b/src/main/java/net/fabricmc/loom/task/GenerateSourcesTask.java
@@ -35,6 +35,7 @@ import java.util.stream.Collectors;
 
 import javax.inject.Inject;
 
+import org.gradle.api.tasks.InputFile;
 import org.gradle.api.tasks.TaskAction;
 
 import net.fabricmc.loom.LoomGradleExtension;
@@ -49,6 +50,8 @@ import net.fabricmc.stitch.util.StitchUtil;
 public class GenerateSourcesTask extends AbstractLoomTask {
 	public final LoomDecompiler decompiler;
 
+	private File inputJar;
+
 	@Inject
 	public GenerateSourcesTask(LoomDecompiler decompiler) {
 		this.decompiler = decompiler;
@@ -65,26 +68,18 @@ public class GenerateSourcesTask extends AbstractLoomTask {
 						.stream().map(File::toPath).collect(Collectors.toSet());
 
 		DecompilationMetadata metadata = new DecompilationMetadata(threads, javaDocs, libraries);
-		Path compiledJar = getExtension().getMappingsProvider().mappedProvider.getMappedJar().toPath();
+		Path runtimeJar = getExtension().getMappingsProvider().mappedProvider.getMappedJar().toPath();
 		Path sourcesDestination = getMappedJarFileWithSuffix("-sources.jar").toPath();
 		Path linemap = getMappedJarFileWithSuffix("-sources.lmap").toPath();
-		decompiler.decompile(compiledJar, sourcesDestination, linemap, metadata);
+		decompiler.decompile(inputJar.toPath(), sourcesDestination, linemap, metadata);
 
 		if (Files.exists(linemap)) {
 			Path linemappedJarDestination = getMappedJarFileWithSuffix("-linemapped.jar").toPath();
 
-			remapLineNumbers(compiledJar, linemap, linemappedJarDestination);
-
-			// In order for IDEs to recognize the new line mappings, we need to overwrite the existing compiled jar
-			// with the linemapped one. In the name of not destroying the existing jar, we will copy it to somewhere else.
-			Path unlinemappedJar = getMappedJarFileWithSuffix("-unlinemapped.jar").toPath();
+			// Line map the actually jar used to run the game, not the one used to decompile
+			remapLineNumbers(runtimeJar, linemap, linemappedJarDestination);
 
-			// The second time genSources is ran, we want to keep the existing unlinemapped jar.
-			if (!Files.exists(unlinemappedJar)) {
-				Files.copy(compiledJar, unlinemappedJar);
-			}
-
-			Files.copy(linemappedJarDestination, compiledJar, StandardCopyOption.REPLACE_EXISTING);
+			Files.copy(linemappedJarDestination, runtimeJar, StandardCopyOption.REPLACE_EXISTING);
 			Files.delete(linemappedJarDestination);
 		}
 	}
@@ -117,4 +112,14 @@ public class GenerateSourcesTask extends AbstractLoomTask {
 
 		return new File(path.substring(0, path.length() - 4) + suffix);
 	}
+
+	@InputFile
+	public File getInputJar() {
+		return inputJar;
+	}
+
+	public GenerateSourcesTask setInputJar(File inputJar) {
+		this.inputJar = inputJar;
+		return this;
+	}
 }
diff --git a/src/main/java/net/fabricmc/loom/task/LoomTasks.java b/src/main/java/net/fabricmc/loom/task/LoomTasks.java
index 198f92e6..1b1de6a7 100644
--- a/src/main/java/net/fabricmc/loom/task/LoomTasks.java
+++ b/src/main/java/net/fabricmc/loom/task/LoomTasks.java
@@ -24,6 +24,8 @@
 
 package net.fabricmc.loom.task;
 
+import java.io.File;
+
 import com.google.common.base.Preconditions;
 import org.gradle.api.Project;
 import org.gradle.api.tasks.TaskContainer;
@@ -31,6 +33,7 @@ import org.gradle.api.tasks.TaskContainer;
 import net.fabricmc.loom.LoomGradleExtension;
 import net.fabricmc.loom.api.decompilers.LoomDecompiler;
 import net.fabricmc.loom.configuration.ide.RunConfigSettings;
+import net.fabricmc.loom.configuration.providers.mappings.MappingsProvider;
 import net.fabricmc.loom.decompilers.fernflower.FabricFernFlowerDecompiler;
 
 public final class LoomTasks {
@@ -110,10 +113,30 @@ public final class LoomTasks {
 		LoomGradleExtension extension = project.getExtensions().getByType(LoomGradleExtension.class);
 
 		project.afterEvaluate(p -> {
+			MappingsProvider mappingsProvider = extension.getMappingsProvider();
+			File inputJar = mappingsProvider.mappedProvider.getMappedJar();
+
+			if (mappingsProvider.hasUnpickDefinitions()) {
+				File outputJar = mappingsProvider.mappedProvider.getUnpickedJar();
+
+				tasks.register("unpickJar", UnpickJarTask.class, unpickJarTask -> {
+					unpickJarTask.setUnpickDefinition(mappingsProvider.getUnpickDefinitionsFile());
+					unpickJarTask.setInputJar(mappingsProvider.mappedProvider.getMappedJar());
+					unpickJarTask.setOutputJar(outputJar);
+				});
+
+				inputJar = outputJar;
+			}
+
 			for (LoomDecompiler decompiler : extension.getDecompilers()) {
 				String taskName = decompiler instanceof FabricFernFlowerDecompiler ? "genSources" : "genSourcesWith" + decompiler.name();
 				// decompiler will be passed to the constructor of GenerateSourcesTask
-				tasks.register(taskName, GenerateSourcesTask.class, decompiler);
+				GenerateSourcesTask generateSourcesTask = tasks.register(taskName, GenerateSourcesTask.class, decompiler).get();
+				generateSourcesTask.setInputJar(inputJar);
+
+				if (mappingsProvider.hasUnpickDefinitions()) {
+					generateSourcesTask.dependsOn(tasks.getByName("unpickJar"));
+				}
 			}
 		});
 	}
diff --git a/src/main/java/net/fabricmc/loom/task/UnpickJarTask.java b/src/main/java/net/fabricmc/loom/task/UnpickJarTask.java
new file mode 100644
index 00000000..e27d2ff2
--- /dev/null
+++ b/src/main/java/net/fabricmc/loom/task/UnpickJarTask.java
@@ -0,0 +1,126 @@
+/*
+ * 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.task;
+
+import java.io.File;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.file.Files;
+
+import org.gradle.api.tasks.InputFile;
+import org.gradle.api.tasks.Internal;
+import org.gradle.api.tasks.JavaExec;
+import org.gradle.api.tasks.OutputFile;
+
+import net.fabricmc.loom.LoomGradleExtension;
+import net.fabricmc.loom.configuration.providers.LaunchProvider;
+import net.fabricmc.loom.util.Constants;
+
+public class UnpickJarTask extends JavaExec {
+	File inputJar;
+	File unpickDefinition;
+
+	File outputJar;
+
+	public UnpickJarTask() {
+		getOutputs().upToDateWhen(e -> false);
+		classpath(getProject().getConfigurations().getByName(Constants.Configurations.UNPICK_CLASSPATH));
+		setMain("daomephsta.unpick.cli.Main");
+	}
+
+	@Override
+	public void exec() {
+		fileArg(getInputJar(), getOutputJar(), getUnpickDefinition());
+		fileArg(getConstantJar());
+
+		// Classpath
+		fileArg(getExtension().getMinecraftMappedProvider().getMappedJar());
+		fileArg(getMinecraftDependencies());
+
+		writeUnpickLogConfig();
+		systemProperty("java.util.logging.config.file", getExtension().getUnpickLoggingConfigFile().getAbsolutePath());
+
+		super.exec();
+	}
+
+	private void writeUnpickLogConfig() {
+		try (InputStream is = LaunchProvider.class.getClassLoader().getResourceAsStream("unpick-logging.properties")) {
+			Files.deleteIfExists(getExtension().getUnpickLoggingConfigFile().toPath());
+			Files.copy(is, getExtension().getUnpickLoggingConfigFile().toPath());
+		} catch (IOException e) {
+			throw new RuntimeException("Failed to copy unpick logging config", e);
+		}
+	}
+
+	private File[] getMinecraftDependencies() {
+		return getProject().getConfigurations().getByName(Constants.Configurations.MINECRAFT_DEPENDENCIES)
+				.resolve().toArray(new File[0]);
+	}
+
+	private File getConstantJar() {
+		return getProject().getConfigurations().getByName(Constants.Configurations.MAPPING_CONSTANTS).getSingleFile();
+	}
+
+	@InputFile
+	public File getInputJar() {
+		return inputJar;
+	}
+
+	public UnpickJarTask setInputJar(File inputJar) {
+		this.inputJar = inputJar;
+		return this;
+	}
+
+	@InputFile
+	public File getUnpickDefinition() {
+		return unpickDefinition;
+	}
+
+	public UnpickJarTask setUnpickDefinition(File unpickDefinition) {
+		this.unpickDefinition = unpickDefinition;
+		return this;
+	}
+
+	@OutputFile
+	public File getOutputJar() {
+		return outputJar;
+	}
+
+	public UnpickJarTask setOutputJar(File outputJar) {
+		this.outputJar = outputJar;
+		return this;
+	}
+
+	private void fileArg(File... files) {
+		for (File file : files) {
+			args(file.getAbsolutePath());
+		}
+	}
+
+	@Internal
+	protected LoomGradleExtension getExtension() {
+		return getProject().getExtensions().getByType(LoomGradleExtension.class);
+	}
+}
diff --git a/src/main/java/net/fabricmc/loom/util/Constants.java b/src/main/java/net/fabricmc/loom/util/Constants.java
index fb4ac480..beff10df 100644
--- a/src/main/java/net/fabricmc/loom/util/Constants.java
+++ b/src/main/java/net/fabricmc/loom/util/Constants.java
@@ -79,6 +79,8 @@ public class Constants {
 		public static final String LOOM_DEVELOPMENT_DEPENDENCIES = "loomDevelopmentDependencies";
 		@Deprecated // Not to be used in gradle 7+
 		public static final String COMPILE = "compile";
+		public static final String MAPPING_CONSTANTS = "mappingsConstants";
+		public static final String UNPICK_CLASSPATH = "unpick";
 
 		private Configurations() {
 		}