Re-use forge accesstransformer tool for legacy forge

3 files changed, 123 insertions, 100 deletions

diff --git a/src/main/java/net/fabricmc/loom/configuration/providers/forge/MinecraftPatchedProvider.java b/src/main/java/net/fabricmc/loom/configuration/providers/forge/MinecraftPatchedProvider.java
index 05a216bd..f11907dc 100644
--- a/src/main/java/net/fabricmc/loom/configuration/providers/forge/MinecraftPatchedProvider.java
+++ b/src/main/java/net/fabricmc/loom/configuration/providers/forge/MinecraftPatchedProvider.java
@@ -477,6 +477,26 @@ public class MinecraftPatchedProvider extends DependencyProvider {
 	}
 
 	private void accessTransformForge(Logger logger) throws Exception {
+		List<byte[]> ats = new ArrayList<>();
+
+		for (File jar : ImmutableList.of(getForgeJar(), getForgeUserdevJar(), minecraftMergedPatchedSrgJar)) {
+			byte[] atBytes = ZipUtils.unpackNullable(jar.toPath(), Constants.Forge.ACCESS_TRANSFORMER_PATH);
+
+			if (atBytes != null) {
+				ats.add(atBytes);
+			}
+		}
+
+		if (usesProjectCache()) {
+			for (File projectAt : projectAts) {
+				ats.add(Files.readAllBytes(projectAt.toPath()));
+			}
+		}
+
+		accessTransformForge(logger, minecraftMergedPatchedSrgJar, minecraftMergedPatchedSrgAtJar, ats);
+	}
+
+	protected void accessTransformForge(Logger logger, File input, File target, List<byte[]> ats) throws Exception {
 		MinecraftProviderImpl minecraftProvider = getExtension().getMinecraftProvider();
 		List<File> toDelete = new ArrayList<>();
 		String atDependency = Constants.Dependencies.ACCESS_TRANSFORMERS + (minecraftProvider.isNewerThan21w39a() ? Constants.Dependencies.Versions.ACCESS_TRANSFORMERS_NEW : Constants.Dependencies.Versions.ACCESS_TRANSFORMERS);
@@ -485,8 +505,6 @@ public class MinecraftPatchedProvider extends DependencyProvider {
 
 		logger.lifecycle(":access transforming minecraft");
 
-		File input = minecraftMergedPatchedSrgJar;
-		File target = minecraftMergedPatchedSrgAtJar;
 		Files.deleteIfExists(target.toPath());
 
 		List<String> args = new ArrayList<>();
@@ -495,23 +513,12 @@ public class MinecraftPatchedProvider extends DependencyProvider {
 		args.add("--outJar");
 		args.add(target.getAbsolutePath());
 
-		for (File jar : ImmutableList.of(getForgeJar(), getForgeUserdevJar(), minecraftMergedPatchedSrgJar)) {
-			byte[] atBytes = ZipUtils.unpackNullable(jar.toPath(), Constants.Forge.ACCESS_TRANSFORMER_PATH);
-
-			if (atBytes != null) {
-				File tmpFile = File.createTempFile("at-conf", ".cfg");
-				toDelete.add(tmpFile);
-				Files.write(tmpFile.toPath(), atBytes, StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING);
-				args.add("--atFile");
-				args.add(tmpFile.getAbsolutePath());
-			}
-		}
-
-		if (usesProjectCache()) {
-			for (File projectAt : projectAts) {
-				args.add("--atFile");
-				args.add(projectAt.getAbsolutePath());
-			}
+		for (byte[] atBytes : ats) {
+			File tmpFile = File.createTempFile("at-conf", ".cfg");
+			toDelete.add(tmpFile);
+			Files.write(tmpFile.toPath(), atBytes, StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING);
+			args.add("--atFile");
+			args.add(tmpFile.getAbsolutePath());
 		}
 
 		getProject().javaexec(spec -> {
diff --git a/src/main/java/net/fabricmc/loom/configuration/providers/forge/fg2/MinecraftLegacyPatchedProvider.java b/src/main/java/net/fabricmc/loom/configuration/providers/forge/fg2/MinecraftLegacyPatchedProvider.java
index e87004c2..e74e463e 100644
--- a/src/main/java/net/fabricmc/loom/configuration/providers/forge/fg2/MinecraftLegacyPatchedProvider.java
+++ b/src/main/java/net/fabricmc/loom/configuration/providers/forge/fg2/MinecraftLegacyPatchedProvider.java
@@ -24,11 +24,13 @@
 
 package net.fabricmc.loom.configuration.providers.forge.fg2;
 
-import java.io.BufferedReader;
+import java.io.BufferedWriter;
 import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStreamReader;
+import java.io.OutputStreamWriter;
 import java.net.URI;
 import java.nio.file.FileSystem;
 import java.nio.file.FileSystems;
@@ -37,6 +39,7 @@ import java.nio.file.Path;
 import java.nio.file.StandardCopyOption;
 import java.nio.file.StandardOpenOption;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
@@ -44,8 +47,9 @@ import java.util.stream.Stream;
 
 import com.google.common.base.Stopwatch;
 import com.google.common.collect.ImmutableMap;
-import net.md_5.specialsource.AccessChange;
-import net.md_5.specialsource.AccessMap;
+import org.cadixdev.at.AccessTransformSet;
+import org.cadixdev.at.io.AccessTransformFormats;
+import org.cadixdev.lorenz.MappingSet;
 import org.gradle.api.Project;
 import org.gradle.api.logging.Logger;
 import org.objectweb.asm.AnnotationVisitor;
@@ -63,6 +67,8 @@ import net.fabricmc.loom.util.FileSystemUtil;
 import net.fabricmc.loom.util.ThreadingUtils;
 import net.fabricmc.loom.util.TinyRemapperHelper;
 import net.fabricmc.loom.util.ZipUtils;
+import net.fabricmc.loom.util.srg.AccessTransformSetMapper;
+import net.fabricmc.lorenztiny.TinyMappingsReader;
 import net.fabricmc.mappingio.tree.MappingTree;
 import net.fabricmc.stitch.merge.JarMerger;
 
@@ -224,26 +230,30 @@ public class MinecraftLegacyPatchedProvider extends MinecraftPatchedProvider {
 	}
 
 	private void accessTransformForge(Logger logger) throws Exception {
-		Stopwatch stopwatch = Stopwatch.createStarted();
-
-		logger.lifecycle(":access transforming minecraft");
-
-		MappingTree mappingTree = getExtension().getMappingsProvider().getMappingsWithSrg();
-
-		AccessMap accessMap = new LegacyAccessMap();
+		// Load all applicable access transformers
+		AccessTransformSet accessTransformSet = AccessTransformSet.create();
 
 		byte[] forgeAt = ZipUtils.unpack(forgeJar.toPath(), "forge_at.cfg");
-		accessMap.loadAccessTransformer(new BufferedReader(new InputStreamReader(new ByteArrayInputStream(forgeAt))));
+		AccessTransformFormats.FML.read(new InputStreamReader(new ByteArrayInputStream(forgeAt)), accessTransformSet);
 
 		for (File projectAt : projectAts) {
-			accessMap.loadAccessTransformer(projectAt);
+			AccessTransformFormats.FML.read(projectAt.toPath(), accessTransformSet);
 		}
 
-		Files.copy(minecraftMergedPatchedJar.toPath(), minecraftMergedPatchedAtJar.toPath(), StandardCopyOption.REPLACE_EXISTING);
-
-		modifyClasses(minecraftMergedPatchedAtJar, writer -> new AccessTransformingVisitor(accessMap, mappingTree, writer));
-
-		logger.lifecycle(":access transformed minecraft in " + stopwatch.stop());
+		// Remap them from srg to official mappings
+		MappingTree mappingTree = getExtension().getMappingsProvider().getMappingsWithSrg();
+		MappingSet mappingSet = new TinyMappingsReader(mappingTree, "srg", "official").read();
+		accessTransformSet = AccessTransformSetMapper.remap(accessTransformSet, mappingSet);
+
+		ByteArrayOutputStream remappedOut = new ByteArrayOutputStream();
+		// TODO the extra BufferedWriter wrapper and closing can be removed once https://github.com/CadixDev/at/issues/6 is fixed
+		BufferedWriter writer = new BufferedWriter(new OutputStreamWriter(remappedOut));
+		AccessTransformFormats.FML.write(writer, accessTransformSet);
+		writer.close();
+		byte[] remappedAt = remappedOut.toByteArray();
+
+		// And finally, apply them to the merged+patched jar
+		accessTransformForge(logger, minecraftMergedPatchedJar, minecraftMergedPatchedAtJar, Collections.singletonList(remappedAt));
 	}
 
 	private void modifyClasses(File jarFile, Function<ClassVisitor, ClassVisitor> func) throws Exception {
@@ -380,68 +390,4 @@ public class MinecraftLegacyPatchedProvider extends MinecraftPatchedProvider {
 			}
 		}
 	}
-
-	private static class LegacyAccessMap extends AccessMap {
-		@Override
-		public void addAccessChange(String key, AccessChange accessChange) {
-			// Forge's AT separates fields/methods from their owner by a space but we require a slash
-			int spaceIdx = key.indexOf(' ');
-
-			if (spaceIdx != -1 && key.charAt(spaceIdx + 1) != '(') {
-				key = key.replaceFirst(" ", "/");
-			}
-
-			super.addAccessChange(key, accessChange);
-		}
-	}
-
-	private static class AccessTransformingVisitor extends ClassVisitor {
-		private final AccessMap accessMap;
-		private final MappingTree mappingTree;
-		private final int src;
-		private final int dst;
-
-		private MappingTree.ClassMapping classMapping;
-		private String mappedClassName;
-
-		private AccessTransformingVisitor(AccessMap accessMap, MappingTree mappingTree, ClassVisitor classVisitor) {
-			super(Opcodes.ASM9, classVisitor);
-			this.accessMap = accessMap;
-			this.mappingTree = mappingTree;
-			this.src = mappingTree.getNamespaceId("official");
-			this.dst = mappingTree.getNamespaceId("srg");
-		}
-
-		@Override
-		public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
-			classMapping = mappingTree.getClass(name, src);
-			mappedClassName = classMapping != null ? classMapping.getName(dst) : name;
-			access = accessMap.applyClassAccess(mappedClassName, access);
-			super.visit(version, access, name, signature, superName, interfaces);
-		}
-
-		@Override
-		public FieldVisitor visitField(int access, String name, String descriptor, String signature, Object value) {
-			MappingTree.FieldMapping field = classMapping != null ? classMapping.getField(name, descriptor, src) : null;
-			String mappedName = field != null ? field.getName(dst) : name;
-			access = accessMap.applyFieldAccess(mappedClassName, mappedName, access);
-			return super.visitField(access, name, descriptor, signature, value);
-		}
-
-		@Override
-		public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
-			MappingTree.MethodMapping method = classMapping != null ? classMapping.getMethod(name, desc, src) : null;
-			String mappedName = method != null ? method.getName(dst) : name;
-			String mappedDesc = method != null ? method.getDesc(dst) : mappingTree.mapDesc(desc, src, dst);
-			access = accessMap.applyMethodAccess(mappedClassName, mappedName, mappedDesc, access);
-			return super.visitMethod(access, name, desc, signature, exceptions);
-		}
-
-		@Override
-		public void visitInnerClass(String name, String outerName, String innerName, int access) {
-			MappingTree.ClassMapping classMapping = mappingTree.getClass(name, src);
-			access = accessMap.applyClassAccess(classMapping != null ? classMapping.getName(dst) : name, access);
-			super.visitInnerClass(name, outerName, innerName, access);
-		}
-	}
 }
diff --git a/src/main/java/net/fabricmc/loom/util/srg/AccessTransformSetMapper.java b/src/main/java/net/fabricmc/loom/util/srg/AccessTransformSetMapper.java
new file mode 100644
index 00000000..d17eabfb
--- /dev/null
+++ b/src/main/java/net/fabricmc/loom/util/srg/AccessTransformSetMapper.java
@@ -0,0 +1,70 @@
+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2018 Minecrell (https://github.com/Minecrell)
+ *
+ * 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.util.Objects;
+import java.util.Optional;
+
+import org.cadixdev.at.AccessTransformSet;
+import org.cadixdev.bombe.type.signature.MethodSignature;
+import org.cadixdev.lorenz.MappingSet;
+import org.cadixdev.lorenz.model.ClassMapping;
+import org.cadixdev.lorenz.model.FieldMapping;
+import org.cadixdev.lorenz.model.Mapping;
+import org.cadixdev.lorenz.model.MethodMapping;
+
+// TODO from https://github.com/CadixDev/at/blob/c2b92fc26cf26e64d3ca3b35abf3364d4b95e6c3/src/main/java/org/cadixdev/at/impl/AccessTransformSetMapper.java
+//      remove once https://github.com/CadixDev/at/issues/7 is fixed
+public final class AccessTransformSetMapper {
+	private AccessTransformSetMapper() {
+	}
+
+	public static AccessTransformSet remap(AccessTransformSet set, MappingSet mappings) {
+		Objects.requireNonNull(set, "set");
+		Objects.requireNonNull(mappings, "mappings");
+
+		AccessTransformSet remapped = AccessTransformSet.create();
+		set.getClasses().forEach((className, classSet) -> {
+			Optional<? extends ClassMapping<?, ?>> mapping = mappings.getClassMapping(className);
+			remap(mappings, mapping, classSet, remapped.getOrCreateClass(mapping.map(Mapping::getFullDeobfuscatedName).orElse(className)));
+		});
+		return remapped;
+	}
+
+	private static void remap(MappingSet mappings, Optional<? extends ClassMapping<?, ?>> mapping, AccessTransformSet.Class set, AccessTransformSet.Class remapped) {
+		remapped.merge(set.get());
+		remapped.mergeAllFields(set.allFields());
+		remapped.mergeAllMethods(set.allMethods());
+
+		set.getFields().forEach((name, transform) ->
+				remapped.mergeField(mapping.flatMap(m -> m.getFieldMapping(name))
+						.map(FieldMapping::getDeobfuscatedName).orElse(name), transform));
+
+		set.getMethods().forEach((signature, transform) ->
+				remapped.mergeMethod(mapping.flatMap(m -> m.getMethodMapping(signature))
+						.map(MethodMapping::getDeobfuscatedSignature)
+						.orElseGet(() -> new MethodSignature(signature.getName(), mappings.deobfuscate(signature.getDescriptor()))), transform));
+	}
+}