4 files changed, 251 insertions, 38 deletions

diff --git a/src/main/java/cuchaz/enigma/command/ComposeMappingsCommand.java b/src/main/java/cuchaz/enigma/command/ComposeMappingsCommand.java
new file mode 100644
index 0000000..1f6a069
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/command/ComposeMappingsCommand.java
@@ -0,0 +1,37 @@
+package cuchaz.enigma.command;
+
+import cuchaz.enigma.throwables.MappingParseException;
+import cuchaz.enigma.translation.mapping.EntryMapping;
+import cuchaz.enigma.translation.mapping.tree.EntryTree;
+import cuchaz.enigma.utils.Utils;
+
+import java.io.IOException;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+
+public class ComposeMappingsCommand extends Command {
+    public ComposeMappingsCommand() {
+        super("compose-mappings");
+    }
+
+    @Override
+    public String getUsage() {
+        return "<left-format> <left> <right-format> <right> <result-format> <result> <keep-mode>";
+    }
+
+    @Override
+    public boolean isValidArgument(int length) {
+        return length == 7;
+    }
+
+    @Override
+    public void run(String... args) throws IOException, MappingParseException {
+        EntryTree<EntryMapping> left = MappingCommandsUtil.read(args[0], Paths.get(args[1]));
+        EntryTree<EntryMapping> right = MappingCommandsUtil.read(args[2], Paths.get(args[3]));
+        EntryTree<EntryMapping> result = MappingCommandsUtil.compose(left, right, args[6].equals("left") || args[6].equals("both"), args[6].equals("right") || args[6].equals("both"));
+
+        Path output = Paths.get(args[5]);
+        Utils.delete(output);
+        MappingCommandsUtil.write(result, args[4], output);
+    }
+}
diff --git a/src/main/java/cuchaz/enigma/command/ConvertMappingsCommand.java b/src/main/java/cuchaz/enigma/command/ConvertMappingsCommand.java
index 75d3791..775bd3e 100644
--- a/src/main/java/cuchaz/enigma/command/ConvertMappingsCommand.java
+++ b/src/main/java/cuchaz/enigma/command/ConvertMappingsCommand.java
@@ -1,47 +1,35 @@
 package cuchaz.enigma.command;
 
+import cuchaz.enigma.throwables.MappingParseException;
 import cuchaz.enigma.translation.mapping.EntryMapping;
-import cuchaz.enigma.translation.mapping.serde.MappingFormat;
 import cuchaz.enigma.translation.mapping.tree.EntryTree;
+import cuchaz.enigma.utils.Utils;
 
-import java.io.File;
+import java.io.IOException;
 import java.nio.file.Path;
-import java.util.Locale;
+import java.nio.file.Paths;
 
 public class ConvertMappingsCommand extends Command {
-
-        public ConvertMappingsCommand() {
-                super("convertmappings");
-        }
-
-        @Override
-        public String getUsage() {
-                return "<enigma mappings> <converted mappings> <ENIGMA_FILE|ENIGMA_DIRECTORY|SRG_FILE>";
-        }
-
-        @Override
-        public boolean isValidArgument(int length) {
-                return length == 3;
-        }
-
-        @Override
-        public void run(String... args) throws Exception {
-                Path fileMappings = getReadablePath(getArg(args, 0, "enigma mappings", true));
-                File result = getWritableFile(getArg(args, 1, "converted mappings", true));
-                String name = getArg(args, 2, "format desc", true);
-                MappingFormat saveFormat;
-                try {
-                        saveFormat = MappingFormat.valueOf(name.toUpperCase(Locale.ROOT));
-                } catch (IllegalArgumentException e) {
-                        throw new IllegalArgumentException(name + "is not a valid mapping format!");
-                }
-
-                System.out.println("Reading mappings...");
-
-                MappingFormat readFormat = chooseEnigmaFormat(fileMappings);
-                EntryTree<EntryMapping> mappings = readFormat.read(fileMappings, new ConsoleProgressListener());
-                System.out.println("Saving new mappings...");
-
-                saveFormat.write(mappings, result.toPath(), new ConsoleProgressListener());
-        }
+    public ConvertMappingsCommand() {
+        super("convert-mappings");
+    }
+
+    @Override
+    public String getUsage() {
+        return "<source-format> <source> <result-format> <result>";
+    }
+
+    @Override
+    public boolean isValidArgument(int length) {
+        return length == 4;
+    }
+
+    @Override
+    public void run(String... args) throws IOException, MappingParseException {
+        EntryTree<EntryMapping> mappings = MappingCommandsUtil.read(args[0], Paths.get(args[1]));
+
+        Path output = Paths.get(args[3]);
+        Utils.delete(output);
+        MappingCommandsUtil.write(mappings, args[2], output);
+    }
 }
diff --git a/src/main/java/cuchaz/enigma/command/InvertMappingsCommand.java b/src/main/java/cuchaz/enigma/command/InvertMappingsCommand.java
new file mode 100644
index 0000000..bfe8308
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/command/InvertMappingsCommand.java
@@ -0,0 +1,36 @@
+package cuchaz.enigma.command;
+
+import cuchaz.enigma.throwables.MappingParseException;
+import cuchaz.enigma.translation.mapping.EntryMapping;
+import cuchaz.enigma.translation.mapping.tree.EntryTree;
+import cuchaz.enigma.utils.Utils;
+
+import java.io.IOException;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+
+public class InvertMappingsCommand extends Command {
+    public InvertMappingsCommand() {
+        super("invert-mappings");
+    }
+
+    @Override
+    public String getUsage() {
+        return "<source-format> <source> <result-format> <result>";
+    }
+
+    @Override
+    public boolean isValidArgument(int length) {
+        return length == 4;
+    }
+
+    @Override
+    public void run(String... args) throws IOException, MappingParseException {
+        EntryTree<EntryMapping> source = MappingCommandsUtil.read(args[0], Paths.get(args[1]));
+        EntryTree<EntryMapping> result = MappingCommandsUtil.invert(source);
+
+        Path output = Paths.get(args[3]);
+        Utils.delete(output);
+        MappingCommandsUtil.write(result, args[2], output);
+    }
+}
diff --git a/src/main/java/cuchaz/enigma/command/MappingCommandsUtil.java b/src/main/java/cuchaz/enigma/command/MappingCommandsUtil.java
new file mode 100644
index 0000000..4679575
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/command/MappingCommandsUtil.java
@@ -0,0 +1,152 @@
+package cuchaz.enigma.command;
+
+import cuchaz.enigma.ProgressListener;
+import cuchaz.enigma.throwables.MappingParseException;
+import cuchaz.enigma.translation.MappingTranslator;
+import cuchaz.enigma.translation.Translator;
+import cuchaz.enigma.translation.mapping.EntryMapping;
+import cuchaz.enigma.translation.mapping.VoidEntryResolver;
+import cuchaz.enigma.translation.mapping.serde.*;
+import cuchaz.enigma.translation.mapping.tree.EntryTree;
+import cuchaz.enigma.translation.mapping.tree.EntryTreeNode;
+import cuchaz.enigma.translation.mapping.tree.HashEntryTree;
+import cuchaz.enigma.translation.representation.entry.ClassEntry;
+import cuchaz.enigma.translation.representation.entry.Entry;
+import cuchaz.enigma.translation.representation.entry.FieldEntry;
+import cuchaz.enigma.translation.representation.entry.MethodEntry;
+
+import java.io.IOException;
+import java.nio.file.Path;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+public final class MappingCommandsUtil {
+    public static void main(String[] args) throws Exception {
+        new InvertMappingsCommand().run(
+                "enigma",
+                "D:\\IdeaProjects\\yarn\\mappings",
+                "enigma",
+                "D:\\IdeaProjects\\Enigma\\converted");
+    }
+
+    private MappingCommandsUtil() {}
+
+    public static EntryTree<EntryMapping> invert(EntryTree<EntryMapping> mappings) {
+        Translator translator = new MappingTranslator(mappings, VoidEntryResolver.INSTANCE);
+        EntryTree<EntryMapping> result = new HashEntryTree<>();
+
+        for (EntryTreeNode<EntryMapping> node : mappings) {
+            Entry<?> leftEntry = node.getEntry();
+            EntryMapping leftMapping = node.getValue();
+
+            if (!(leftEntry instanceof ClassEntry || leftEntry instanceof MethodEntry || leftEntry instanceof FieldEntry)) {
+                result.insert(translator.translate(leftEntry), leftMapping);
+                continue;
+            }
+
+            Entry<?> rightEntry = translator.translate(leftEntry);
+
+            result.insert(rightEntry, leftMapping == null ? null : new EntryMapping(leftEntry.getName())); // TODO: leftMapping.withName once javadoc PR is merged
+        }
+
+        return result;
+    }
+
+    @SuppressWarnings("unchecked")
+    public static EntryTree<EntryMapping> compose(EntryTree<EntryMapping> left, EntryTree<EntryMapping> right, boolean keepLeftOnly, boolean keepRightOnly) {
+        Translator leftTranslator = new MappingTranslator(left, VoidEntryResolver.INSTANCE);
+        EntryTree<EntryMapping> result = new HashEntryTree<>();
+        Map<Entry<?>, Entry<?>> rightToLeft = new HashMap<>();
+        Set<Entry<?>> addedMappings = new HashSet<>();
+
+        for (EntryTreeNode<EntryMapping> node : left) {
+            Entry<?> leftEntry = node.getEntry();
+            EntryMapping leftMapping = node.getValue();
+
+            Entry<?> rightEntry = leftTranslator.translate(leftEntry);
+            rightToLeft.put(rightEntry, leftEntry);
+
+            EntryMapping rightMapping = right.get(rightEntry);
+            if (rightMapping != null) {
+                result.insert(leftEntry, rightMapping);
+                addedMappings.add(rightEntry);
+            } else if (keepLeftOnly) {
+                result.insert(leftEntry, leftMapping);
+            }
+        }
+
+        if (keepRightOnly) {
+            for (EntryTreeNode<EntryMapping> node : right) {
+                Entry<?> rightEntry = node.getEntry();
+                EntryMapping rightMapping = node.getValue();
+
+                if (addedMappings.contains(rightEntry)) {
+                    continue;
+                }
+
+                Entry<?> parent = rightEntry.getParent();
+                Entry<?> correctEntry = rightEntry;
+                if (rightToLeft.containsKey(parent)) {
+                    correctEntry = ((Entry<Entry<?>>) rightEntry).withParent(rightToLeft.get(parent));
+                }
+
+                result.insert(correctEntry, rightMapping);
+                rightToLeft.put(rightEntry, correctEntry);
+            }
+        }
+        return result;
+    }
+
+    public static EntryTree<EntryMapping> read(String type, Path path) throws MappingParseException, IOException {
+        if (type.equals("enigma")) {
+            return EnigmaMappingsReader.DIRECTORY.read(path, ProgressListener.none());
+        }
+
+        if (type.equals("tiny")) {
+            return TinyMappingsReader.INSTANCE.read(path, ProgressListener.none());
+        }
+
+        MappingFormat format = null;
+        try {
+            format = MappingFormat.valueOf(type.toUpperCase());
+        } catch (IllegalArgumentException ignored) {}
+
+        if (format != null) {
+            return format.getReader().read(path, ProgressListener.none());
+        }
+
+        throw new IllegalArgumentException("no reader for " + type);
+    }
+
+    public static void write(EntryTree<EntryMapping> mappings, String type, Path path) {
+        if (type.equals("enigma")) {
+            EnigmaMappingsWriter.DIRECTORY.write(mappings, path, ProgressListener.none());
+            return;
+        }
+
+        if (type.startsWith("tiny")) {
+            String[] split = type.split(":");
+
+            if (split.length != 3) {
+                throw new IllegalArgumentException("specify column names as 'tiny:from_column:to_column'");
+            }
+
+            new TinyMappingsWriter(split[1], split[2]).write(mappings, path, ProgressListener.none());
+            return;
+        }
+
+        MappingFormat format = null;
+        try {
+            format = MappingFormat.valueOf(type.toUpperCase());
+        } catch (IllegalArgumentException ignored) {}
+
+        if (format != null) {
+            format.getWriter().write(mappings, path, ProgressListener.none());
+            return;
+        }
+
+        throw new IllegalArgumentException("no writer for " + type);
+    }
+}