Add diffing

1 files changed, 206 insertions, 0 deletions

diff --git a/src/main/java/moe/nea/zwirn/TinyBaseMerger.java b/src/main/java/moe/nea/zwirn/TinyBaseMerger.java
new file mode 100644
index 0000000..96ddf2a
--- /dev/null
+++ b/src/main/java/moe/nea/zwirn/TinyBaseMerger.java
@@ -0,0 +1,206 @@
+package moe.nea.zwirn;
+
+import net.fabricmc.stitch.commands.tinyv2.*;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.*;
+import java.util.function.BiFunction;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+public abstract class TinyBaseMerger {
+    protected final @NotNull TinyFile base;
+    protected final @NotNull TinyFile overlay;
+    protected final @NotNull String sharedNamespace;
+    protected final int baseSharedIndex;
+    protected final int overlaySharedIndex;
+    protected final @NotNull Map<@NotNull String, @NotNull TinyClass> overlayLUT;
+    protected final @NotNull Map<@NotNull String, @NotNull TinyClass> baseLUT;
+    protected final Integer[] baseToOverlayIndexMap;
+
+    protected TinyBaseMerger(TinyFile base, TinyFile overlay, String sharedNamespace) {
+        this.base = base;
+        this.overlay = overlay;
+        this.sharedNamespace = sharedNamespace;
+        this.baseSharedIndex = base.getHeader().getNamespaces().indexOf(sharedNamespace);
+        this.overlaySharedIndex = overlay.getHeader().getNamespaces().indexOf(sharedNamespace);
+        this.overlayLUT = overlay.getClassEntries().stream()
+                .collect(Collectors.toMap(it -> it.getClassNames().get(this.overlaySharedIndex), Function.identity()));
+        this.baseLUT = base.getClassEntries().stream()
+                .collect(Collectors.toMap(it -> it.getClassNames().get(this.baseSharedIndex), Function.identity()));
+        this.baseToOverlayIndexMap = base.getHeader().getNamespaces()
+                .stream().map(it -> overlay.getHeader().getNamespaces().indexOf(it))
+                .map(it -> it < 0 ? null : it)
+                .toArray(Integer[]::new);
+    }
+
+    public @NotNull TinyFile merge() {
+        var baseHeader = base.getHeader();
+        return new TinyFile(
+                new TinyHeader(baseHeader.getNamespaces(), baseHeader.getMajorVersion(), baseHeader.getMinorVersion(), baseHeader.getProperties()),
+                mergeClasses()
+        );
+    }
+
+
+    protected abstract @NotNull List<@NotNull String> mergeNames(
+            @NotNull List<@NotNull String> baseNames,
+            @NotNull List<@NotNull String> overlayNames);
+
+    protected <T> List<String> mergeNames(@Nullable T baseObject, @Nullable T overlayObject, @NotNull Function<@NotNull T, @NotNull List<@NotNull String>> nameExtractor) {
+        if (baseObject == null && overlayObject == null)
+            throw new IllegalArgumentException("Either baseObject or overlayObject must not be null");
+        if (overlayObject == null)
+            return nameExtractor.apply(baseObject);
+        if (baseObject == null) {
+            if (base.getHeader().getNamespaces().size() != overlay.getHeader().getNamespaces().size())
+                throw new IllegalStateException("Can only merge unnamed base objects if both files have the same namespaces");
+            List<String> names = new ArrayList<>();
+            for (int i = 0; i < base.getHeader().getNamespaces().size(); i++) {
+                names.add("");
+            }
+            List<@NotNull String> mergedNames = mergeNames(names, nameExtractor.apply(overlayObject));
+            for (var name : mergedNames) {
+                if (name.isEmpty())
+                    throw new RuntimeException("Invariant violated: unnamed overlay name");
+            }
+            return mergedNames;
+        }
+        return mergeNames(nameExtractor.apply(baseObject), nameExtractor.apply(overlayObject));
+    }
+
+    protected abstract <T> @NotNull List<@NotNull String> mergeComments(@Nullable T baseObject, @Nullable T overlayObject, @NotNull Function<@NotNull T, ? extends @NotNull Collection<@NotNull String>> commentExtractor);
+
+    protected TinyClass mergeClass(@Nullable TinyClass baseClass, @Nullable TinyClass overlayClass) {
+        return new TinyClass(
+                mergeNames(baseClass, overlayClass, TinyClass::getClassNames),
+                mergeMethods(baseClass, overlayClass),
+                mergeFields(baseClass, overlayClass),
+                mergeComments(baseClass, overlayClass, TinyClass::getComments)
+        );
+    }
+
+    protected Collection<TinyField> mergeFields(@Nullable TinyClass baseClass, @Nullable TinyClass overlayClass) {
+        return mergeChildren(baseClass, overlayClass,
+                it -> it.getFieldNames().get(baseSharedIndex),
+                it -> it.getFieldNames().get(overlaySharedIndex),
+                TinyClass::getFields,
+                this::mergeField);
+    }
+
+    protected TinyField mergeField(@Nullable TinyField baseField, @Nullable TinyField overlayField) {
+        var mergedNames = mergeNames(baseField, overlayField, TinyField::getFieldNames);
+        return new TinyField(
+                mergedNames.get(0),
+                mergedNames,
+                mergeComments(baseField, overlayField, TinyField::getComments)
+        );
+    }
+
+    protected Collection<TinyMethod> mergeMethods(@Nullable TinyClass baseClass, @Nullable TinyClass overlayClass) {
+        return mergeChildren(baseClass, overlayClass,
+                it -> it.getMethodNames().get(baseSharedIndex),
+                it -> it.getMethodNames().get(overlaySharedIndex),
+                TinyClass::getMethods,
+                this::mergeMethod);
+    }
+
+    protected @NotNull TinyMethod mergeMethod(@Nullable TinyMethod baseMethod, @Nullable TinyMethod overlayMethod) {
+        var mergedNames = mergeNames(baseMethod, overlayMethod, TinyMethod::getMethodNames);
+        return new TinyMethod(
+                mergedNames.get(0),
+                mergedNames,
+                mergeChildren(baseMethod, overlayMethod,
+                        TinyMethodParameter::getLvIndex,
+                        TinyMethodParameter::getLvIndex,
+                        TinyMethod::getParameters,
+                        this::mergeMethodParameter),
+                mergeChildren(baseMethod, overlayMethod,
+                        TinyLocalVariable::getLvIndex,
+                        TinyLocalVariable::getLvIndex,
+                        TinyMethod::getLocalVariables,
+                        this::mergeLocalVariable),
+                mergeComments(baseMethod, overlayMethod, TinyMethod::getComments)
+        );
+    }
+
+    protected TinyLocalVariable mergeLocalVariable(
+            @Nullable TinyLocalVariable baseVariable,
+            @Nullable TinyLocalVariable overlayVariable) {
+        if (baseVariable == null && overlayVariable == null)
+            throw new NullPointerException();
+        if (overlayVariable != null && baseVariable != null) {
+            if (overlayVariable.getLvIndex() != baseVariable.getLvIndex())
+                throw new IllegalArgumentException("lvIndex must match on local variables with the same shared name");
+            if (overlayVariable.getLvStartOffset() != baseVariable.getLvStartOffset())
+                throw new IllegalArgumentException("lvStartOffset must match on local variables with the same shared name");
+            if (overlayVariable.getLvTableIndex() != baseVariable.getLvTableIndex())
+                throw new IllegalArgumentException("lvTableIndex must match on local variables with the same shared name");
+        }
+        return new TinyLocalVariable(
+                baseVariable == null ? overlayVariable.getLvIndex() : baseVariable.getLvIndex(),
+                baseVariable == null ? overlayVariable.getLvStartOffset() : baseVariable.getLvStartOffset(),
+                baseVariable == null ? overlayVariable.getLvTableIndex() : baseVariable.getLvTableIndex(),
+                mergeNames(baseVariable, overlayVariable, TinyLocalVariable::getLocalVariableNames),
+                mergeComments(baseVariable, overlayVariable, TinyLocalVariable::getComments)
+        );
+    }
+
+    protected TinyMethodParameter mergeMethodParameter(
+            @Nullable TinyMethodParameter baseParam,
+            @Nullable TinyMethodParameter overlayParam) {
+        if (baseParam == null && overlayParam == null)
+            throw new NullPointerException();
+        if (overlayParam != null && baseParam != null && overlayParam.getLvIndex() != baseParam.getLvIndex())
+            throw new IllegalArgumentException("lvIndex must match on method parameters with the same shared name");
+        return new TinyMethodParameter(
+                baseParam != null ? baseParam.getLvIndex() : overlayParam.getLvIndex(),
+                mergeNames(baseParam, overlayParam, TinyMethodParameter::getParameterNames),
+                mergeComments(baseParam, overlayParam, TinyMethodParameter::getComments)
+        );
+    }
+
+    protected <B, C, K> List<C> mergeChildren(
+            @Nullable B base,
+            @Nullable B overlay,
+            @NotNull Function<@NotNull C, @NotNull K> sharedBaseName,
+            @NotNull Function<@NotNull C, @NotNull K> sharedOverlayName,
+            @NotNull Function<@NotNull B, ? extends @NotNull Collection<@NotNull C>> childExtractor,
+            @NotNull BiFunction<@Nullable C, @Nullable C, @NotNull C> childMerger
+    ) {
+        Map<K, C> baseIndex =
+                base == null ? Collections.emptyMap() : childExtractor.apply(base).stream()
+                        .collect(Collectors.toMap(sharedBaseName, Function.identity()));
+        Map<K, C> overlayIndex =
+                overlay == null ? Collections.emptyMap() : childExtractor.apply(overlay).stream()
+                        .collect(Collectors.toMap(sharedOverlayName, Function.identity()));
+        var results = new ArrayList<C>();
+        if (base != null)
+            for (var baseObject : childExtractor.apply(base)) {
+                var sharedName = sharedBaseName.apply(baseObject);
+                var overlayObject = overlayIndex.get(sharedName);
+                var mergedClass = childMerger.apply(baseObject, overlayObject);
+                results.add(mergedClass);
+            }
+        if (overlay != null)
+            for (var overlayObject : childExtractor.apply(overlay)) {
+                var sharedName = sharedOverlayName.apply(overlayObject);
+                if (baseIndex.containsKey(sharedName))
+                    continue;
+                results.add(overlayObject);
+            }
+        return results;
+    }
+
+    protected List<TinyClass> mergeClasses() {
+        return (mergeChildren(
+                base,
+                overlay,
+                it -> it.getClassNames().get(baseSharedIndex),
+                it -> it.getClassNames().get(overlaySharedIndex),
+                TinyFile::getClassEntries,
+                this::mergeClass
+        ));
+    }
+}