1 files changed, 18 insertions, 42 deletions

diff --git a/src/main/java/net/fabricmc/loom/build/MixinRefmapHelper.java b/src/main/java/net/fabricmc/loom/build/MixinRefmapHelper.java
index bc775746..458a5ffb 100644
--- a/src/main/java/net/fabricmc/loom/build/MixinRefmapHelper.java
+++ b/src/main/java/net/fabricmc/loom/build/MixinRefmapHelper.java
@@ -25,71 +25,47 @@
 package net.fabricmc.loom.build;
 
 import java.io.File;
-import java.io.IOException;
-import java.io.InputStreamReader;
 import java.nio.file.Path;
-import java.util.HashSet;
-import java.util.Set;
+import java.util.Objects;
+import java.util.stream.Stream;
 import java.util.zip.ZipEntry;
 
 import com.google.gson.JsonObject;
+import org.gradle.api.Project;
 import org.zeroturnaround.zip.ZipUtil;
 import org.zeroturnaround.zip.transform.StringZipEntryTransformer;
 import org.zeroturnaround.zip.transform.ZipEntryTransformerEntry;
 
+import net.fabricmc.loom.LoomGradleExtension;
 import net.fabricmc.loom.LoomGradlePlugin;
+import net.fabricmc.loom.extension.MixinApExtension;
 
 public final class MixinRefmapHelper {
 	private MixinRefmapHelper() { }
 
-	public static boolean addRefmapName(String filename, Path outputPath) {
+	public static boolean addRefmapName(Project project, Path outputPath) {
+		MixinApExtension mixin = LoomGradleExtension.get(project).getMixinApExtension();
 		File output = outputPath.toFile();
-		Set<String> mixinFilenames = findMixins(output, true);
 
-		if (mixinFilenames.size() > 0) {
-			return ZipUtil.transformEntries(output, mixinFilenames.stream().map((f) -> new ZipEntryTransformerEntry(f, new StringZipEntryTransformer("UTF-8") {
+		return mixin.getMixinSourceSetsStream().map(sourceSet -> {
+			MixinApExtension.MixinInformationContainer container = Objects.requireNonNull(
+					MixinApExtension.getMixinInformationContainer(sourceSet)
+			);
+			Stream<String> mixinJsonNames = container.getMixinJsonNames();
+			String refmapName = container.getRefmapName();
+
+			return ZipUtil.transformEntries(output, mixinJsonNames.map(f -> new ZipEntryTransformerEntry(f, new StringZipEntryTransformer("UTF-8") {
 				@Override
-				protected String transform(ZipEntry zipEntry, String input) throws IOException {
+				protected String transform(ZipEntry zipEntry, String input) {
 					JsonObject json = LoomGradlePlugin.GSON.fromJson(input, JsonObject.class);
 
 					if (!json.has("refmap")) {
-						json.addProperty("refmap", filename);
+						json.addProperty("refmap", refmapName);
 					}
 
 					return LoomGradlePlugin.GSON.toJson(json);
 				}
 			})).toArray(ZipEntryTransformerEntry[]::new));
-		} else {
-			return false;
-		}
-	}
-
-	private static Set<String> findMixins(File output, boolean onlyWithoutRefmap) {
-		// first, identify all of the mixin files
-		Set<String> mixinFilename = new HashSet<>();
-		// TODO: this is a lovely hack
-		ZipUtil.iterate(output, (stream, entry) -> {
-			if (!entry.isDirectory() && entry.getName().endsWith(".json") && !entry.getName().contains("/") && !entry.getName().contains("\\")) {
-				// JSON file in root directory
-				try (InputStreamReader inputStreamReader = new InputStreamReader(stream)) {
-					JsonObject json = LoomGradlePlugin.GSON.fromJson(inputStreamReader, JsonObject.class);
-
-					if (json != null) {
-						boolean hasMixins = json.has("mixins") && json.get("mixins").isJsonArray();
-						boolean hasClient = json.has("client") && json.get("client").isJsonArray();
-						boolean hasServer = json.has("server") && json.get("server").isJsonArray();
-
-						if (json.has("package") && (hasMixins || hasClient || hasServer)) {
-							if (!onlyWithoutRefmap || !json.has("refmap") || !json.has("minVersion")) {
-								mixinFilename.add(entry.getName());
-							}
-						}
-					}
-				} catch (Exception ignored) {
-					// ...
-				}
-			}
-		});
-		return mixinFilename;
+		}).reduce(false, Boolean::logicalOr);
 	}
 }