Fix support for Forge 1.18

Signed-off-by: shedaniel <daniel@shedaniel.me>

4 files changed, 148 insertions, 10 deletions

diff --git a/src/main/java/net/fabricmc/loom/configuration/providers/MinecraftProviderImpl.java b/src/main/java/net/fabricmc/loom/configuration/providers/MinecraftProviderImpl.java
index 1967bb72..c8987bfe 100644
--- a/src/main/java/net/fabricmc/loom/configuration/providers/MinecraftProviderImpl.java
+++ b/src/main/java/net/fabricmc/loom/configuration/providers/MinecraftProviderImpl.java
@@ -52,6 +52,7 @@ import net.fabricmc.loom.util.Constants;
 import net.fabricmc.loom.util.MirrorUtil;
 import net.fabricmc.loom.util.DownloadUtil;
 import net.fabricmc.loom.util.HashedDownloadUtil;
+import net.fabricmc.loom.util.ZipUtils;
 import net.fabricmc.stitch.merge.JarMerger;
 
 public class MinecraftProviderImpl extends DependencyProvider implements MinecraftProvider {
@@ -67,6 +68,7 @@ public class MinecraftProviderImpl extends DependencyProvider implements Minecra
 	public File minecraftServerJar;
 	// The extracted server jar from the boostrap, only exists in >=21w39a
 	public File minecraftExtractedServerJar;
+	private Boolean isNewerThan21w39a;
 	private File minecraftMergedJar;
 	private File versionManifestJson;
 	private File experimentalVersionsJson;
@@ -332,6 +334,26 @@ public class MinecraftProviderImpl extends DependencyProvider implements Minecra
 		}
 	}
 
+	public File getMinecraftServerJar() {
+		if (isNewerThan21w39a()) {
+			try {
+				return getServerJarToMerge(getProject().getLogger());
+			} catch (IOException e) {
+				throw new RuntimeException(e);
+			}
+		}
+
+		return minecraftServerJar;
+	}
+
+	public boolean isNewerThan21w39a() {
+		if (isNewerThan21w39a != null) {
+			return isNewerThan21w39a;
+		}
+		
+		return isNewerThan21w39a = ZipUtils.contains(minecraftServerJar.toPath(), "META-INF/versions.list");
+	}
+
 	public File getMergedJar() {
 		return minecraftMergedJar;
 	}
diff --git a/src/main/java/net/fabricmc/loom/configuration/providers/forge/ForgeUserdevProvider.java b/src/main/java/net/fabricmc/loom/configuration/providers/forge/ForgeUserdevProvider.java
index a0d247b5..fecfa4d3 100644
--- a/src/main/java/net/fabricmc/loom/configuration/providers/forge/ForgeUserdevProvider.java
+++ b/src/main/java/net/fabricmc/loom/configuration/providers/forge/ForgeUserdevProvider.java
@@ -135,7 +135,6 @@ public class ForgeUserdevProvider extends DependencyProvider {
 			}
 		}
 
-		// TODO: Read launch configs from the JSON too
 		// TODO: Should I copy the patches from here as well?
 		//       That'd require me to run the "MCP environment" fully up to merging.
 		for (Map.Entry<String, JsonElement> entry : json.getAsJsonObject("runs").entrySet()) {
diff --git a/src/main/java/net/fabricmc/loom/configuration/providers/forge/McpConfigProvider.java b/src/main/java/net/fabricmc/loom/configuration/providers/forge/McpConfigProvider.java
index ae6e0e87..e357a15d 100644
--- a/src/main/java/net/fabricmc/loom/configuration/providers/forge/McpConfigProvider.java
+++ b/src/main/java/net/fabricmc/loom/configuration/providers/forge/McpConfigProvider.java
@@ -24,20 +24,30 @@
 
 package net.fabricmc.loom.configuration.providers.forge;
 
+import java.io.ByteArrayInputStream;
+import java.io.File;
 import java.io.IOException;
 import java.io.Reader;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.StandardCopyOption;
 import java.nio.file.StandardOpenOption;
+import java.util.List;
 import java.util.function.Consumer;
+import java.util.jar.Attributes;
+import java.util.jar.Manifest;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
 
 import com.google.gson.Gson;
+import com.google.gson.JsonElement;
 import com.google.gson.JsonObject;
 import org.gradle.api.Project;
+import org.gradle.api.file.FileCollection;
 
 import net.fabricmc.loom.configuration.DependencyProvider;
 import net.fabricmc.loom.util.Constants;
+import net.fabricmc.loom.util.DependencyDownloader;
 import net.fabricmc.loom.util.ZipUtils;
 
 public class McpConfigProvider extends DependencyProvider {
@@ -46,6 +56,7 @@ public class McpConfigProvider extends DependencyProvider {
 	private Path mappings;
 	private Boolean official;
 	private String mappingsPath;
+	private RemapAction remapAction;
 
 	public McpConfigProvider(Project project) {
 		super(project);
@@ -70,6 +81,22 @@ public class McpConfigProvider extends DependencyProvider {
 
 		official = json.has("official") && json.getAsJsonPrimitive("official").getAsBoolean();
 		mappingsPath = json.get("data").getAsJsonObject().get("mappings").getAsString();
+
+		if (json.has("functions")) {
+			JsonObject functions = json.getAsJsonObject("functions");
+
+			if (functions.has("rename")) {
+				remapAction = new ConfigDefinedRemapAction(getProject(), functions.getAsJsonObject("rename"));
+			}
+		}
+
+		if (remapAction == null) {
+			throw new RuntimeException("Could not find remap action, this is probably a version Architectury Loom does not support!");
+		}
+	}
+
+	public RemapAction getRemapAction() {
+		return remapAction;
 	}
 
 	private void init(String version) throws IOException {
@@ -111,4 +138,95 @@ public class McpConfigProvider extends DependencyProvider {
 	public String getTargetConfig() {
 		return Constants.Configurations.MCP_CONFIG;
 	}
+
+	public interface RemapAction {
+		FileCollection getClasspath();
+
+		String getMainClass();
+
+		List<String> getArgs(Path input, Path output, Path mappings, FileCollection libraries);
+	}
+
+	public static class ConfigDefinedRemapAction implements RemapAction {
+		private final Project project;
+		private final String name;
+		private final File mainClasspath;
+		private final FileCollection classpath;
+		private final List<String> args;
+		private boolean hasLibraries;
+
+		public ConfigDefinedRemapAction(Project project, JsonObject json) {
+			this.project = project;
+			this.name = json.get("version").getAsString();
+			this.mainClasspath = DependencyDownloader.download(project, this.name, false, true)
+					.getSingleFile();
+			this.classpath = DependencyDownloader.download(project, this.name, true, true);
+			this.args = StreamSupport.stream(json.getAsJsonArray("args").spliterator(), false)
+					.map(JsonElement::getAsString)
+					.collect(Collectors.toList());
+			for (int i = 1; i < this.args.size(); i++) {
+				if (this.args.get(i).equals("{libraries}")) {
+					this.args.remove(i);
+					this.args.remove(i - 1);
+					this.hasLibraries = true;
+					break;
+				}
+			}
+		}
+
+		@Override
+		public FileCollection getClasspath() {
+			return classpath;
+		}
+
+		@Override
+		public String getMainClass() {
+			try {
+				byte[] manifestBytes = ZipUtils.unpackNullable(mainClasspath.toPath(), "META-INF/MANIFEST.MF");
+
+				if (manifestBytes == null) {
+					throw new RuntimeException("Could not find MANIFEST.MF in " + mainClasspath + "!");
+				}
+
+				Manifest manifest = new Manifest(new ByteArrayInputStream(manifestBytes));
+				Attributes attributes = manifest.getMainAttributes();
+				String value = attributes.getValue(Attributes.Name.MAIN_CLASS);
+
+				if (value == null) {
+					throw new RuntimeException("Could not find main class in " + mainClasspath + "!");
+				} else {
+					return value;
+				}
+			} catch (IOException e) {
+				throw new RuntimeException(e);
+			}
+		}
+
+		@Override
+		public List<String> getArgs(Path input, Path output, Path mappings, FileCollection libraries) {
+			List<String> args = this.args.stream()
+					.map(str -> {
+						return switch (str) {
+							case "{input}" -> input.toAbsolutePath().toString();
+							case "{output}" -> output.toAbsolutePath().toString();
+							case "{mappings}" -> mappings.toAbsolutePath().toString();
+							default -> str;
+						};
+					})
+					.collect(Collectors.toList());
+
+			if (hasLibraries) {
+				for (File file : libraries) {
+					args.add("-e=" + file.getAbsolutePath());
+				}
+			}
+
+			return args;
+		}
+
+		@Override
+		public String toString() {
+			return this.name;
+		}
+	}
 }
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 02a5cce4..a74b49d0 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
@@ -87,6 +87,7 @@ import org.objectweb.asm.tree.ClassNode;
 
 import net.fabricmc.loom.configuration.DependencyProvider;
 import net.fabricmc.loom.configuration.providers.MinecraftProviderImpl;
+import net.fabricmc.loom.extension.LoomGradleExtensionImpl;
 import net.fabricmc.loom.util.Constants;
 import net.fabricmc.loom.util.DependencyDownloader;
 import net.fabricmc.loom.util.FileSystemUtil;
@@ -338,20 +339,17 @@ public class MinecraftPatchedProvider extends DependencyProvider {
 
 	private void createSrgJars(Logger logger) throws Exception {
 		MinecraftProviderImpl minecraftProvider = getExtension().getMinecraftProvider();
-		String dep = getExtension().isForgeAndOfficial() ? Constants.Dependencies.VIGNETTE + Constants.Dependencies.Versions.VIGNETTE
-				: Constants.Dependencies.SPECIAL_SOURCE + Constants.Dependencies.Versions.SPECIAL_SOURCE + ":shaded";
-		FileCollection classpath = DependencyDownloader.download(getProject(), dep, true, true);
-		produceSrgJar(getExtension().isForgeAndOfficial(), minecraftProvider.minecraftClientJar.toPath(), minecraftProvider.minecraftServerJar.toPath(), classpath);
+		produceSrgJar(getExtension().isForgeAndOfficial(), minecraftProvider.minecraftClientJar.toPath(), minecraftProvider.getMinecraftServerJar().toPath());
 	}
 
-	private void produceSrgJar(boolean official, Path clientJar, Path serverJar, FileCollection classpath) throws IOException {
+	private void produceSrgJar(boolean official, Path clientJar, Path serverJar) throws IOException {
 		Path tmpSrg = getToSrgMappings();
 		Set<File> mcLibs = getProject().getConfigurations().getByName(Constants.Configurations.MINECRAFT_DEPENDENCIES).resolve();
 
 		ThreadingUtils.run(() -> {
-			Files.copy(SpecialSourceExecutor.produceSrgJar(getExtension().isForgeAndNotOfficial(), getProject(), "client", classpath, mcLibs, clientJar, tmpSrg), minecraftClientSrgJar.toPath());
+			Files.copy(SpecialSourceExecutor.produceSrgJar(getExtension().getMcpConfigProvider().getRemapAction(), getProject(), "client", mcLibs, clientJar, tmpSrg), minecraftClientSrgJar.toPath());
 		}, () -> {
-				Files.copy(SpecialSourceExecutor.produceSrgJar(getExtension().isForgeAndNotOfficial(), getProject(), "server", classpath, mcLibs, serverJar, tmpSrg), minecraftServerSrgJar.toPath());
+				Files.copy(SpecialSourceExecutor.produceSrgJar(getExtension().getMcpConfigProvider().getRemapAction(), getProject(), "server", mcLibs, serverJar, tmpSrg), minecraftServerSrgJar.toPath());
 			});
 	}
 
@@ -481,8 +479,9 @@ public class MinecraftPatchedProvider extends DependencyProvider {
 	}
 
 	private void accessTransformForge(Logger logger) throws Exception {
+		MinecraftProviderImpl minecraftProvider = getExtension().getMinecraftProvider();
 		List<File> toDelete = new ArrayList<>();
-		String atDependency = Constants.Dependencies.ACCESS_TRANSFORMERS + Constants.Dependencies.Versions.ACCESS_TRANSFORMERS;
+		String atDependency = Constants.Dependencies.ACCESS_TRANSFORMERS + (minecraftProvider.isNewerThan21w39a() ? Constants.Dependencies.Versions.ACCESS_TRANSFORMERS_NEW : Constants.Dependencies.Versions.ACCESS_TRANSFORMERS);
 		FileCollection classpath = DependencyDownloader.download(getProject(), atDependency);
 		Stopwatch stopwatch = Stopwatch.createStarted();
 
@@ -655,7 +654,7 @@ public class MinecraftPatchedProvider extends DependencyProvider {
 			// Copy resources
 			MinecraftProviderImpl minecraftProvider = getExtension().getMinecraftProvider();
 			copyNonClassFiles(minecraftProvider.minecraftClientJar, minecraftMergedPatchedSrgJar);
-			copyNonClassFiles(minecraftProvider.minecraftServerJar, minecraftMergedPatchedSrgJar);
+			copyNonClassFiles(minecraftProvider.getMinecraftServerJar(), minecraftMergedPatchedSrgJar);
 		}
 	}