Async line number remapping (#1074)

src/main/java/net/fabricmc/loom/decompilers/LineNumberRemapper.java

@@ -26,12 +26,11 @@ package net.fabricmc.loom.decompilers;
 
 import java.io.IOException;
 import java.io.InputStream;
-import java.nio.file.FileVisitResult;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.nio.file.SimpleFileVisitor;
 import java.nio.file.StandardCopyOption;
-import java.nio.file.attribute.BasicFileAttributes;
+import java.util.HashSet;
+import java.util.Set;
 
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.ClassVisitor;
@@ -41,31 +40,30 @@ import org.objectweb.asm.MethodVisitor;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import net.fabricmc.loom.util.AsyncZipProcessor;
 import net.fabricmc.loom.util.Constants;
 
 public record LineNumberRemapper(ClassLineNumbers lineNumbers) {
 	private static final Logger LOGGER = LoggerFactory.getLogger(LineNumberRemapper.class);
 
 	public void process(Path input, Path output) throws IOException {
-		Files.walkFileTree(input, new SimpleFileVisitor<>() {
+		AsyncZipProcessor.processEntries(input, output, new AsyncZipProcessor() {
+			private final Set<Path> createdParents = new HashSet<>();
+
 			@Override
-			public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
-				String rel = input.relativize(file).toString();
-				Path dst = output.resolve(rel);
+			public void processEntryAsync(Path file, Path dst) throws IOException {
 				Path parent = dst.getParent();
 
-				if (parent != null) {
-					Files.createDirectories(parent);
+				synchronized (createdParents) {
+					if (parent != null && createdParents.add(parent)) {
+						Files.createDirectories(parent);
+					}
 				}
 
-				String fName = file.getFileName().toString();
+				String fileName = file.getFileName().toString();
 
-				if (fName.endsWith(".class")) {
-					if (Files.exists(dst)) {
-						Files.delete(dst);
-					}
-
-					String idx = rel.substring(0, rel.length() - 6);
+				if (fileName.endsWith(".class")) {
+					String idx = fileName.substring(0, fileName.length() - 6);
 
 					LOGGER.debug("Remapping line numbers for class: " + idx);
 
@@ -82,13 +80,12 @@ public record LineNumberRemapper(ClassLineNumbers lineNumbers) {
 
 							reader.accept(new LineNumberVisitor(Constants.ASM_VERSION, writer, lineNumbers.lineMap().get(idx)), 0);
 							Files.write(dst, writer.toByteArray());
-							return FileVisitResult.CONTINUE;
+							return;
 						}
 					}
 				}
 
 				Files.copy(file, dst, StandardCopyOption.REPLACE_EXISTING);
-				return FileVisitResult.CONTINUE;
 			}
 		});
 	}

src/main/java/net/fabricmc/loom/task/GenerateSourcesTask.java

@@ -488,11 +488,7 @@ public abstract class GenerateSourcesTask extends AbstractLoomTask {
 	private void remapLineNumbers(ClassLineNumbers lineNumbers, Path inputJar, Path outputJar) throws IOException {
 		Objects.requireNonNull(lineNumbers, "lineNumbers");
 		final var remapper = new LineNumberRemapper(lineNumbers);
-
-		try (FileSystemUtil.Delegate inFs = FileSystemUtil.getJarFileSystem(inputJar, false);
-				FileSystemUtil.Delegate outFs = FileSystemUtil.getJarFileSystem(outputJar, true)) {
-			remapper.process(inFs.get().getPath("/"), outFs.get().getPath("/"));
-		}
+		remapper.process(inputJar, outputJar);
 	}
 
 	private void doWork(@Nullable IPCServer ipcServer, Path inputJar, Path outputJar, Path linemapFile, @Nullable Path existingJar) {

src/main/java/net/fabricmc/loom/util/AsyncZipProcessor.java

@@ -0,0 +1,109 @@
+/*
+ * This file is part of fabric-loom, licensed under the MIT License (MIT).
+ *
+ * Copyright (c) 2024 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;
+
+import java.io.IOException;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.nio.file.FileVisitResult;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.SimpleFileVisitor;
+import java.nio.file.attribute.BasicFileAttributes;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+
+import org.gradle.api.JavaVersion;
+
+public interface AsyncZipProcessor {
+	static void processEntries(Path inputZip, Path outputZip, AsyncZipProcessor processor) throws IOException {
+		try (FileSystemUtil.Delegate inFs = FileSystemUtil.getJarFileSystem(inputZip, false);
+				FileSystemUtil.Delegate outFs = FileSystemUtil.getJarFileSystem(outputZip, true)) {
+			final Path inRoot = inFs.get().getPath("/");
+			final Path outRoot = outFs.get().getPath("/");
+
+			List<CompletableFuture<Void>> futures = new ArrayList<>();
+			final ExecutorService executor = getExecutor();
+
+			Files.walkFileTree(inRoot, new SimpleFileVisitor<>() {
+				@Override
+				public FileVisitResult visitFile(Path inputFile, BasicFileAttributes attrs) throws IOException {
+					final CompletableFuture<Void> future = CompletableFuture.supplyAsync(() -> {
+						try {
+							final String rel = inRoot.relativize(inputFile).toString();
+							final Path outputFile = outRoot.resolve(rel);
+							processor.processEntryAsync(inputFile, outputFile);
+						} catch (IOException e) {
+							throw new CompletionException(e);
+						}
+
+						return null;
+					});
+
+					futures.add(future);
+					return FileVisitResult.CONTINUE;
+				}
+			});
+
+			// Wait for all futures to complete
+			for (CompletableFuture<Void> future : futures) {
+				try {
+					future.join();
+				} catch (CompletionException e) {
+					if (e.getCause() instanceof IOException ioe) {
+						throw ioe;
+					}
+
+					throw new RuntimeException("Failed to process zip", e.getCause());
+				}
+			}
+
+			executor.shutdown();
+		}
+	}
+
+	/**
+	 * On Java 21 return the virtual thread pool, otherwise return a fixed thread pool.
+	 */
+	private static ExecutorService getExecutor() {
+		if (JavaVersion.current().isCompatibleWith(JavaVersion.VERSION_21)) {
+			// I'm not sure if this is actually faster, but its fun to use loom in loom :D
+			try {
+				Method m = Executors.class.getMethod("newVirtualThreadPerTaskExecutor");
+				return (ExecutorService) m.invoke(null);
+			} catch (NoSuchMethodException | InvocationTargetException | IllegalAccessException e) {
+				throw new RuntimeException("Failed to create virtual thread executor", e);
+			}
+		}
+
+		return Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());
+	}
+
+	void processEntryAsync(Path inputEntry, Path outputEntry) throws IOException;
+}

src/test/groovy/net/fabricmc/loom/test/integration/DecompileTest.groovy

@@ -77,7 +77,7 @@ class DecompileTest extends Specification implements GradleProjectTestTrait {
 
 	def "decompile cache"() {
 		setup:
-		def gradle = gradleProject(project: "minimalBase", version: PRE_RELEASE_GRADLE)
+		def gradle = gradleProject(project: "minimalBase", version: PRE_RELEASE_GRADLE, gradleHomeDir: File.createTempDir())
 		gradle.buildSrc("decompile")
 		gradle.buildGradle << '''
                 dependencies {

src/test/groovy/net/fabricmc/loom/test/unit/AsyncZipProcessorTest.groovy

@@ -0,0 +1,78 @@
+/*
+ * This file is part of fabric-loom, licensed under the MIT License (MIT).
+ *
+ * Copyright (c) 2024 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.test.unit
+
+import java.nio.file.Files
+import java.nio.file.Path
+
+import spock.lang.Specification
+
+import net.fabricmc.loom.test.util.ZipTestUtils
+import net.fabricmc.loom.util.AsyncZipProcessor
+import net.fabricmc.loom.util.ZipUtils
+
+class AsyncZipProcessorTest extends Specification {
+	def "process async"() {
+		given:
+		def inputZip = ZipTestUtils.createZip(createEntries())
+		def outputZip = ZipTestUtils.createZip(Collections.emptyMap())
+		Files.delete(outputZip)
+
+		when:
+		// Process the input zip asynchronously, converting all entries to uppercase
+		AsyncZipProcessor.processEntries(inputZip, outputZip) { Path inputEntry, Path outputEntry ->
+			def str = Files.readString(inputEntry)
+			Files.writeString(outputEntry, str.toUpperCase())
+		}
+
+		then:
+		ZipUtils.unpack(outputZip, "file1.txt") == "FILE1".bytes
+		ZipUtils.unpack(outputZip, "file500.txt") == "FILE500".bytes
+		ZipUtils.unpack(outputZip, "file800.txt") == "FILE800".bytes
+	}
+
+	def "re throws"() {
+		given:
+		def inputZip = ZipTestUtils.createZip(createEntries())
+		def outputZip = ZipTestUtils.createZip(Collections.emptyMap())
+		Files.delete(outputZip)
+
+		when:
+		AsyncZipProcessor.processEntries(inputZip, outputZip) { Path inputEntry, Path outputEntry ->
+			throw new IOException("Test exception")
+		}
+
+		then:
+		thrown(IOException)
+	}
+
+	Map<String, String> createEntries(int count = 1000) {
+		Map<String, String> entries = [:]
+		for (int i = 0; i < count; i++) {
+			entries.put("file" + i + ".txt", "file$i")
+		}
+		return entries
+	}
+}

src/test/groovy/net/fabricmc/loom/test/util/GradleProjectTestTrait.groovy

@@ -45,7 +45,7 @@ trait GradleProjectTestTrait {
 		String gradleVersion = options.version as String ?: LoomTestConstants.DEFAULT_GRADLE
 		String warningMode = options.warningMode as String ?: "fail"
 		File projectDir = options.projectDir as File ?: options.sharedFiles ? sharedProjectDir : File.createTempDir()
-		File gradleHomeDir = gradleHomeDir
+		File gradleHomeDir = options.gradleHomeDir as File ?: gradleHomeDir
 
 		setupProject(options, projectDir)