WIP AT -> JarProcessor

1 files changed, 66 insertions, 69 deletions

diff --git a/src/main/java/net/fabricmc/loom/util/srg/AtRemapper.java b/src/main/java/net/fabricmc/loom/util/srg/AtRemapper.java
index 467c15a4..e14afeb9 100644
--- a/src/main/java/net/fabricmc/loom/util/srg/AtRemapper.java
+++ b/src/main/java/net/fabricmc/loom/util/srg/AtRemapper.java
@@ -24,111 +24,108 @@
 
 package net.fabricmc.loom.util.srg;
 
+import static me.shedaniel.architectury.refmapremapper.utils.DescriptorRemapper.remapDescriptor;
+
 import java.io.IOException;
-import java.io.StringReader;
 import java.nio.file.Path;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.function.UnaryOperator;
+import java.util.Objects;
 import java.util.zip.ZipEntry;
 
+import me.shedaniel.architectury.refmapremapper.remapper.SimpleReferenceRemapper;
 import org.apache.logging.log4j.util.Strings;
 import org.gradle.api.logging.Logger;
+import org.jetbrains.annotations.Nullable;
 import org.zeroturnaround.zip.ZipUtil;
 import org.zeroturnaround.zip.transform.StringZipEntryTransformer;
-import org.zeroturnaround.zip.transform.ZipEntryTransformerEntry;
 
-import net.fabricmc.loom.util.function.CollectionUtil;
 import net.fabricmc.mapping.tree.TinyTree;
 
 /**
- * Remaps AT classes from SRG to Yarn.
+ * Remaps AT contents.
  *
  * @author Juuz
  */
 public final class AtRemapper {
-	public static void remap(Logger logger, Path jar, TinyTree mappings) throws IOException {
-		ZipUtil.transformEntries(jar.toFile(), new ZipEntryTransformerEntry[] {(new ZipEntryTransformerEntry("META-INF/accesstransformer.cfg", new StringZipEntryTransformer() {
+	public static void remapSrgToNamed(Logger logger, Path jar, TinyTree mappings) throws IOException {
+		ZipUtil.transformEntry(jar.toFile(), "META-INF/accesstransformer.cfg", new StringZipEntryTransformer() {
 			@Override
 			protected String transform(ZipEntry zipEntry, String input) {
-				String[] lines = input.split("\n");
-				List<String> output = new ArrayList<>(lines.length);
-
-				for (int i = 0; i < lines.length; i++) {
-					String line = lines[i].trim();
-
-					if (line.startsWith("#") || Strings.isBlank(line)) {
-						output.add(i, line);
-						continue;
+				return remapAt(logger, input, new SimpleReferenceRemapper.Remapper() {
+					@Override
+					@Nullable
+					public String mapClass(String value) {
+						return mappings.getClasses().stream()
+								.filter(classDef -> Objects.equals(classDef.getName("srg"), value))
+								.findFirst()
+								.map(classDef -> classDef.getName("named"))
+								.orElse(null);
 					}
 
-					String[] parts = line.split("\\s+");
-
-					if (parts.length < 2) {
-						logger.warn("Invalid AT Line: " + line);
-						output.add(i, line);
-						continue;
+					@Override
+					@Nullable
+					public String mapMethod(@Nullable String className, String methodName, String methodDescriptor) {
+						return null;
 					}
 
-					String name = parts[1].replace('.', '/');
-					parts[1] = CollectionUtil.find(
-							mappings.getClasses(),
-							def -> def.getName("srg").equals(name)
-					).map(def -> def.getName("named")).orElse(name).replace('/', '.');
-
-					if (parts.length >= 3) {
-						if (parts[2].contains("(")) {
-							parts[2] = parts[2].substring(0, parts[2].indexOf('(')) + remapDescriptor(parts[2].substring(parts[2].indexOf('(')), s -> {
-								return CollectionUtil.find(
-										mappings.getClasses(),
-										def -> def.getName("srg").equals(s)
-								).map(def -> def.getName("named")).orElse(s);
-							});
-						}
+					@Override
+					@Nullable
+					public String mapField(@Nullable String className, String fieldName, String fieldDescriptor) {
+						return null;
 					}
+				});
+			}
+		});
+	}
+
+	public static String remapAt(Logger logger, String sourceAt, SimpleReferenceRemapper.Remapper remapper) {
+		String[] lines = sourceAt.split("\n");
+		StringBuilder builder = new StringBuilder();
 
-					output.add(i, String.join(" ", parts));
+		for (String line : lines) {
+			{
+				int indexOf = line.indexOf('#');
+
+				if (indexOf != -1) {
+					line = line.substring(0, indexOf);
 				}
 
-				return String.join("\n", output);
+				line = line.trim();
 			}
-		}))});
-	}
 
-	private static String remapDescriptor(String original, UnaryOperator<String> classMappings) {
-		try {
-			StringReader reader = new StringReader(original);
-			StringBuilder result = new StringBuilder();
-			boolean insideClassName = false;
-			StringBuilder className = new StringBuilder();
+			if (Strings.isBlank(line)) {
+				builder.append(line).append('\n');
+				continue;
+			}
 
-			while (true) {
-				int c = reader.read();
+			String[] parts = line.split("\\s+");
 
-				if (c == -1) {
-					break;
-				}
+			if (parts.length < 2) {
+				logger.warn("Invalid AT Line: " + line);
+				builder.append(line).append('\n');
+				continue;
+			}
 
-				if ((char) c == ';') {
-					insideClassName = false;
-					result.append(classMappings.apply(className.toString()));
-				}
+			String originalClassName = parts[1].replace('.', '/');
+			parts[1] = either(remapper.mapClass(originalClassName), parts[1]).replace('/', '.');
 
-				if (insideClassName) {
-					className.append((char) c);
+			if (parts.length >= 3) {
+				if (parts[2].contains("(")) {
+					String methodName = parts[2].substring(0, parts[2].indexOf('('));
+					String methodDescriptor = parts[2].substring(parts[2].indexOf('('));
+					parts[2] = either(remapper.mapMethod(originalClassName, methodName, methodDescriptor), methodName)
+							+ remapDescriptor(methodDescriptor, it -> either(remapper.mapClass(it), it));
 				} else {
-					result.append((char) c);
-				}
-
-				if (!insideClassName && (char) c == 'L') {
-					insideClassName = true;
-					className.setLength(0);
+					parts[2] = either(remapper.mapField(originalClassName, parts[2], null), parts[2]);
 				}
 			}
 
-			return result.toString();
-		} catch (IOException e) {
-			throw new AssertionError(e);
+			builder.append(String.join(" ", parts)).append('\n');
 		}
+
+		return builder.toString();
+	}
+
+	private static <T> T either(@Nullable T first, @Nullable T second) {
+		return first == null ? second : first;
 	}
 }