Add searge enricher

4 files changed, 186 insertions, 0 deletions

diff --git a/build.gradle.kts b/build.gradle.kts
index c273390..0596e52 100644
--- a/build.gradle.kts
+++ b/build.gradle.kts
@@ -17,6 +17,7 @@ dependencies {
     compileOnly("org.jetbrains:annotations:24.0.0")
     implementation("net.minecraftforge:srgutils:0.4.13")
     implementation("net.fabricmc:tiny-remapper:0.8.6")
+    implementation("de.siegmar:fastcsv:3.0.0")
     testImplementation("org.junit.jupiter:junit-jupiter:5.9.2")
 }
 
diff --git a/src/main/java/moe/nea/zwirn/EnrichSeargeWithMCP.java b/src/main/java/moe/nea/zwirn/EnrichSeargeWithMCP.java
new file mode 100644
index 0000000..d91e424
--- /dev/null
+++ b/src/main/java/moe/nea/zwirn/EnrichSeargeWithMCP.java
@@ -0,0 +1,157 @@
+package moe.nea.zwirn;
+
+import de.siegmar.fastcsv.reader.CsvReader;
+import net.fabricmc.stitch.commands.tinyv2.*;
+import org.jetbrains.annotations.NotNull;
+
+import java.io.IOException;
+import java.nio.file.Path;
+import java.util.*;
+import java.util.function.Function;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+
+public class EnrichSeargeWithMCP {
+    private final TinyFile searge;
+    private final List<MCPField> fields;
+    private final List<MCPMethod> methods;
+    private final List<MCPParam> params;
+    private final Map<String, MCPField> fieldMap;
+    private final Map<Integer, List<MCPParam>> paramMap;
+    private final Map<String, MCPMethod> methodMap;
+
+    // TODO: parse joined.exc for constructor indexes parameters
+    public EnrichSeargeWithMCP(@NotNull TinyFile searge, Path fields, Path methods, Path params) throws IOException {
+        this.searge = searge;
+        this.fields = readFields(fields);
+        this.methods = readMethods(methods);
+        this.params = readParams(params);
+        this.fieldMap = this.fields.stream().collect(Collectors.toMap(MCPField::searge, Function.identity()));
+        this.methodMap = this.methods.stream().collect(Collectors.toMap(MCPMethod::searge, Function.identity()));
+        this.paramMap = this.params.stream().filter(it -> it.methodId() != null).collect(Collectors.groupingBy(MCPParam::methodId, Collectors.toList()));
+    }
+
+    record MCPParam(
+            String searge,
+            String name
+    ) {
+        public Integer methodId() {
+            var matcher = SEARGE_PARAM_ID_PATTERN.matcher(searge);
+            if (!matcher.matches())
+                return -1;
+            return Integer.parseInt(matcher.group(1));
+        }
+
+        public Integer lvIndexHeuristic() {
+            var matcher = SEARGE_PARAM_ID_PATTERN.matcher(searge);
+            if (!matcher.matches())
+                return -1;
+            return Integer.parseInt(matcher.group(2));
+        }
+
+        static final Pattern SEARGE_PARAM_ID_PATTERN = Pattern.compile("^p_([0-9]+)_([0-9]+)_$");
+
+    }
+
+    record MCPField(String searge, String name, String desc) {
+    }
+
+    record MCPMethod(String searge, String name, String desc) {
+        public int methodId() {
+            var matcher = SEARGE_METHOD_ID_PATTERN.matcher(searge);
+            if (!matcher.matches())
+                throw new IllegalStateException("Searge name does not contain method id");
+            return Integer.parseInt(matcher.group(1));
+        }
+
+        static final Pattern SEARGE_METHOD_ID_PATTERN = Pattern.compile("^func_([0-9]+)_.+$");
+    }
+
+
+    private static List<MCPParam> readParams(Path params) throws IOException {
+        try (var csvReader = CsvReader.builder().ofNamedCsvRecord(params)
+                .stream()) {
+            // Header: param,name,side
+            return csvReader.map(
+                    it -> new MCPParam(it.getField("param"), it.getField("name"))
+            ).collect(Collectors.toList());
+        }
+    }
+
+    private static List<MCPField> readFields(Path fields) throws IOException {
+        try (var csvReader = CsvReader.builder().ofNamedCsvRecord(fields)
+                .stream()) {
+            // Header: searge,name,side,desc
+            return csvReader.map(
+                    it -> new MCPField(it.getField("searge"), it.getField("name"), it.getField("desc").replace("\\n", "\n"))
+            ).collect(Collectors.toList());
+        }
+    }
+
+    private static List<MCPMethod> readMethods(Path methods) throws IOException {
+        try (var csvReader = CsvReader.builder().ofNamedCsvRecord(methods)
+                .stream()) {
+            // Header: searge,name,side,desc
+            return csvReader.map(
+                    it -> new MCPMethod(it.getField("searge"), it.getField("name"), it.getField("desc").replace("\\n", "\n"))
+            ).collect(Collectors.toList());
+        }
+    }
+
+    public TinyFile mergeTinyFile() {
+        return new TinyFile(
+                new TinyHeader(
+                        Arrays.asList("notch", "searge", "mcp"),
+                        2, 0, new HashMap<>()
+                ),
+                searge.getClassEntries()
+                        .stream().map(this::mergeClass)
+                        .collect(Collectors.toList())
+        );
+    }
+
+    private TinyClass mergeClass(TinyClass tinyClass) {
+        return new TinyClass(
+                Arrays.asList(
+                        tinyClass.getClassNames().get(0), tinyClass.getClassNames().get(1),
+                        tinyClass.getClassNames().get(1) // MCP does not handle class names. those are done in searge
+                ),
+                tinyClass.getMethods().stream().map(this::mergeMethod).collect(Collectors.toList()),
+                tinyClass.getFields().stream().map(this::mergeField).collect(Collectors.toList()),
+                Arrays.asList() // Searge doesn't have comments
+        );
+    }
+
+    private TinyField mergeField(TinyField tinyField) {
+        var srg = tinyField.getFieldNames().get(1);
+        var mcpField = fieldMap.get(srg);
+        return new TinyField(
+                tinyField.getFieldDescriptorInFirstNamespace(),
+                Arrays.asList(tinyField.getFieldNames().get(0), srg, mcpField == null ? srg : mcpField.name()),
+                mcpField == null ? Arrays.asList() : Arrays.asList(mcpField.desc)// TODO: handle empty comment
+        );
+    }
+
+    private TinyMethod mergeMethod(TinyMethod tinyMethod) {
+        var srg = tinyMethod.getMethodNames().get(1);
+        var mcpMethod = methodMap.get(srg);
+        List<TinyMethodParameter> params = new ArrayList<>();
+        if (mcpMethod != null) {
+            var mcpParams = paramMap.get(mcpMethod.methodId());
+            if (mcpParams != null) for (var param : mcpParams) {
+                params.add(new TinyMethodParameter(
+                        param.lvIndexHeuristic(),
+                        Arrays.asList("p" + param.lvIndexHeuristic(), param.searge(), param.name()),
+                        Arrays.asList()
+                ));
+            }
+        }
+        return new TinyMethod(
+                tinyMethod.getMethodDescriptorInFirstNamespace(),
+                Arrays.asList(tinyMethod.getMethodNames().get(0), srg, mcpMethod == null ? srg : mcpMethod.name()),
+                params,
+                Arrays.asList(),
+                mcpMethod == null ? Arrays.asList() : Arrays.asList(mcpMethod.desc) // TODO: handle empty comment
+        );
+    }
+}
diff --git a/src/main/java/moe/nea/zwirn/Zwirn.java b/src/main/java/moe/nea/zwirn/Zwirn.java
index 746c993..3612007 100644
--- a/src/main/java/moe/nea/zwirn/Zwirn.java
+++ b/src/main/java/moe/nea/zwirn/Zwirn.java
@@ -3,6 +3,10 @@ package moe.nea.zwirn;
 import net.fabricmc.stitch.commands.tinyv2.TinyFile;
 import org.jetbrains.annotations.NotNull;
 
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Arrays;
 import java.util.List;
 
 public class Zwirn {
@@ -17,6 +21,21 @@ public class Zwirn {
         return new TinyMerger(base, overlay, sharedNamespace).merge();
     }
 
+    public static @NotNull TinyFile enrichSeargeWithMCP(@NotNull TinyFile searge, @NotNull Path mcpArchiveRoot) throws IOException {
+        if (!searge.getHeader().getNamespaces().equals(Arrays.asList("left", "right")))
+            throw new IllegalArgumentException("Searge namespaces need to be left and right");
+        var fields = mcpArchiveRoot.resolve("fields.csv");
+        var methods = mcpArchiveRoot.resolve("methods.csv");
+        var params = mcpArchiveRoot.resolve("params.csv");
+        if (!Files.exists(fields))
+            throw new IllegalArgumentException("Missing fields.csv");
+        if (!Files.exists(methods))
+            throw new IllegalArgumentException("Missing methods.csv");
+        if (!Files.exists(params))
+            throw new IllegalArgumentException("Missing params.csv");
+        return new EnrichSeargeWithMCP(searge, fields, methods, params).mergeTinyFile();
+    }
+
     public static @NotNull TinyFile createOverlayTinyFile(
             @NotNull TinyFile base, @NotNull TinyFile overlay,
             @NotNull List<@NotNull String> retainedNamespaces,
diff --git a/src/test/java/moe/nea/zwirn/ZwirnTest.java b/src/test/java/moe/nea/zwirn/ZwirnTest.java
index b3ed4f5..c0db320 100644
--- a/src/test/java/moe/nea/zwirn/ZwirnTest.java
+++ b/src/test/java/moe/nea/zwirn/ZwirnTest.java
@@ -4,6 +4,7 @@ import net.fabricmc.stitch.commands.tinyv2.*;
 import org.junit.jupiter.api.Test;
 
 import java.io.IOException;
+import java.nio.file.FileSystems;
 import java.nio.file.Path;
 import java.util.Arrays;
 import java.util.HashMap;
@@ -20,6 +21,14 @@ class ZwirnTest {
         TinyV2Writer.write(unmerged, Path.of("unmerged.tiny"));
     }
 
+    @Test
+    void mergeMCP() throws IOException {
+        try (var fs = FileSystems.newFileSystem(Path.of("mcp.zip"))) {
+            var merged = Zwirn.enrichSeargeWithMCP(TinyV2Reader.read(Path.of("searge.tiny")), fs.getPath("/"));
+            TinyV2Writer.write(merged, Path.of("mcp.tiny"));
+        }
+    }
+
 
     TinyFile getBaseFile() {
         return new TinyFile(