Add constructor statusHEADmaster

3 files changed, 127 insertions, 4 deletions

diff --git a/src/main/java/moe/nea/zwirn/ConstructorEnrichmentTask.java b/src/main/java/moe/nea/zwirn/ConstructorEnrichmentTask.java
new file mode 100644
index 0000000..687aeb9
--- /dev/null
+++ b/src/main/java/moe/nea/zwirn/ConstructorEnrichmentTask.java
@@ -0,0 +1,100 @@
+package moe.nea.zwirn;
+
+import net.fabricmc.stitch.commands.tinyv2.TinyClass;
+import net.fabricmc.stitch.commands.tinyv2.TinyFile;
+import net.fabricmc.stitch.commands.tinyv2.TinyMethod;
+import net.fabricmc.stitch.commands.tinyv2.TinyMethodParameter;
+import org.jetbrains.annotations.NotNull;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.*;
+import java.util.stream.Collectors;
+
+public class ConstructorEnrichmentTask {
+    private final TinyFile searge;
+    private final List<McpConstructor> mcpConstructors;
+    private final Map<String, List<McpConstructor>> indexedConstructors;
+    private final SimpleRemapper remapper;
+
+    public ConstructorEnrichmentTask(@NotNull TinyFile searge, Path joinedExc) {
+        this.searge = searge;
+        this.mcpConstructors = readExc(joinedExc);
+        this.indexedConstructors = this.mcpConstructors.stream().collect(Collectors.groupingBy(McpConstructor::owningClass, Collectors.toList()));
+        this.remapper = new SimpleRemapper(searge, "right", "left");
+    }
+
+    public TinyFile enrich() {
+        return new TinyFile(
+                searge.getHeader(),
+                searge.getClassEntries()
+                        .stream().map(this::enrichClass).collect(Collectors.toList())
+        );
+    }
+
+    private TinyClass enrichClass(TinyClass tinyClass) {
+        List<TinyMethod> methods = new ArrayList<>(tinyClass.getMethods());
+        List<McpConstructor> constructors = indexedConstructors.getOrDefault(tinyClass.getClassNames().get(1).replace('/', '.'), Collections.emptyList());
+        for (McpConstructor constructor : constructors) {
+            methods.add(new TinyMethod(
+                    remapper.remapMethodDescriptor(constructor.methodDescriptorInSrg()),
+                    Arrays.asList("<init>", "<init>"),
+                    constructor.getTinyParameters(),
+                    Arrays.asList(),
+                    Arrays.asList()
+            ));
+        }
+        return new TinyClass(
+                tinyClass.getClassNames(),
+                methods,
+                tinyClass.getFields(),
+                tinyClass.getComments()
+        );
+    }
+
+    record McpConstructor(
+            String owningClass,
+            String methodDescriptorInSrg,
+            List<String> parameterNames
+    ) {
+        public Collection<TinyMethodParameter> getTinyParameters() {
+            List<TinyMethodParameter> params = new ArrayList<>();
+            int i = 0;
+            for (String parameterName : parameterNames) {
+                int lvIndex = ++i;
+                params.add(new TinyMethodParameter(lvIndex, Arrays.asList("p" + lvIndex, parameterName), Arrays.asList()));
+            }
+            return params;
+        }
+    }
+
+    private static List<McpConstructor> readExc(Path joinedExc) {
+        var prop = new Properties();
+        try (var reader = Files.newBufferedReader(joinedExc)) {
+            prop.load(reader);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+        List<McpConstructor> cons = new ArrayList<>();
+        for (Map.Entry<Object, Object> entry : prop.entrySet()) {
+            var target = (String) entry.getKey();
+            var signatureInfo = (String) entry.getValue();
+            if (!signatureInfo.contains("|")) continue;
+            var targetParts = target.split("\\.");
+            assert targetParts.length == 2;
+            var className = targetParts[0].replace('/', '.');
+            var constructorDescriptor = "(" + targetParts[1].split("\\(")[1];
+            if (!targetParts[1].startsWith("<init>"))
+                continue;
+            int pipeIndex = signatureInfo.indexOf('|');
+            var exceptions = signatureInfo.substring(0, pipeIndex);
+            var parameterNameString = signatureInfo.substring(pipeIndex + 1);
+            var parameterNames = parameterNameString.isBlank()
+                    ? Arrays.<String>asList()
+                    : Arrays.asList(parameterNameString.split(","));
+            cons.add(new McpConstructor(className, constructorDescriptor, parameterNames));
+        }
+        return cons;
+    }
+}
diff --git a/src/main/java/moe/nea/zwirn/EnrichSeargeWithMCP.java b/src/main/java/moe/nea/zwirn/EnrichSeargeWithMCP.java
index d91e424..8631e6a 100644
--- a/src/main/java/moe/nea/zwirn/EnrichSeargeWithMCP.java
+++ b/src/main/java/moe/nea/zwirn/EnrichSeargeWithMCP.java
@@ -6,7 +6,10 @@ import org.jetbrains.annotations.NotNull;
 
 import java.io.IOException;
 import java.nio.file.Path;
-import java.util.*;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 import java.util.function.Function;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
@@ -19,6 +22,7 @@ public class EnrichSeargeWithMCP {
     private final Map<String, MCPField> fieldMap;
     private final Map<Integer, List<MCPParam>> paramMap;
     private final Map<String, MCPMethod> methodMap;
+    private final Map<String, MCPParam> constructorParamMap;
 
     // TODO: parse joined.exc for constructor indexes parameters
     public EnrichSeargeWithMCP(@NotNull TinyFile searge, Path fields, Path methods, Path params) throws IOException {
@@ -29,6 +33,7 @@ public class EnrichSeargeWithMCP {
         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()));
+        this.constructorParamMap = this.params.stream().collect(Collectors.toMap(MCPParam::searge, Function.identity()));
     }
 
     record MCPParam(
@@ -135,21 +140,30 @@ public class EnrichSeargeWithMCP {
     private TinyMethod mergeMethod(TinyMethod tinyMethod) {
         var srg = tinyMethod.getMethodNames().get(1);
         var mcpMethod = methodMap.get(srg);
-        List<TinyMethodParameter> params = new ArrayList<>();
+        Map<Integer, TinyMethodParameter> params = new HashMap<>();
         if (mcpMethod != null) {
             var mcpParams = paramMap.get(mcpMethod.methodId());
             if (mcpParams != null) for (var param : mcpParams) {
-                params.add(new TinyMethodParameter(
+                params.put(param.lvIndexHeuristic(), new TinyMethodParameter(
                         param.lvIndexHeuristic(),
                         Arrays.asList("p" + param.lvIndexHeuristic(), param.searge(), param.name()),
                         Arrays.asList()
                 ));
             }
         }
+        for (TinyMethodParameter parameter : tinyMethod.getParameters()) {
+            MCPParam mcpParam = constructorParamMap.get(parameter.getParameterNames().get(1));
+            if (mcpParam != null)
+                params.put(parameter.getLvIndex(), new TinyMethodParameter(
+                        parameter.getLvIndex(),
+                        Arrays.asList(parameter.getParameterNames().get(0), parameter.getParameterNames().get(1), mcpParam.name()),
+                        parameter.getComments()
+                ));
+        }
         return new TinyMethod(
                 tinyMethod.getMethodDescriptorInFirstNamespace(),
                 Arrays.asList(tinyMethod.getMethodNames().get(0), srg, mcpMethod == null ? srg : mcpMethod.name()),
-                params,
+                params.values(),
                 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 c5e6661..47680b2 100644
--- a/src/main/java/moe/nea/zwirn/Zwirn.java
+++ b/src/main/java/moe/nea/zwirn/Zwirn.java
@@ -40,6 +40,15 @@ public class Zwirn {
         return new RenameTask(tinyFile, newNamespaceOrder).rename();
     }
 
+    public static @NotNull TinyFile enrichSeargeWithConstructors(@NotNull TinyFile searge, @NotNull Path mcpSrgRoot) {
+        var joinedExc = mcpSrgRoot.resolve("joined.exc");
+        if (!Files.exists(joinedExc))
+            throw new IllegalArgumentException("joined.exc not found in MCP srg jar");
+        if (!searge.getHeader().getNamespaces().equals(Arrays.asList("left", "right")))
+            throw new IllegalArgumentException("Searge namespaces need to be left and right");
+        return new ConstructorEnrichmentTask(searge, joinedExc).enrich();
+    }
+
     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");