Renaming and fixing method descriptors sometimes

8 files changed, 286 insertions, 7 deletions

diff --git a/build.gradle.kts b/build.gradle.kts
index 50f3fe5..1761713 100644
--- a/build.gradle.kts
+++ b/build.gradle.kts
@@ -24,6 +24,7 @@ dependencies {
 
 java {
     this.toolchain.languageVersion.set(JavaLanguageVersion.of(17))
+    withSourcesJar()
 }
 
 tasks.test {
diff --git a/src/main/java/moe/nea/zwirn/Descriptor.java b/src/main/java/moe/nea/zwirn/Descriptor.java
new file mode 100644
index 0000000..a121264
--- /dev/null
+++ b/src/main/java/moe/nea/zwirn/Descriptor.java
@@ -0,0 +1,154 @@
+package moe.nea.zwirn;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public interface Descriptor {
+    Descriptor mapClassName(SimpleRemapper remapper);
+
+    String toJvmDescriptor();
+
+    interface Type extends Descriptor {
+        Type mapClassName(SimpleRemapper remapper);
+
+        static Type readType(GoodStringReader reader) {
+            if (reader.peekChar() == 'V') {
+                reader.nextChar();
+                return VoidType.VOID;
+            }
+            return Field.readField(reader);
+        }
+    }
+
+    interface Field extends Type {
+        Field mapClassName(SimpleRemapper remapper);
+
+        static Field readField(GoodStringReader reader) {
+            switch (reader.nextChar()) {
+                case 'L':
+                    return new Class(reader.readUntil(';').replace('/', '.'));
+                case 'Z':
+                    return DefaultField.BOOLEAN;
+                case 'B':
+                    return DefaultField.BYTE;
+                case 'I':
+                    return DefaultField.INT;
+                case 'D':
+                    return DefaultField.DOUBLE;
+                case 'J':
+                    return DefaultField.LONG;
+                case 'S':
+                    return DefaultField.SHORT;
+                case 'F':
+                    return DefaultField.FLOAT;
+                case 'C':
+                    return DefaultField.CHAR;
+                case '[':
+                    return new Array(readField(reader));
+            }
+            throw new IllegalStateException("Unknown type");
+        }
+    }
+
+    record Method(List<Field> argumentTypes, Type returnType) implements Descriptor {
+
+        @Override
+        public Descriptor mapClassName(SimpleRemapper remapper) {
+            List<Field> newArgumentTypes = new ArrayList<>(argumentTypes.size());
+            for (Field argumentType : argumentTypes) {
+                newArgumentTypes.add(argumentType.mapClassName(remapper));
+            }
+            return new Method(newArgumentTypes, returnType.mapClassName(remapper));
+        }
+
+        @Override
+        public String toJvmDescriptor() {
+            StringBuilder sb = new StringBuilder();
+            sb.append("(");
+            for (Field argumentType : argumentTypes) {
+                sb.append(argumentType.toJvmDescriptor());
+            }
+            sb.append(")");
+            sb.append(returnType.toJvmDescriptor());
+            return sb.toString();
+        }
+
+        public static Method readMethod(GoodStringReader reader) {
+            if (reader.nextChar() != '(')
+                throw new IllegalStateException("Expected (");
+            List<Field> argument = new ArrayList<>();
+            while (reader.peekChar() != ')') {
+                argument.add(Field.readField(reader));
+            }
+            reader.nextChar(); // Consume )
+            return new Method(argument, Type.readType(reader));
+        }
+    }
+
+    record Array(Field field) implements Field {
+
+        @Override
+        public String toJvmDescriptor() {
+            return "[" + field.toJvmDescriptor();
+        }
+
+        @Override
+        public Array mapClassName(SimpleRemapper remapper) {
+            return new Array(field.mapClassName(remapper));
+        }
+    }
+
+    record Class(String dottedName) implements Field {
+        @Override
+        public String toJvmDescriptor() {
+            return "L" + dottedName.replace('.', '/') + ";";
+        }
+
+        @Override
+        public Field mapClassName(SimpleRemapper remapper) {
+            return new Class(remapper.remapClass(dottedName));
+        }
+    }
+
+
+    enum VoidType implements Type {
+        VOID;
+
+        @Override
+        public String toJvmDescriptor() {
+            return "V";
+        }
+
+        @Override
+        public VoidType mapClassName(SimpleRemapper remapper) {
+            return this;
+        }
+    }
+
+    enum DefaultField implements Field {
+        BOOLEAN("Z"),
+        BYTE("B"),
+        INT("I"),
+        DOUBLE("D"),
+        LONG("J"),
+        SHORT("S"),
+        FLOAT("F"),
+        CHAR("C");
+        private final String jvmDescriptor;
+
+        DefaultField(String jvmDescriptor) {
+            this.jvmDescriptor = jvmDescriptor;
+        }
+
+        @Override
+        public String toJvmDescriptor() {
+            return jvmDescriptor;
+        }
+
+        @Override
+        public DefaultField mapClassName(SimpleRemapper remapper) {
+            return this;
+        }
+    }
+
+}
diff --git a/src/main/java/moe/nea/zwirn/GoodStringReader.java b/src/main/java/moe/nea/zwirn/GoodStringReader.java
new file mode 100644
index 0000000..a396120
--- /dev/null
+++ b/src/main/java/moe/nea/zwirn/GoodStringReader.java
@@ -0,0 +1,46 @@
+package moe.nea.zwirn;
+
+import java.util.Stack;
+
+public class GoodStringReader {
+    final String source;
+    Stack<Integer> stack = new Stack<>();
+    int index = 0;
+
+    public GoodStringReader(String source) {
+        this.source = source;
+    }
+
+    public void push() {
+        stack.push(index);
+    }
+
+    public void reset() {
+        index = stack.pop();
+    }
+
+    public void discard() {
+        stack.pop();
+    }
+
+    public char nextChar() {
+        return source.charAt(index++);
+    }
+
+    public String readUntil(char... cs) {
+        int minI = -1;
+        for (char c : cs) {
+            int i = source.indexOf(c, index);
+            if (i < 0) continue;
+            minI = minI < 0 ? i : Math.min(minI, i);
+        }
+        if (minI < 0) return null;
+        int startIndex = index;
+        index = minI + 1;
+        return source.substring(startIndex, index - 1);
+    }
+
+    public char peekChar() {
+        return source.charAt(index);
+    }
+}
diff --git a/src/main/java/moe/nea/zwirn/RenameTask.java b/src/main/java/moe/nea/zwirn/RenameTask.java
index 89c15c4..a9a2153 100644
--- a/src/main/java/moe/nea/zwirn/RenameTask.java
+++ b/src/main/java/moe/nea/zwirn/RenameTask.java
@@ -4,6 +4,8 @@ import net.fabricmc.stitch.commands.tinyv2.*;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
 import java.util.stream.Collectors;
 
@@ -11,6 +13,7 @@ public class RenameTask {
     private final TinyFile tinyFile;
     private final List<Zwirn.RenameCommand> newNamespaceOrder;
     private final int[] namespaceRemapOrder;
+    private final SimpleRemapper remapper;
 
     public RenameTask(@NotNull TinyFile tinyFile, @NotNull List<Zwirn.RenameCommand> newNamespaceOrder) {
         this.tinyFile = tinyFile;
@@ -18,8 +21,26 @@ public class RenameTask {
         namespaceRemapOrder = newNamespaceOrder.stream().mapToInt(
                 it -> tinyFile.getHeader().getNamespaces().indexOf(it.oldNamespaceName())
         ).toArray();
+        this.remapper = new SimpleRemapper(
+                classOnlyRename(), "__old", "__new"
+        );
     }
 
+    public TinyFile classOnlyRename() {
+        List<String> namespaces = new ArrayList<>();
+        namespaces.add("__old");
+        namespaces.add("__new");
+        return new TinyFile(
+                new TinyHeader(namespaces, 2, 0, new HashMap<>()),
+                tinyFile.getClassEntries().stream().map(
+                        it -> new TinyClass(
+                                Arrays.asList(
+                                        it.getClassNames().get(0),
+                                        it.getClassNames().get(namespaceRemapOrder[0]))
+                        )
+                ).collect(Collectors.toList())
+        );
+    }
 
     private List<String> rename(List<String> strings) {
         List<String> newNames = new ArrayList<>(namespaceRemapOrder.length);
@@ -48,7 +69,7 @@ public class RenameTask {
     private TinyField renameField(TinyField tinyField) {
         var names = rename(tinyField.getFieldNames());
         return new TinyField(
-                names.get(0),
+                remapper.remapFieldDescriptor(tinyField.getFieldDescriptorInFirstNamespace()),
                 names,
                 tinyField.getComments()
         );
@@ -57,7 +78,7 @@ public class RenameTask {
     private TinyMethod renameMethod(TinyMethod tinyMethod) {
         var names = rename(tinyMethod.getMethodNames());
         return new TinyMethod(
-                names.get(0),
+                remapper.remapMethodDescriptor(tinyMethod.getMethodDescriptorInFirstNamespace()),
                 names,
                 tinyMethod.getParameters().stream().map(this::renameMethodParameter).collect(Collectors.toList()),
                 tinyMethod.getLocalVariables().stream().map(this::renameVariable).collect(Collectors.toList()),
diff --git a/src/main/java/moe/nea/zwirn/SimpleRemapper.java b/src/main/java/moe/nea/zwirn/SimpleRemapper.java
new file mode 100644
index 0000000..55463f1
--- /dev/null
+++ b/src/main/java/moe/nea/zwirn/SimpleRemapper.java
@@ -0,0 +1,33 @@
+package moe.nea.zwirn;
+
+import net.fabricmc.stitch.commands.tinyv2.TinyFile;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class SimpleRemapper {
+    private final Map<String, String> lookup = new HashMap<>();
+
+    public SimpleRemapper(TinyFile file, String sourceNs, String targetNs) {
+        int sourceIndex = file.getHeader().getNamespaces().indexOf(sourceNs);
+        int targetIndex = file.getHeader().getNamespaces().indexOf(targetNs);
+        for (var classEntry : file.getClassEntries()) {
+            lookup.put(classEntry.getClassNames().get(sourceIndex).replace('/', '.'),
+                    classEntry.getClassNames().get(targetIndex).replace('/', '.'));
+        }
+    }
+
+    public String remapClass(String dottedName) {
+        return lookup.getOrDefault(dottedName, dottedName);
+    }
+
+    public String remapFieldDescriptor(String field) {
+        if ("null".equals(field)) return "null";
+        return Descriptor.Field.readField(new GoodStringReader(field)).mapClassName(this).toJvmDescriptor();
+    }
+
+    public String remapMethodDescriptor(String method) {
+        if ("null".equals(method)) return "null";
+        return Descriptor.Method.readMethod(new GoodStringReader(method)).mapClassName(this).toJvmDescriptor();
+    }
+}
diff --git a/src/main/java/moe/nea/zwirn/TinyDiffer.java b/src/main/java/moe/nea/zwirn/TinyDiffer.java
index f5fc471..2c274f7 100644
--- a/src/main/java/moe/nea/zwirn/TinyDiffer.java
+++ b/src/main/java/moe/nea/zwirn/TinyDiffer.java
@@ -23,7 +23,7 @@ public class TinyDiffer {
         this.retainedNamespaces = retainedNamespaces;
         this.baseSharedIndex = base.getHeader().getNamespaces().indexOf(retainedNamespaces.get(0));// TODO: shared namespace argument
         this.overlaySharedIndex = overlay.getHeader().getNamespaces().indexOf(retainedNamespaces.get(0));
-
+        // TODO: remap descriptors
         retainedToNormalLookup = retainedNamespaces.stream()
                 .mapToInt(it -> base.getHeader().getNamespaces().indexOf(it))
                 .toArray();
diff --git a/src/main/java/moe/nea/zwirn/TinyMerger.java b/src/main/java/moe/nea/zwirn/TinyMerger.java
index 97b245b..6a1c5ab 100644
--- a/src/main/java/moe/nea/zwirn/TinyMerger.java
+++ b/src/main/java/moe/nea/zwirn/TinyMerger.java
@@ -19,6 +19,7 @@ class TinyMerger {
     private final @NotNull Map<@NotNull String, @NotNull TinyClass> baseLUT;
     private final Integer[] baseToOverlayIndexMap;
     private final @NotNull List<@NotNull TinyClass> entries = new ArrayList<>();
+    private final SimpleRemapper remapper;
 
     public TinyMerger(TinyFile base, TinyFile overlay, String sharedNamespace) {
         this.base = base;
@@ -34,6 +35,29 @@ class TinyMerger {
                 .stream().map(it -> overlay.getHeader().getNamespaces().indexOf(it))
                 .map(it -> it < 0 ? null : it)
                 .toArray(Integer[]::new);
+        remapper = new SimpleRemapper(classOnlyMerge(), "__old", base.getHeader().getNamespaces().get(0));
+    }
+
+    public TinyFile classOnlyMerge() {
+        List<String> namespaces = new ArrayList<>();
+        namespaces.add("__old");
+        namespaces.addAll(base.getHeader().getNamespaces());
+        return new TinyFile(
+                new TinyHeader(namespaces, 2, 0, new HashMap<>()),
+                mergeChildren(
+                        base,
+                        overlay,
+                        it -> it.getClassNames().get(baseSharedIndex),
+                        it -> it.getClassNames().get(overlaySharedIndex),
+                        TinyFile::getClassEntries,
+                        (tinyClass, tinyClass2) -> {
+                            List<String> mergedNames = new ArrayList<>();
+                            mergedNames.add(tinyClass.getClassNames().get(0));
+                            mergedNames.addAll(mergeNames(tinyClass, tinyClass2, TinyClass::getClassNames));
+                            return new TinyClass(mergedNames);
+                        }
+                )
+        );
     }
 
     public @NotNull TinyFile merge() {
@@ -113,7 +137,7 @@ class TinyMerger {
     private TinyField mergeField(@Nullable TinyField baseField, @Nullable TinyField overlayField) {
         var mergedNames = mergeNames(baseField, overlayField, TinyField::getFieldNames);
         return new TinyField(
-                mergedNames.get(0),
+                remapper.remapFieldDescriptor(baseField != null ? baseField.getFieldDescriptorInFirstNamespace() : overlayField.getFieldDescriptorInFirstNamespace()),
                 mergedNames,
                 mergeComments(baseField, overlayField, TinyField::getComments)
         );
@@ -130,7 +154,7 @@ class TinyMerger {
     private @NotNull TinyMethod mergeMethod(@Nullable TinyMethod baseMethod, @Nullable TinyMethod overlayMethod) {
         var mergedNames = mergeNames(baseMethod, overlayMethod, TinyMethod::getMethodNames);
         return new TinyMethod(
-                mergedNames.get(0),
+                remapper.remapMethodDescriptor(baseMethod != null ? baseMethod.getMethodDescriptorInFirstNamespace() : overlayMethod.getMethodDescriptorInFirstNamespace()),
                 mergedNames,
                 mergeChildren(baseMethod, overlayMethod,
                         TinyMethodParameter::getLvIndex,
diff --git a/src/test/java/moe/nea/zwirn/ZwirnTest.java b/src/test/java/moe/nea/zwirn/ZwirnTest.java
index c0db320..549e2d7 100644
--- a/src/test/java/moe/nea/zwirn/ZwirnTest.java
+++ b/src/test/java/moe/nea/zwirn/ZwirnTest.java
@@ -41,7 +41,7 @@ class ZwirnTest {
                                 Arrays.asList("A", "ClassA", "SomeClass"),
                                 Arrays.asList(
                                         new TinyMethod(
-                                                "a",
+                                                "(LA;)V",
                                                 Arrays.asList("a", "method_a", "doSomething"),
                                                 Arrays.asList(
                                                         new TinyMethodParameter(
@@ -54,7 +54,7 @@ class ZwirnTest {
                                                 Arrays.asList("method comment")
                                         )
                                 ),
-                                Arrays.asList(new TinyField("a",
+                                Arrays.asList(new TinyField("LA;",
                                         Arrays.asList("a", "field_a", "myField"), Arrays.asList("Field comment"))),
                                 Arrays.asList("some comment")
                         ),