Read loader access transformers from userdev config

src/main/java/dev/architectury/loom/forge/UserdevConfig.java

@@ -1,13 +1,17 @@
 package dev.architectury.loom.forge;
 
+import java.io.IOException;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Function;
 
+import com.mojang.datafixers.util.Either;
 import com.mojang.serialization.Codec;
 import com.mojang.serialization.codecs.RecordCodecBuilder;
 
 import net.fabricmc.loom.configuration.providers.forge.ForgeRunTemplate;
+import net.fabricmc.loom.util.IOFunction;
 
 public record UserdevConfig(
 		String mcp,
@@ -20,7 +24,8 @@ public record UserdevConfig(
 		BinaryPatcherConfig binpatcher,
 		List<String> libraries,
 		Map<String, ForgeRunTemplate> runs,
-		List<String> sass
+		List<String> sass,
+		AccessTransformerLocation ats
 ) {
 	public static final Codec<UserdevConfig> CODEC = RecordCodecBuilder.create(instance -> instance.group(
 			Codec.STRING.fieldOf("mcp").forGetter(UserdevConfig::mcp),
@@ -33,7 +38,8 @@ public record UserdevConfig(
 			BinaryPatcherConfig.CODEC.fieldOf("binpatcher").forGetter(UserdevConfig::binpatcher),
 			Codec.STRING.listOf().fieldOf("libraries").forGetter(UserdevConfig::libraries),
 			ForgeRunTemplate.MAP_CODEC.fieldOf("runs").forGetter(UserdevConfig::runs),
-			Codec.STRING.listOf().optionalFieldOf("sass", List.of()).forGetter(UserdevConfig::sass)
+			Codec.STRING.listOf().optionalFieldOf("sass", List.of()).forGetter(UserdevConfig::sass),
+			AccessTransformerLocation.CODEC.fieldOf("ats").forGetter(UserdevConfig::ats)
 	).apply(instance, UserdevConfig::new));
 
 	public record BinaryPatcherConfig(String dependency, List<String> args) {
@@ -42,4 +48,38 @@ public record BinaryPatcherConfig(String dependency, List<String> args) {
 				Codec.STRING.listOf().fieldOf("args").forGetter(BinaryPatcherConfig::args)
 		).apply(instance, BinaryPatcherConfig::new));
 	}
+
+	public sealed interface AccessTransformerLocation {
+		Codec<AccessTransformerLocation> CODEC = Codec.either(Codec.STRING, Codec.STRING.listOf()).xmap(
+				either -> either.map(Directory::new, FileList::new),
+				location -> location.visit(Either::left, Either::right)
+		);
+
+		<T> T visit(Function<String, T> ifDirectory, Function<List<String>, T> ifFileList);
+		<T> T visitIo(IOFunction<String, T> ifDirectory, IOFunction<List<String>, T> ifFileList) throws IOException;
+
+		record Directory(String path) implements AccessTransformerLocation {
+			@Override
+			public <T> T visit(Function<String, T> ifDirectory, Function<List<String>, T> ifFileList) {
+				return ifDirectory.apply(path);
+			}
+
+			@Override
+			public <T> T visitIo(IOFunction<String, T> ifDirectory, IOFunction<List<String>, T> ifFileList) throws IOException {
+				return ifDirectory.apply(path);
+			}
+		}
+
+		record FileList(List<String> paths) implements AccessTransformerLocation {
+			@Override
+			public <T> T visit(Function<String, T> ifDirectory, Function<List<String>, T> ifFileList) {
+				return ifFileList.apply(paths);
+			}
+
+			@Override
+			public <T> T visitIo(IOFunction<String, T> ifDirectory, IOFunction<List<String>, T> ifFileList) throws IOException {
+				return ifFileList.apply(paths);
+			}
+		}
+	}
 }

src/main/java/net/fabricmc/loom/configuration/accesstransformer/AccessTransformerJarProcessor.java

@@ -97,6 +97,7 @@ public AccessTransformerJarProcessor(String name, Project project, Iterable<File
 			final byte[] bytes;
 
 			try {
+				// TODO: Shouldn't we check for the mods.toml AT list on Neo?
 				bytes = localMod.getSource().read(Constants.Forge.ACCESS_TRANSFORMER_PATH);
 			} catch (FileNotFoundException | NoSuchFileException e) {
 				continue;

src/main/java/net/fabricmc/loom/configuration/providers/forge/MinecraftPatchedProvider.java

@@ -1,7 +1,7 @@
 /*
  * This file is part of fabric-loom, licensed under the MIT License (MIT).
  *
- * Copyright (c) 2020-2023 FabricMC
+ * Copyright (c) 2020-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
@@ -32,6 +32,7 @@
 import java.io.UncheckedIOException;
 import java.nio.file.FileSystem;
 import java.nio.file.Files;
+import java.nio.file.NoSuchFileException;
 import java.nio.file.Path;
 import java.nio.file.StandardCopyOption;
 import java.nio.file.StandardOpenOption;
@@ -369,34 +370,53 @@ public static void accessTransform(Project project, Path input, Path target) thr
 		project.getLogger().lifecycle(":access transforming minecraft");
 
 		LoomGradleExtension extension = LoomGradleExtension.get(project);
-		List<Path> atSources = List.of(
-				extension.getForgeUniversalProvider().getForge().toPath(),
-				extension.getForgeUserdevProvider().getUserdevJar().toPath(),
-				((ForgeMinecraftProvider) extension.getMinecraftProvider())
-						.getPatchedProvider()
-						.getMinecraftPatchedIntermediateJar()
-		);
-
+		Path userdevJar = extension.getForgeUserdevProvider().getUserdevJar().toPath();
 		Files.deleteIfExists(target);
 
 		try (var tempFiles = new TempFiles()) {
 			AccessTransformerJarProcessor.executeAt(project, input, target, args -> {
-				for (Path jar : atSources) {
-					byte[] atBytes = ZipUtils.unpackNullable(jar, Constants.Forge.ACCESS_TRANSFORMER_PATH);
-
-					if (atBytes != null) {
-						Path tmpFile = tempFiles.file("at-conf", ".cfg");
-						Files.write(tmpFile, atBytes, StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING);
-						args.add("--atFile");
-						args.add(tmpFile.toAbsolutePath().toString());
-					}
+				for (String atFile : extractAccessTransformers(userdevJar, extension.getForgeUserdevProvider().getConfig().ats(), tempFiles)) {
+					args.add("--atFile");
+					args.add(atFile);
 				}
 			});
 		}
 
 		project.getLogger().lifecycle(":access transformed minecraft in " + stopwatch.stop());
 	}
 
+	private static List<String> extractAccessTransformers(Path jar, UserdevConfig.AccessTransformerLocation location, TempFiles tempFiles) throws IOException {
+		final List<String> extracted = new ArrayList<>();
+
+		try (FileSystemUtil.Delegate fs = FileSystemUtil.getJarFileSystem(jar)) {
+			for (Path atFile : getAccessTransformerPaths(fs, location)) {
+				byte[] atBytes;
+
+				try {
+					atBytes = Files.readAllBytes(atFile);
+				} catch (NoSuchFileException e) {
+					continue;
+				}
+
+				Path tmpFile = tempFiles.file("at-conf", ".cfg");
+				Files.write(tmpFile, atBytes, StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING);
+				extracted.add(tmpFile.toAbsolutePath().toString());
+			}
+		}
+
+		return extracted;
+	}
+
+	private static List<Path> getAccessTransformerPaths(FileSystemUtil.Delegate fs, UserdevConfig.AccessTransformerLocation location) throws IOException {
+		return location.visitIo(directory -> {
+			Path dirPath = fs.getPath(directory);
+
+			try (Stream<Path> paths = Files.list(dirPath)) {
+				return paths.toList();
+			}
+		}, paths -> paths.stream().map(fs::getPath).toList());
+	}
+
 	private void remapPatchedJar(ServiceFactory serviceFactory) throws Exception {
 		logger.lifecycle(":remapping minecraft (TinyRemapper, srg -> official)");
 		Path mcInput = minecraftPatchedIntermediateAtJar;