Use TinyRemapper directly

5 files changed, 39 insertions, 214 deletions

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 3be5e9f8..90af6c80 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
@@ -29,17 +29,15 @@ import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
-import java.io.UncheckedIOException;
 import java.net.URI;
 import java.nio.file.FileSystem;
 import java.nio.file.FileSystems;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.nio.file.StandardCopyOption;
-import java.nio.file.StandardOpenOption;
+import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Consumer;
@@ -52,15 +50,10 @@ import dev.architectury.tinyremapper.IMappingProvider;
 import dev.architectury.tinyremapper.TinyRemapper;
 import org.gradle.api.Project;
 import org.jetbrains.annotations.Nullable;
-import org.objectweb.asm.ClassReader;
-import org.objectweb.asm.ClassWriter;
-import org.objectweb.asm.commons.ClassRemapper;
 
 import net.fabricmc.loom.configuration.DependencyProvider;
 import net.fabricmc.loom.configuration.providers.MinecraftProvider;
 import net.fabricmc.loom.configuration.providers.mappings.MappingsProvider;
-import net.fabricmc.loom.configuration.providers.minecraft.tr.CompiledMappedClassRemapper;
-import net.fabricmc.loom.configuration.providers.minecraft.tr.MappingsCompiled;
 import net.fabricmc.loom.configuration.providers.minecraft.tr.OutputRemappingHandler;
 import net.fabricmc.loom.util.Constants;
 import net.fabricmc.loom.util.DownloadUtil;
@@ -155,27 +148,38 @@ public class MinecraftMappedProvider extends DependencyProvider {
 		TinyRemapper remapper = getTinyRemapper();
 		remapper.readClassPath(libraries);
 		remapper.prepareClasses();
-		remapper.readInputs(input);
 
 		Path tmpAssets = Files.createTempFile("tmpAssets", null);
 		Files.deleteIfExists(tmpAssets);
 		tmpAssets.toFile().deleteOnExit();
 
+		List<byte[]> inputByteList = new ArrayList<>();
+
 		try (FileSystemUtil.FileSystemDelegate inputFs = FileSystemUtil.getJarFileSystem(input, false)) {
 			ThreadingUtils.TaskCompleter taskCompleter = ThreadingUtils.taskCompleter();
 
 			try (FileSystemUtil.FileSystemDelegate assetsFs = FileSystemUtil.getJarFileSystem(tmpAssets, true)) {
 				for (Path path : (Iterable<? extends Path>) Files.walk(inputFs.get().getPath("/"))::iterator) {
-					if (Files.isRegularFile(path) && !path.getFileName().toString().endsWith(".class")) {
-						Path p = assetsFs.get().getPath(path.toString());
-
-						if (p.getParent() != null) {
-							Files.createDirectories(p.getParent());
+					if (Files.isRegularFile(path)) {
+						if (path.getFileName().toString().endsWith(".class")) {
+							taskCompleter.add(() -> {
+								byte[] bytes = Files.readAllBytes(path);
+
+								synchronized (inputByteList) {
+									inputByteList.add(bytes);
+								}
+							});
+						} else {
+							Path p = assetsFs.get().getPath(path.toString());
+
+							if (p.getParent() != null) {
+								Files.createDirectories(p.getParent());
+							}
+
+							taskCompleter.add(() -> {
+								Files.copy(path, p);
+							});
 						}
-
-						taskCompleter.add(() -> {
-							Files.copy(path, p);
-						});
 					}
 				}
 
@@ -183,40 +187,19 @@ public class MinecraftMappedProvider extends DependencyProvider {
 			}
 		}
 
-		Files.copy(tmpAssets, outputMapped, StandardCopyOption.REPLACE_EXISTING);
-		FileSystemUtil.FileSystemDelegate systemMapped = FileSystemUtil.getJarFileSystem(outputMapped, true);
-		ThreadingUtils.TaskCompleter mappedRemapper = ThreadingUtils.taskCompleter().onComplete(stopwatch -> {
-			getProject().getLogger().lifecycle(":remapped minecraft (AsmRemapper, intermediary -> named) in " + stopwatch);
-		});
-		MappingsCompiled compiledMapped = new MappingsCompiled(getMappings(input, "intermediary", "named"));
+		byte[][] inputBytes = inputByteList.toArray(new byte[0][0]);
 
-		for (String toM : getExtension().isForge() ? Arrays.asList("intermediary", "srg") : Collections.singletonList("intermediary")) {
-			Path output = "srg".equals(toM) ? outputSrg : outputIntermediary;
+		for (String toM : getExtension().isForge() ? Arrays.asList("intermediary", "srg", "named") : Arrays.asList("intermediary", "named")) {
+			Path output = "named".equals(toM) ? outputMapped : "srg".equals(toM) ? outputSrg : outputIntermediary;
 			Stopwatch stopwatch = Stopwatch.createStarted();
 			getProject().getLogger().lifecycle(":remapping minecraft (TinyRemapper, " + fromM + " -> " + toM + ")");
 
+			remapper.readInputs(inputBytes);
 			remapper.replaceMappings(getMappings(input, fromM, toM));
-			OutputRemappingHandler.remap(remapper, tmpAssets, output, toM.equals("intermediary") ? (path, bytes) -> {
-				try {
-					Path fsPath = systemMapped.get().getPath(compiledMapped.map(path) + ".class");
-
-					if (fsPath.getParent() != null) {
-						Files.createDirectories(fsPath.getParent());
-					}
-
-					mappedRemapper.add(() -> {
-						ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_MAXS);
-						ClassRemapper classRemapper = new CompiledMappedClassRemapper(writer, compiledMapped);
-						new ClassReader(bytes).accept(classRemapper, ClassReader.EXPAND_FRAMES);
-
-						Files.write(fsPath, writer.toByteArray(), StandardOpenOption.CREATE);
-					});
-				} catch (IOException e) {
-					throw new UncheckedIOException(e);
-				}
-			} : null);
+			OutputRemappingHandler.remap(remapper, tmpAssets, output);
 
 			getProject().getLogger().lifecycle(":remapped minecraft (TinyRemapper, " + fromM + " -> " + toM + ") in " + stopwatch);
+			remapper.removeInput();
 
 			if (getExtension().isForge() && !"srg".equals(toM)) {
 				getProject().getLogger().info(":running forge finalising tasks");
@@ -255,11 +238,6 @@ public class MinecraftMappedProvider extends DependencyProvider {
 		}
 
 		remapper.finish();
-
-		getProject().getLogger().lifecycle(":remapping minecraft (AsmRemapper, intermediary -> named)");
-		mappedRemapper.complete();
-
-		systemMapped.close();
 	}
 
 	public TinyRemapper getTinyRemapper() throws IOException {
diff --git a/src/main/java/net/fabricmc/loom/configuration/providers/minecraft/tr/CompiledMappedClassRemapper.java b/src/main/java/net/fabricmc/loom/configuration/providers/minecraft/tr/CompiledMappedClassRemapper.java
deleted file mode 100644
index e9b484f4..00000000
--- a/src/main/java/net/fabricmc/loom/configuration/providers/minecraft/tr/CompiledMappedClassRemapper.java
+++ /dev/null
@@ -1,51 +0,0 @@
-package net.fabricmc.loom.configuration.providers.minecraft.tr;
-
-import java.util.HashMap;
-
-import dev.architectury.tinyremapper.AsmClassRemapper;
-import org.objectweb.asm.ClassVisitor;
-import org.objectweb.asm.MethodVisitor;
-import org.objectweb.asm.Opcodes;
-import org.objectweb.asm.commons.ClassRemapper;
-import org.objectweb.asm.commons.MethodRemapper;
-import org.objectweb.asm.tree.MethodNode;
-
-public class CompiledMappedClassRemapper extends ClassRemapper {
-	private final MappingsCompiled compiled;
-	private String lastName;
-	private MethodNode lastMethod;
-
-	public CompiledMappedClassRemapper(ClassVisitor classVisitor, MappingsCompiled compiled) {
-		super(Opcodes.ASM9, classVisitor, compiled);
-
-		this.compiled = compiled;
-	}
-
-	@Override
-	public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
-		lastName = name;
-		this.compiled.lastSuperClass = superName;
-		this.compiled.lastInterfaces = interfaces;
-		super.visit(version, access, name, signature, superName, interfaces);
-	}
-
-	@Override
-	public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
-		lastMethod = new MethodNode(api, access, name, descriptor, signature, exceptions);
-		return super.visitMethod(access, name, descriptor, signature, exceptions);
-	}
-
-	@Override
-	protected MethodVisitor createMethodRemapper(MethodVisitor methodVisitor) {
-		return new MethodRemapper(api, lastMethod, remapper) {
-			@Override
-			public void visitEnd() {
-				lastMethod.localVariables = null;
-				lastMethod.parameters = null;
-				AsmClassRemapper.AsmMethodRemapper.processLocals(compiled, lastName, lastMethod, false, true, new HashMap<>());
-				lastMethod.visitEnd();
-				lastMethod.accept(methodVisitor);
-			}
-		};
-	}
-}
diff --git a/src/main/java/net/fabricmc/loom/configuration/providers/minecraft/tr/MappingsCompiled.java b/src/main/java/net/fabricmc/loom/configuration/providers/minecraft/tr/MappingsCompiled.java
deleted file mode 100644
index c4ebb67e..00000000
--- a/src/main/java/net/fabricmc/loom/configuration/providers/minecraft/tr/MappingsCompiled.java
+++ /dev/null
@@ -1,107 +0,0 @@
-package net.fabricmc.loom.configuration.providers.minecraft.tr;
-
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Set;
-
-import com.google.common.collect.HashBasedTable;
-import com.google.common.collect.Table;
-import dev.architectury.tinyremapper.IMappingProvider;
-import dev.architectury.tinyremapper.MethodRemapperProvider;
-import org.objectweb.asm.commons.Remapper;
-
-public class MappingsCompiled extends Remapper implements MethodRemapperProvider {
-	private final Map<String, String> classes;
-	private final Map<String, String> fields;
-	private final Map<String, String> methods;
-	private final Table<String, Integer, String> methodArgs;
-	String lastSuperClass;
-	String[] lastInterfaces;
-
-	public MappingsCompiled(Set<IMappingProvider> mappings) {
-		this.classes = new HashMap<>();
-		this.fields = new HashMap<>();
-		this.methods = new HashMap<>();
-		this.methodArgs = HashBasedTable.create();
-
-		for (IMappingProvider mapping : mappings) {
-			mapping.load(new IMappingProvider.MappingAcceptor() {
-				@Override
-				public void acceptClass(String srcName, String dstName) {
-					classes.put(srcName, dstName);
-				}
-
-				@Override
-				public void acceptMethod(IMappingProvider.Member method, String dstName) {
-					methods.put(method.name, dstName);
-				}
-
-				@Override
-				public void acceptMethodArg(IMappingProvider.Member method, int lvIndex, String dstName) {
-					methodArgs.put(method.owner + "|" + method.name, lvIndex, dstName);
-				}
-
-				@Override
-				public void acceptMethodVar(IMappingProvider.Member method, int lvIndex, int startOpIdx, int asmIndex, String dstName) {
-				}
-
-				@Override
-				public void acceptField(IMappingProvider.Member field, String dstName) {
-					fields.put(field.name, dstName);
-				}
-			});
-		}
-	}
-
-	@Override
-	public String map(String name) {
-		return classes.getOrDefault(name, name);
-	}
-
-	@Override
-	public String mapFieldName(String owner, String name, String descriptor) {
-		return mapField(name);
-	}
-
-	public String mapField(String name) {
-		return fields.getOrDefault(name, name);
-	}
-
-	@Override
-	public String mapMethodName(String owner, String name, String descriptor) {
-		return mapMethod(name);
-	}
-
-	public String mapMethod(String name) {
-		return methods.getOrDefault(name, name);
-	}
-
-	public String mapMethodArg(String methodOwner, String methodName, int lvIndex, String def) {
-		String arg = methodArgs.get(methodOwner + "|" + methodName, lvIndex);
-		if (arg != null) return arg;
-
-		if (lastSuperClass != null) {
-			arg = methodArgs.get(lastSuperClass + "|" + methodName, lvIndex);
-			if (arg != null) return arg;
-		}
-
-		if (lastInterfaces != null) {
-			for (String lastInterface : lastInterfaces) {
-				arg = methodArgs.get(lastInterface + "|" + methodName, lvIndex);
-				if (arg != null) return arg;
-			}
-		}
-
-		return def;
-	}
-
-	@Override
-	public String mapMethodVar(String methodOwner, String methodName, String methodDesc, int lvIndex, int startOpIdx, int asmIndex, String name) {
-		return name;
-	}
-
-	@Override
-	public String mapMethodArg(String methodOwner, String methodName, String methodDesc, int lvIndex, String name) {
-		return mapMethodArg(methodOwner, methodName, lvIndex, name);
-	}
-}
diff --git a/src/main/java/net/fabricmc/loom/configuration/providers/minecraft/tr/OutputRemappingHandler.java b/src/main/java/net/fabricmc/loom/configuration/providers/minecraft/tr/OutputRemappingHandler.java
index d3a6c69d..f69beb0c 100644
--- a/src/main/java/net/fabricmc/loom/configuration/providers/minecraft/tr/OutputRemappingHandler.java
+++ b/src/main/java/net/fabricmc/loom/configuration/providers/minecraft/tr/OutputRemappingHandler.java
@@ -17,8 +17,7 @@ import net.fabricmc.loom.util.ThreadingUtils;
 
 public class OutputRemappingHandler {
 	public static void remap(TinyRemapper remapper, Path assets, Path output) throws IOException {
-		remap(remapper, assets, output, (path, bytes) -> {
-		});
+		remap(remapper, assets, output, null);
 	}
 
 	public static void remap(TinyRemapper remapper, Path assets, Path output, BiConsumer<String, byte[]> then) throws IOException {
diff --git a/src/main/java/net/fabricmc/loom/util/ThreadingUtils.java b/src/main/java/net/fabricmc/loom/util/ThreadingUtils.java
index f0c9ff9c..53c2d221 100644
--- a/src/main/java/net/fabricmc/loom/util/ThreadingUtils.java
+++ b/src/main/java/net/fabricmc/loom/util/ThreadingUtils.java
@@ -130,7 +130,7 @@ public class ThreadingUtils {
 		return new TaskCompleter();
 	}
 
-	public static class TaskCompleter {
+	public static class TaskCompleter implements Function<Throwable, Void> {
 		Stopwatch stopwatch = Stopwatch.createUnstarted();
 		List<CompletableFuture<?>> tasks = new ArrayList<>();
 		ExecutorService service = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());
@@ -147,7 +147,7 @@ public class ThreadingUtils {
 				} catch (Throwable throwable) {
 					throw new RuntimeException(throwable);
 				}
-			}, service));
+			}, service).exceptionally(this));
 
 			return this;
 		}
@@ -159,7 +159,7 @@ public class ThreadingUtils {
 
 		public void complete() {
 			try {
-				CompletableFuture.allOf(tasks.toArray(new CompletableFuture[0])).get();
+				CompletableFuture.allOf(tasks.toArray(new CompletableFuture[0])).exceptionally(this).get();
 				service.shutdownNow();
 				stopwatch.stop();
 
@@ -170,5 +170,11 @@ public class ThreadingUtils {
 				throw new RuntimeException(e);
 			}
 		}
+
+		@Override
+		public Void apply(Throwable throwable) {
+			throwable.printStackTrace();
+			return null;
+		}
 	}
 }