Fix occassional diffing problems

4 files changed, 113 insertions, 11 deletions

diff --git a/src/main/java/moe/nea/zwirn/FieldSignatureFixer.java b/src/main/java/moe/nea/zwirn/FieldSignatureFixer.java
new file mode 100644
index 0000000..7554a42
--- /dev/null
+++ b/src/main/java/moe/nea/zwirn/FieldSignatureFixer.java
@@ -0,0 +1,70 @@
+package moe.nea.zwirn;
+
+import net.fabricmc.stitch.commands.tinyv2.TinyClass;
+import net.fabricmc.stitch.commands.tinyv2.TinyField;
+import net.fabricmc.stitch.commands.tinyv2.TinyFile;
+import org.jetbrains.annotations.NotNull;
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.FieldVisitor;
+import org.objectweb.asm.Opcodes;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public class FieldSignatureFixer {
+    private final TinyFile tinyFile;
+    private final Path classRoot;
+
+    public FieldSignatureFixer(@NotNull TinyFile tinyFile, @NotNull Path classRoot) {
+        this.tinyFile = tinyFile;
+        this.classRoot = classRoot;
+    }
+
+    public TinyFile fix() {
+        return new TinyFile(tinyFile.getHeader(),
+                tinyFile.getClassEntries().stream().map(this::fixClass).collect(Collectors.toList())
+        );
+    }
+
+    private TinyClass fixClass(TinyClass tinyClass) {
+        var path = tinyClass.getClassNames().get(0).replace('.', '/') + ".class";
+        var classPath = classRoot.resolve(path);
+        Map<String, String> fieldNameToDescriptorMap = new HashMap<>();
+        if (Files.exists(classPath))
+            try (var is = Files.newInputStream(classPath)) {
+                var reader = new ClassReader(is);
+                var visitor = new FieldSignatureVisitor();
+                reader.accept(visitor, 0);
+                fieldNameToDescriptorMap.putAll(visitor.fieldNameToDescriptorMap);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        return new TinyClass(
+                tinyClass.getClassNames(),
+                tinyClass.getMethods(),
+                tinyClass.getFields().stream()
+                        .map(it -> new TinyField(fieldNameToDescriptorMap.get(it.getFieldNames().get(0)), it.getFieldNames(), it.getComments()))
+                        .collect(Collectors.toList()),
+                tinyClass.getComments()
+        );
+    }
+
+    static class FieldSignatureVisitor extends ClassVisitor {
+        protected FieldSignatureVisitor() {
+            super(Opcodes.ASM9);
+        }
+
+        public Map<String, String> fieldNameToDescriptorMap = new HashMap<>();
+
+        @Override
+        public FieldVisitor visitField(int access, String name, String descriptor, String signature, Object value) {
+            fieldNameToDescriptorMap.put(name, descriptor);
+            return super.visitField(access, name, descriptor, signature, value);
+        }
+    }
+}
diff --git a/src/main/java/moe/nea/zwirn/TinyDiffer.java b/src/main/java/moe/nea/zwirn/TinyDiffer.java
index 2c274f7..5a3c20d 100644
--- a/src/main/java/moe/nea/zwirn/TinyDiffer.java
+++ b/src/main/java/moe/nea/zwirn/TinyDiffer.java
@@ -9,6 +9,8 @@ import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
+import static moe.nea.zwirn.ZwirnUtil.TODO;
+
 public class TinyDiffer {
     private final TinyFile base;
     private final TinyFile overlay;
@@ -27,6 +29,8 @@ public class TinyDiffer {
         retainedToNormalLookup = retainedNamespaces.stream()
                 .mapToInt(it -> base.getHeader().getNamespaces().indexOf(it))
                 .toArray();
+        if (baseSharedIndex != 0 || overlaySharedIndex != 0)
+            TODO("add support for remapping descriptors");
     }
 
 
@@ -158,7 +162,10 @@ public class TinyDiffer {
             @Nullable TinyClass overlayClass) {
         var names = diffNamesWithMappings(baseClass, overlayClass);
         var fields = diffChildrenByMapping(baseClass, overlayClass, TinyClass::getFields, this::diffField);
-        var methods = diffChildrenByMapping(baseClass, overlayClass, TinyClass::getMethods, this::diffMethod);
+        var methods = diffChildren(baseClass, overlayClass,
+                it -> it.getMethodNames().get(baseSharedIndex) + it.getMethodDescriptorInFirstNamespace(),
+                it -> it.getMethodNames().get(overlaySharedIndex) + it.getMethodDescriptorInFirstNamespace(),
+                TinyClass::getMethods, this::diffMethod);
         var comments = diffComments(baseClass, overlayClass, TinyClass::getComments);
         return Retained.keep(
                 new TinyClass(
@@ -181,6 +188,7 @@ public class TinyDiffer {
             return Retained.empty(new ArrayList<>(commentExtractor.apply(baseObject)));
         var comments = new ArrayList<>(commentExtractor.apply(overlayObject));
         comments.removeAll(commentExtractor.apply(baseObject));
+        comments.removeIf(String::isEmpty);
         if (comments.isEmpty())
             return Retained.empty(comments);
         return Retained.strong(comments);
@@ -193,7 +201,7 @@ public class TinyDiffer {
         var comments = diffComments(baseMethod, overlayMethod, TinyMethod::getComments);
         return Retained.keep(
                 new TinyMethod(
-                        names.diff.get(0),
+                        baseMethod.getMethodDescriptorInFirstNamespace(),
                         names.diff,
                         params.diff,
                         variables.diff,
@@ -253,7 +261,7 @@ public class TinyDiffer {
         var comments = diffComments(baseField, overlayField, TinyField::getComments);
         return Retained.keep(
                 new TinyField(
-                        names.diff.get(0),
+                        baseField.getFieldDescriptorInFirstNamespace(),
                         names.diff,
                         comments.diff
                 ),
diff --git a/src/main/java/moe/nea/zwirn/TinyMerger.java b/src/main/java/moe/nea/zwirn/TinyMerger.java
index 6a1c5ab..2b3d4cc 100644
--- a/src/main/java/moe/nea/zwirn/TinyMerger.java
+++ b/src/main/java/moe/nea/zwirn/TinyMerger.java
@@ -19,7 +19,8 @@ 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;
+    private final SimpleRemapper remapperBase;
+    private final SimpleRemapper remapperOverlay;
 
     public TinyMerger(TinyFile base, TinyFile overlay, String sharedNamespace) {
         this.base = base;
@@ -35,12 +36,15 @@ 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));
+        var classMerged = classOnlyMerge();
+        remapperBase = new SimpleRemapper(classMerged, "__old_base", base.getHeader().getNamespaces().get(0));
+        remapperOverlay = new SimpleRemapper(classMerged, "__old_shared", base.getHeader().getNamespaces().get(0));
     }
 
     public TinyFile classOnlyMerge() {
         List<String> namespaces = new ArrayList<>();
-        namespaces.add("__old");
+        namespaces.add("__old_base");
+        namespaces.add("__old_shared");
         namespaces.addAll(base.getHeader().getNamespaces());
         return new TinyFile(
                 new TinyHeader(namespaces, 2, 0, new HashMap<>()),
@@ -52,7 +56,16 @@ class TinyMerger {
                         TinyFile::getClassEntries,
                         (tinyClass, tinyClass2) -> {
                             List<String> mergedNames = new ArrayList<>();
-                            mergedNames.add(tinyClass.getClassNames().get(0));
+                            if (tinyClass != null) {
+                                mergedNames.add(tinyClass.getClassNames().get(baseSharedIndex));
+                            } else {
+                                mergedNames.add(tinyClass2.getClassNames().get(overlaySharedIndex));
+                            }
+                            if (tinyClass2 != null) {
+                                mergedNames.add(tinyClass2.getClassNames().get(overlaySharedIndex));
+                            } else {
+                                mergedNames.add(tinyClass.getClassNames().get(baseSharedIndex));
+                            }
                             mergedNames.addAll(mergeNames(tinyClass, tinyClass2, TinyClass::getClassNames));
                             return new TinyClass(mergedNames);
                         }
@@ -137,7 +150,8 @@ class TinyMerger {
     private TinyField mergeField(@Nullable TinyField baseField, @Nullable TinyField overlayField) {
         var mergedNames = mergeNames(baseField, overlayField, TinyField::getFieldNames);
         return new TinyField(
-                remapper.remapFieldDescriptor(baseField != null ? baseField.getFieldDescriptorInFirstNamespace() : overlayField.getFieldDescriptorInFirstNamespace()),
+                baseField != null ? remapperBase.remapFieldDescriptor(baseField.getFieldDescriptorInFirstNamespace())
+                        : remapperOverlay.remapFieldDescriptor(overlayField.getFieldDescriptorInFirstNamespace()),
                 mergedNames,
                 mergeComments(baseField, overlayField, TinyField::getComments)
         );
@@ -145,8 +159,8 @@ class TinyMerger {
 
     private Collection<TinyMethod> mergeMethods(@Nullable TinyClass baseClass, @Nullable TinyClass overlayClass) {
         return mergeChildren(baseClass, overlayClass,
-                it -> it.getMethodNames().get(baseSharedIndex),
-                it -> it.getMethodNames().get(overlaySharedIndex),
+                it -> it.getMethodNames().get(baseSharedIndex) + remapperBase.remapMethodDescriptor(it.getMethodDescriptorInFirstNamespace()),
+                it -> it.getMethodNames().get(overlaySharedIndex) + remapperOverlay.remapMethodDescriptor(it.getMethodDescriptorInFirstNamespace()),
                 TinyClass::getMethods,
                 this::mergeMethod);
     }
@@ -154,7 +168,8 @@ class TinyMerger {
     private @NotNull TinyMethod mergeMethod(@Nullable TinyMethod baseMethod, @Nullable TinyMethod overlayMethod) {
         var mergedNames = mergeNames(baseMethod, overlayMethod, TinyMethod::getMethodNames);
         return new TinyMethod(
-                remapper.remapMethodDescriptor(baseMethod != null ? baseMethod.getMethodDescriptorInFirstNamespace() : overlayMethod.getMethodDescriptorInFirstNamespace()),
+                (baseMethod != null ? remapperBase.remapMethodDescriptor(baseMethod.getMethodDescriptorInFirstNamespace())
+                        : remapperOverlay.remapMethodDescriptor(overlayMethod.getMethodDescriptorInFirstNamespace())),
                 mergedNames,
                 mergeChildren(baseMethod, overlayMethod,
                         TinyMethodParameter::getLvIndex,
diff --git a/src/main/java/moe/nea/zwirn/Zwirn.java b/src/main/java/moe/nea/zwirn/Zwirn.java
index b3a4956..86cb948 100644
--- a/src/main/java/moe/nea/zwirn/Zwirn.java
+++ b/src/main/java/moe/nea/zwirn/Zwirn.java
@@ -1,12 +1,16 @@
 package moe.nea.zwirn;
 
+import net.fabricmc.stitch.commands.tinyv2.TinyClass;
 import net.fabricmc.stitch.commands.tinyv2.TinyFile;
+import net.fabricmc.stitch.commands.tinyv2.TinyHeader;
 import org.jetbrains.annotations.NotNull;
 
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
 
 public class Zwirn {
@@ -67,4 +71,9 @@ public class Zwirn {
             throw new IllegalArgumentException("Shared namespace must be retained");
         return new TinyDiffer(base, overlay, retainedNamespaces).createDiff();
     }
+
+    public static @NotNull TinyFile fixFieldDescriptorsFromJar(@NotNull TinyFile tinyFile, @NotNull Path classRoot) {
+        return new FieldSignatureFixer(tinyFile, classRoot).fix();
+    }
+
 }