Create merger

4 files changed, 281 insertions, 0 deletions

diff --git a/src/main/java/moe/nea/zwirn/Index.java b/src/main/java/moe/nea/zwirn/Index.java
new file mode 100644
index 0000000..a27344e
--- /dev/null
+++ b/src/main/java/moe/nea/zwirn/Index.java
@@ -0,0 +1,30 @@
+package moe.nea.zwirn;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Function;
+
+public class Index<K, V> {
+    Map<K, Set<V>> map = new HashMap<>();
+
+    public void addToIndex(K k, V v) {
+        find(k).add(v);
+    }
+
+    public Set<K> keys() {
+        return map.keySet();
+    }
+
+    public <T> void loadFrom(Iterable<T> it, Function<T, K> keyExtractor, Function<T, V> valueExtractor) {
+        for (T t : it) {
+            addToIndex(keyExtractor.apply(t), valueExtractor.apply(t));
+        }
+    }
+
+    public Set<V> find(K k) {
+        return map.computeIfAbsent(k, ignored -> new HashSet<>());
+    }
+
+}
\ No newline at end of file
diff --git a/src/main/java/moe/nea/zwirn/TinyMerger.java b/src/main/java/moe/nea/zwirn/TinyMerger.java
new file mode 100644
index 0000000..97b245b
--- /dev/null
+++ b/src/main/java/moe/nea/zwirn/TinyMerger.java
@@ -0,0 +1,227 @@
+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;
+
+class TinyMerger {
+    private final @NotNull TinyFile base;
+    private final @NotNull TinyFile overlay;
+    private final @NotNull String sharedNamespace;
+    private final int baseSharedIndex;
+    private final int overlaySharedIndex;
+    private final @NotNull Map<@NotNull String, @NotNull TinyClass> overlayLUT;
+    private final @NotNull Map<@NotNull String, @NotNull TinyClass> baseLUT;
+    private final Integer[] baseToOverlayIndexMap;
+    private final @NotNull List<@NotNull TinyClass> entries = new ArrayList<>();
+
+    public TinyMerger(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() {
+        mergeClasses();
+        var baseHeader = base.getHeader();
+        return new TinyFile(
+                new TinyHeader(baseHeader.getNamespaces(), baseHeader.getMajorVersion(), baseHeader.getMinorVersion(), baseHeader.getProperties()),
+                entries
+        );
+    }
+
+
+    private @NotNull List<@NotNull String> mergeNames(
+            @NotNull List<@NotNull String> baseNames,
+            @NotNull List<@NotNull String> overlayNames) {
+        List<String> mergedNames = new ArrayList<>(baseNames);
+        String sharedName = baseNames.get(baseSharedIndex);
+        for (int i = 0; i < mergedNames.size(); i++) {
+            Integer overlayIndex = this.baseToOverlayIndexMap[i];
+            if (overlayIndex == null) continue;
+            String overlayName = overlayNames.get(overlayIndex);
+            if (overlayName.isEmpty() || sharedName.equals(overlayName))
+                continue;
+            mergedNames.set(i, overlayName);
+        }
+        return mergedNames;
+    }
+
+    private <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));
+    }
+
+    private <T> @NotNull List<@NotNull String> mergeComments(@Nullable T baseObject, @Nullable T overlayObject, @NotNull Function<@NotNull T, ? extends @NotNull Collection<@NotNull String>> commentExtractor) {
+        List<String> comments = new ArrayList<>();
+        if (baseObject != null)
+            comments.addAll(commentExtractor.apply(baseObject));
+        if (overlayObject != null)
+            comments.addAll(commentExtractor.apply(overlayObject));
+        return comments;
+    }
+
+    private 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)
+        );
+    }
+
+    private 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);
+    }
+
+    private 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)
+        );
+    }
+
+    private 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);
+    }
+
+    private @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)
+        );
+    }
+
+    private 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)
+        );
+    }
+
+    private 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)
+        );
+    }
+
+    private <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;
+    }
+
+    private void mergeClasses() {
+        entries.addAll(mergeChildren(
+                base,
+                overlay,
+                it -> it.getClassNames().get(baseSharedIndex),
+                it -> it.getClassNames().get(overlaySharedIndex),
+                TinyFile::getClassEntries,
+                this::mergeClass
+        ));
+    }
+}
diff --git a/src/main/java/moe/nea/zwirn/Zwirn.java b/src/main/java/moe/nea/zwirn/Zwirn.java
new file mode 100644
index 0000000..db616d7
--- /dev/null
+++ b/src/main/java/moe/nea/zwirn/Zwirn.java
@@ -0,0 +1,17 @@
+package moe.nea.zwirn;
+
+import net.fabricmc.stitch.commands.tinyv2.TinyFile;
+import org.jetbrains.annotations.NotNull;
+
+public class Zwirn {
+    public static @NotNull TinyFile mergeTinyFile(
+            @NotNull TinyFile base, @NotNull TinyFile overlay,
+            @NotNull String sharedNamespace) {
+        var namespaces = base.getHeader().getNamespaces();
+        if (!namespaces.containsAll(overlay.getHeader().getNamespaces()))
+            throw new IllegalArgumentException("When merging a tiny file, overlay may not introduce a new namespace.");
+        if (!overlay.getHeader().getNamespaces().contains(sharedNamespace))
+            throw new IllegalArgumentException("When merging a tiny file, overlay must contain the shared namespace");
+        return new TinyMerger(base, overlay, sharedNamespace).merge();
+    }
+}
diff --git a/src/main/java/moe/nea/zwirn/ZwirnUtil.java b/src/main/java/moe/nea/zwirn/ZwirnUtil.java
new file mode 100644
index 0000000..f17eae6
--- /dev/null
+++ b/src/main/java/moe/nea/zwirn/ZwirnUtil.java
@@ -0,0 +1,7 @@
+package moe.nea.zwirn;
+
+public class ZwirnUtil {
+    private static <T> T TODO(String message) {
+        throw new RuntimeException("Not yet implemented: " + message);
+    }
+}