Properly reapply changed ATs, Improve ATs performance.

3 files changed, 124 insertions, 63 deletions

diff --git a/src/main/java/net/fabricmc/loom/providers/MappingsProvider.java b/src/main/java/net/fabricmc/loom/providers/MappingsProvider.java
index 2f45c289..b6d537f7 100644
--- a/src/main/java/net/fabricmc/loom/providers/MappingsProvider.java
+++ b/src/main/java/net/fabricmc/loom/providers/MappingsProvider.java
@@ -201,6 +201,7 @@ public class MappingsProvider extends DependencyProvider {
 
 		if (extension.isForge()) {
 			patchedProvider = new MinecraftPatchedProvider(this, getProject());
+			patchedProvider.provide(dependency, postPopulationScheduler);
 		}
 
 		if (processorManager.active() || (extension.isForge() && patchedProvider.usesProjectCache())) {
@@ -210,9 +211,6 @@ public class MappingsProvider extends DependencyProvider {
 			mappedProvider = new MinecraftMappedProvider(getProject());
 		}
 
-		if (extension.isForge()) {
-			patchedProvider.provide(dependency, postPopulationScheduler);
-		}
 		mappedProvider.initFiles(minecraftProvider, this);
 		mappedProvider.provide(dependency, postPopulationScheduler);
 	}
diff --git a/src/main/java/net/fabricmc/loom/providers/MinecraftPatchedProvider.java b/src/main/java/net/fabricmc/loom/providers/MinecraftPatchedProvider.java
index c88a0fc5..1e6ea739 100644
--- a/src/main/java/net/fabricmc/loom/providers/MinecraftPatchedProvider.java
+++ b/src/main/java/net/fabricmc/loom/providers/MinecraftPatchedProvider.java
@@ -24,6 +24,7 @@
 
 package net.fabricmc.loom.providers;
 
+import com.google.common.base.Predicates;
 import com.google.common.collect.ImmutableMap;
 import com.google.gson.JsonParser;
 import de.oceanlabs.mcp.mcinjector.adaptors.ParameterAnnotationFixer;
@@ -34,7 +35,9 @@ import net.fabricmc.loom.util.srg.SpecialSourceExecutor;
 import net.fabricmc.mapping.tree.TinyTree;
 import net.fabricmc.tinyremapper.OutputConsumerPath;
 import net.fabricmc.tinyremapper.TinyRemapper;
+import net.minecraftforge.accesstransformer.AccessTransformerEngine;
 import net.minecraftforge.accesstransformer.TransformerProcessor;
+import net.minecraftforge.accesstransformer.parser.AccessTransformerList;
 import net.minecraftforge.binarypatcher.ConsoleTool;
 import org.apache.commons.io.FileUtils;
 import org.apache.commons.io.IOUtils;
@@ -51,6 +54,7 @@ import org.objectweb.asm.tree.ClassNode;
 import org.zeroturnaround.zip.ZipUtil;
 
 import java.io.*;
+import java.lang.reflect.Field;
 import java.net.URI;
 import java.net.URL;
 import java.nio.charset.StandardCharsets;
@@ -59,18 +63,27 @@ import java.nio.file.*;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Locale;
+import java.util.Map;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
+import java.util.stream.Stream;
 
 public class MinecraftPatchedProvider extends DependencyProvider {
 	private final MappingsProvider mappingsProvider;
+	// Step 1: Remap Minecraft to SRG
 	private File minecraftClientSrgJar;
 	private File minecraftServerSrgJar;
+	// Step 2: Binary Patch
 	private File minecraftClientPatchedSrgJar;
 	private File minecraftServerPatchedSrgJar;
+	// Step 3: Access Transform
+	private File minecraftClientPatchedSrgATJar;
+	private File minecraftServerPatchedSrgATJar;
+	// Step 4: Remap Patched AT to Official
 	private File minecraftClientPatchedOfficialJar;
 	private File minecraftServerPatchedOfficialJar;
+	// Step 5: Merge
 	private File minecraftMergedPatchedJar;
 	private File projectAtHash;
 	@Nullable
@@ -108,7 +121,7 @@ public class MinecraftPatchedProvider extends DependencyProvider {
 				writeAtHash();
 			}
 
-			atDirty = mismatched && projectAt != null;
+			atDirty = mismatched;
 		}
 
 		MinecraftProvider minecraftProvider = getExtension().getMinecraftProvider();
@@ -117,44 +130,56 @@ public class MinecraftPatchedProvider extends DependencyProvider {
 		String jarSuffix = "-patched-forge-" + patchProvider.forgeVersion;
 		minecraftProvider.setJarSuffix(jarSuffix);
 
-		File cache = usesProjectCache() ? getExtension().getProjectPersistentCache() : getExtension().getUserCache();
+		File globalCache = getExtension().getUserCache();
+		File cache = usesProjectCache() ? getExtension().getProjectPersistentCache() : globalCache;
 
+		minecraftClientSrgJar = new File(globalCache, "minecraft-" + minecraftVersion + "-client-srg.jar");
+		minecraftServerSrgJar = new File(globalCache, "minecraft-" + minecraftVersion + "-server-srg.jar");
+		minecraftClientPatchedSrgJar = new File(globalCache, "minecraft-" + minecraftVersion + "-client-srg" + jarSuffix + ".jar");
+		minecraftServerPatchedSrgJar = new File(globalCache, "minecraft-" + minecraftVersion + "-server-srg" + jarSuffix + ".jar");
+		minecraftClientPatchedSrgATJar = new File(cache, "minecraft-" + minecraftVersion + "-client-srg-at" + jarSuffix + ".jar");
+		minecraftServerPatchedSrgATJar = new File(cache, "minecraft-" + minecraftVersion + "-server-srg-at" + jarSuffix + ".jar");
 		minecraftClientPatchedOfficialJar = new File(cache, "minecraft-" + minecraftVersion + "-client" + jarSuffix + ".jar");
 		minecraftServerPatchedOfficialJar = new File(cache, "minecraft-" + minecraftVersion + "-server" + jarSuffix + ".jar");
-		minecraftClientSrgJar = new File(cache, "minecraft-" + minecraftVersion + "-client-srg.jar");
-		minecraftServerSrgJar = new File(cache, "minecraft-" + minecraftVersion + "-server-srg.jar");
-		minecraftClientPatchedSrgJar = new File(cache, "minecraft-" + minecraftVersion + "-client-srg" + jarSuffix + ".jar");
-		minecraftServerPatchedSrgJar = new File(cache, "minecraft-" + minecraftVersion + "-server-srg" + jarSuffix + ".jar");
 		minecraftMergedPatchedJar = new File(cache, "minecraft-" + minecraftVersion + "-merged" + jarSuffix + ".jar");
 
-		if (isRefreshDeps()) {
-			cleanCache();
+		if (isRefreshDeps() || Stream.of(getGlobalCaches()).anyMatch(Predicates.not(File::exists))) {
+			cleanAllCache();
+		} else if (atDirty || Stream.of(getProjectCache()).anyMatch(Predicates.not(File::exists))) {
+			cleanProjectCache();
 		}
+	}
 
-		if (!minecraftClientSrgJar.exists() || !minecraftServerSrgJar.exists()
-		    || !minecraftClientPatchedSrgJar.exists() || !minecraftServerPatchedSrgJar.exists()
-		    || !minecraftMergedPatchedJar.exists()) {
-			minecraftClientSrgJar.delete();
-			minecraftServerSrgJar.delete();
-			minecraftClientPatchedSrgJar.delete();
-			minecraftServerPatchedSrgJar.delete();
-			minecraftMergedPatchedJar.delete();
+	public void cleanAllCache() {
+		for (File file : getGlobalCaches()) {
+			file.delete();
 		}
+		cleanProjectCache();
 	}
 
-	public void cleanCache() {
-		for (File file : Arrays.asList(
-				projectAtHash,
+	private File[] getGlobalCaches() {
+		return new File[]{
 				minecraftClientSrgJar,
 				minecraftServerSrgJar,
 				minecraftClientPatchedSrgJar,
-				minecraftServerPatchedSrgJar,
+				minecraftServerPatchedSrgJar
+		};
+	}
+
+	public void cleanProjectCache() {
+		for (File file : getProjectCache()) {
+			file.delete();
+		}
+	}
+
+	private File[] getProjectCache() {
+		return new File[]{
+				minecraftClientPatchedSrgATJar,
+				minecraftServerPatchedSrgATJar,
 				minecraftClientPatchedOfficialJar,
 				minecraftServerPatchedOfficialJar,
 				minecraftMergedPatchedJar
-		)) {
-			file.delete();
-		}
+		};
 	}
 
 	@Override
@@ -165,21 +190,30 @@ public class MinecraftPatchedProvider extends DependencyProvider {
 			getProject().getLogger().lifecycle(":found dirty access transformers");
 		}
 
-		if (atDirty || !minecraftClientPatchedOfficialJar.exists() || !minecraftServerPatchedOfficialJar.exists()) {
-			if (!minecraftClientSrgJar.exists() || !minecraftServerSrgJar.exists()) {
-				// Remap official jars to MCPConfig remapped srg jars
-				createSrgJars(getProject().getLogger());
-			}
+		boolean dirty = false;
 
-			if (atDirty || !minecraftClientPatchedSrgJar.exists() || !minecraftServerPatchedSrgJar.exists()) {
-				patchJars(getProject().getLogger());
-				injectForgeClasses(getProject().getLogger());
-			}
+		if (!minecraftClientSrgJar.exists() || !minecraftServerSrgJar.exists()) {
+			dirty = true;
+			// Remap official jars to MCPConfig remapped srg jars
+			createSrgJars(getProject().getLogger());
+		}
+
+		if (!minecraftClientPatchedSrgJar.exists() || !minecraftServerPatchedSrgJar.exists()) {
+			dirty = true;
+			patchJars(getProject().getLogger());
+			injectForgeClasses(getProject().getLogger());
+		}
+
+		if (atDirty || !minecraftClientPatchedSrgATJar.exists() || !minecraftServerPatchedSrgATJar.exists()) {
+			dirty = true;
+			accessTransformForge(getProject().getLogger());
+		}
 
+		if (dirty) {
 			remapPatchedJars(getProject().getLogger());
 		}
 
-		if (atDirty || !minecraftMergedPatchedJar.exists()) {
+		if (dirty || !minecraftMergedPatchedJar.exists()) {
 			mergeJars(getProject().getLogger());
 		}
 	}
@@ -252,12 +286,9 @@ public class MinecraftPatchedProvider extends DependencyProvider {
 
 	private void injectForgeClasses(Logger logger) throws IOException {
 		logger.lifecycle(":injecting forge classes into minecraft");
-		ThreadingUtils.run(() -> {
-			copyAll(getExtension().getForgeUniversalProvider().getForge(), minecraftClientPatchedSrgJar);
-			copyUserdevFiles(getExtension().getForgeUserdevProvider().getUserdevJar(), minecraftClientPatchedSrgJar);
-		}, () -> {
-			copyAll(getExtension().getForgeUniversalProvider().getForge(), minecraftServerPatchedSrgJar);
-			copyUserdevFiles(getExtension().getForgeUserdevProvider().getUserdevJar(), minecraftServerPatchedSrgJar);
+		ThreadingUtils.run(Arrays.asList(Environment.values()), environment -> {
+			copyAll(getExtension().getForgeUniversalProvider().getForge(), environment.patchedSrgJar.apply(this));
+			copyUserdevFiles(getExtension().getForgeUserdevProvider().getUserdevJar(), environment.patchedSrgJar.apply(this));
 		});
 
 		logger.lifecycle(":injecting loom classes into minecraft");
@@ -271,38 +302,73 @@ public class MinecraftPatchedProvider extends DependencyProvider {
 			String side = environment.side();
 			File target = environment.patchedSrgJar.apply(this);
 			walkFileSystems(injection, target, it -> !it.getFileName().toString().equals("MANIFEST.MF"), this::copyReplacing);
+		}
+	}
 
+	private void accessTransformForge(Logger logger) throws Exception {
+		for (Environment environment : Environment.values()) {
+			String side = environment.side();
 			logger.lifecycle(":access transforming minecraft (" + side + ")");
 
-			File atJar = File.createTempFile("at" + side, ".jar");
+			File input = environment.patchedSrgJar.apply(this);
+			File inputCopied = File.createTempFile("at" + side, ".jar");
+			FileUtils.copyFile(input, inputCopied);
+			File target = environment.patchedSrgATJar.apply(this);
+			target.delete();
 			File at = File.createTempFile("at" + side, ".cfg");
-			FileUtils.copyFile(target, atJar);
-			JarUtil.extractFile(atJar, "META-INF/accesstransformer.cfg", at);
+			JarUtil.extractFile(inputCopied, "META-INF/accesstransformer.cfg", at);
 			String[] args = new String[]{
-					"--inJar", atJar.getAbsolutePath(),
+					"--inJar", inputCopied.getAbsolutePath(),
 					"--outJar", target.getAbsolutePath(),
 					"--atFile", at.getAbsolutePath()
 			};
 
-			if (projectAt != null) {
+			if (usesProjectCache()) {
 				args = Arrays.copyOf(args, args.length + 2);
 				args[args.length - 2] = "--atFile";
 				args[args.length - 1] = projectAt.getAbsolutePath();
 			}
 
+			resetAccessTransformerEngine();
 			TransformerProcessor.main(args);
-		};
+			inputCopied.delete();
+		}
 	}
 
-	private enum Environment {
-		CLIENT(provider -> provider.minecraftClientPatchedSrgJar, provider -> provider.minecraftClientPatchedOfficialJar),
-		SERVER(provider -> provider.minecraftServerPatchedSrgJar, provider -> provider.minecraftServerPatchedOfficialJar);
-
-		Function<MinecraftPatchedProvider, File> patchedSrgJar;
-		Function<MinecraftPatchedProvider, File> patchedOfficialJar;
+	private void resetAccessTransformerEngine() throws Exception {
+		// Thank you Forge, I love you
+		Field field = AccessTransformerEngine.class.getDeclaredField("masterList");
+		field.setAccessible(true);
+		AccessTransformerList list = (AccessTransformerList) field.get(AccessTransformerEngine.INSTANCE);
+		field = AccessTransformerList.class.getDeclaredField("accessTransformers");
+		field.setAccessible(true);
+		((Map<?, ?>) field.get(list)).clear();
+	}
 
-		Environment(Function<MinecraftPatchedProvider, File> patchedSrgJar, Function<MinecraftPatchedProvider, File> patchedOfficialJar) {
+	private enum Environment {
+		CLIENT(provider -> provider.minecraftClientSrgJar,
+				provider -> provider.minecraftClientPatchedSrgJar,
+				provider -> provider.minecraftClientPatchedSrgATJar,
+				provider -> provider.minecraftClientPatchedOfficialJar
+		),
+		SERVER(provider -> provider.minecraftServerSrgJar,
+				provider -> provider.minecraftServerPatchedSrgJar,
+				provider -> provider.minecraftServerPatchedSrgATJar,
+				provider -> provider.minecraftServerPatchedOfficialJar
+		);
+
+		final Function<MinecraftPatchedProvider, File> srgJar;
+		final Function<MinecraftPatchedProvider, File> patchedSrgJar;
+		final Function<MinecraftPatchedProvider, File> patchedSrgATJar;
+		final Function<MinecraftPatchedProvider, File> patchedOfficialJar;
+
+		Environment(Function<MinecraftPatchedProvider, File> srgJar,
+				Function<MinecraftPatchedProvider, File> patchedSrgJar,
+				Function<MinecraftPatchedProvider, File> patchedSrgATJar,
+				Function<MinecraftPatchedProvider, File> patchedOfficialJar) {
+			this.srgJar = srgJar;
 			this.patchedSrgJar = patchedSrgJar;
+			this.patchedSrgATJar = patchedSrgATJar;
 			this.patchedOfficialJar = patchedOfficialJar;
 		}
 
@@ -324,7 +390,7 @@ public class MinecraftPatchedProvider extends DependencyProvider {
 			logger.lifecycle(":remapping minecraft (TinyRemapper, " + environment.side() + ", srg -> official)");
 			TinyTree mappingsWithSrg = getExtension().getMappingsProvider().getMappingsWithSrg();
 
-			Path input = environment.patchedSrgJar.apply(this).toPath();
+			Path input = environment.patchedSrgATJar.apply(this).toPath();
 			Path output = environment.patchedOfficialJar.apply(this).toPath();
 
 			Files.deleteIfExists(output);
@@ -356,12 +422,9 @@ public class MinecraftPatchedProvider extends DependencyProvider {
 		PatchProvider patchProvider = getExtension().getPatchProvider();
 		patchJars(minecraftClientSrgJar, minecraftClientPatchedSrgJar, patchProvider.clientPatches);
 		patchJars(minecraftServerSrgJar, minecraftServerPatchedSrgJar, patchProvider.serverPatches);
-		ThreadingUtils.run(() -> {
-			copyMissingClasses(minecraftClientSrgJar, minecraftClientPatchedSrgJar);
-			fixParameterAnnotation(minecraftClientPatchedSrgJar);
-		}, () -> {
-			copyMissingClasses(minecraftServerSrgJar, minecraftServerPatchedSrgJar);
-			fixParameterAnnotation(minecraftServerPatchedSrgJar);
+		ThreadingUtils.run(Arrays.asList(Environment.values()), environment -> {
+			copyMissingClasses(environment.srgJar.apply(this), environment.patchedSrgJar.apply(this));
+			fixParameterAnnotation(environment.patchedSrgJar.apply(this));
 		});
 	}
 
diff --git a/src/main/java/net/fabricmc/loom/task/CleanLoomBinaries.java b/src/main/java/net/fabricmc/loom/task/CleanLoomBinaries.java
index 2b845afb..2981a6ec 100644
--- a/src/main/java/net/fabricmc/loom/task/CleanLoomBinaries.java
+++ b/src/main/java/net/fabricmc/loom/task/CleanLoomBinaries.java
@@ -45,7 +45,7 @@ public class CleanLoomBinaries extends AbstractLoomTask {
 			try {
 				MinecraftPatchedProvider provider = new MinecraftPatchedProvider(extension.getMappingsProvider(), getProject());
 				provider.initFiles();
-				provider.cleanCache();
+				provider.cleanAllCache();
 			} catch (IOException e) {
 				e.printStackTrace();
 			}