7 files changed, 409 insertions, 39 deletions

diff --git a/src/main/java/cuchaz/enigma/CommandMain.java b/src/main/java/cuchaz/enigma/CommandMain.java
index 5b25087..5e7a1af 100644
--- a/src/main/java/cuchaz/enigma/CommandMain.java
+++ b/src/main/java/cuchaz/enigma/CommandMain.java
@@ -83,6 +83,8 @@ public class CommandMain {
                 register(new DeobfuscateCommand());
                 register(new DecompileCommand());
                 register(new ConvertMappingsCommand());
+                register(new ComposeMappingsCommand());
+                register(new InvertMappingsCommand());
                 register(new CheckMappingsCommand());
         }
 
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);
+    }
+}
diff --git a/src/main/java/cuchaz/enigma/translation/mapping/serde/TinyMappingsWriter.java b/src/main/java/cuchaz/enigma/translation/mapping/serde/TinyMappingsWriter.java
new file mode 100644
index 0000000..0a52dad
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/translation/mapping/serde/TinyMappingsWriter.java
@@ -0,0 +1,144 @@
+package cuchaz.enigma.translation.mapping.serde;
+
+import com.google.common.base.Joiner;
+import com.google.common.collect.Lists;
+import cuchaz.enigma.ProgressListener;
+import cuchaz.enigma.translation.MappingTranslator;
+import cuchaz.enigma.translation.Translator;
+import cuchaz.enigma.translation.mapping.EntryMapping;
+import cuchaz.enigma.translation.mapping.MappingDelta;
+import cuchaz.enigma.translation.mapping.VoidEntryResolver;
+import cuchaz.enigma.translation.mapping.tree.EntryTree;
+import cuchaz.enigma.translation.mapping.tree.EntryTreeNode;
+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.BufferedWriter;
+import java.io.IOException;
+import java.io.Writer;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Comparator;
+import java.util.HashSet;
+import java.util.Set;
+
+public class TinyMappingsWriter implements MappingsWriter {
+    private static final String VERSION_CONSTANT = "v1";
+    private static final Joiner TAB_JOINER = Joiner.on('\t');
+
+    // HACK: as of enigma 0.13.1, some fields seem to appear duplicated?
+    private final Set<String> writtenLines = new HashSet<>();
+    private final String nameObf;
+    private final String nameDeobf;
+
+    public TinyMappingsWriter(String nameObf, String nameDeobf) {
+        this.nameObf = nameObf;
+        this.nameDeobf = nameDeobf;
+    }
+
+    @Override
+    public void write(EntryTree<EntryMapping> mappings, MappingDelta<EntryMapping> delta, Path path, ProgressListener progress) {
+        try {
+            Files.deleteIfExists(path);
+            Files.createFile(path);
+        } catch (IOException e) {
+            e.printStackTrace();
+        }
+
+        try (BufferedWriter writer = Files.newBufferedWriter(path, StandardCharsets.UTF_8)) {
+            writeLine(writer, new String[]{VERSION_CONSTANT, nameObf, nameDeobf});
+
+            Lists.newArrayList(mappings).stream()
+                    .map(EntryTreeNode::getEntry).sorted(Comparator.comparing(Object::toString))
+                    .forEach(entry -> writeEntry(writer, mappings, entry));
+        } catch (IOException e) {
+            e.printStackTrace();
+        }
+    }
+
+    private void writeEntry(Writer writer, EntryTree<EntryMapping> mappings, Entry<?> entry) {
+        EntryTreeNode<EntryMapping> node = mappings.findNode(entry);
+        if (node == null) {
+            return;
+        }
+
+        Translator translator = new MappingTranslator(mappings, VoidEntryResolver.INSTANCE);
+
+        EntryMapping mapping = mappings.get(entry);
+        if (mapping != null && !entry.getName().equals(mapping.getTargetName())) {
+            if (entry instanceof ClassEntry) {
+                writeClass(writer, (ClassEntry) entry, translator);
+            } else if (entry instanceof FieldEntry) {
+                writeLine(writer, serializeEntry(entry, mapping.getTargetName()));
+            } else if (entry instanceof MethodEntry) {
+                writeLine(writer, serializeEntry(entry, mapping.getTargetName()));
+            }
+        }
+
+        writeChildren(writer, mappings, node);
+    }
+
+    private void writeChildren(Writer writer, EntryTree<EntryMapping> mappings, EntryTreeNode<EntryMapping> node) {
+        node.getChildren().stream()
+                .filter(e -> e instanceof FieldEntry).sorted()
+                .forEach(child -> writeEntry(writer, mappings, child));
+
+        node.getChildren().stream()
+                .filter(e -> e instanceof MethodEntry).sorted()
+                .forEach(child -> writeEntry(writer, mappings, child));
+
+        node.getChildren().stream()
+                .filter(e -> e instanceof ClassEntry).sorted()
+                .forEach(child -> writeEntry(writer, mappings, child));
+    }
+
+    private void writeClass(Writer writer, ClassEntry entry, Translator translator) {
+        ClassEntry translatedEntry = translator.translate(entry);
+
+        String obfClassName = entry.getFullName();
+        String deobfClassName = translatedEntry.getFullName();
+        writeLine(writer, new String[]{"CLASS", obfClassName, deobfClassName});
+    }
+
+    private void writeLine(Writer writer, String[] data) {
+        try {
+            String line = TAB_JOINER.join(data) + "\n";
+            if (writtenLines.add(line)) {
+                writer.write(line);
+            }
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private String[] serializeEntry(Entry<?> entry, String... extraFields) {
+        String[] data = null;
+
+        if (entry instanceof FieldEntry) {
+            data = new String[4 + extraFields.length];
+            data[0] = "FIELD";
+            data[1] = entry.getContainingClass().getFullName();
+            data[2] = ((FieldEntry) entry).getDesc().toString();
+            data[3] = entry.getName();
+        } else if (entry instanceof MethodEntry) {
+            data = new String[4 + extraFields.length];
+            data[0] = "METHOD";
+            data[1] = entry.getContainingClass().getFullName();
+            data[2] = ((MethodEntry) entry).getDesc().toString();
+            data[3] = entry.getName();
+        } else if (entry instanceof ClassEntry) {
+            data = new String[2 + extraFields.length];
+            data[0] = "CLASS";
+            data[1] = ((ClassEntry) entry).getFullName();
+        }
+
+        if (data != null) {
+            System.arraycopy(extraFields, 0, data, data.length - extraFields.length, extraFields.length);
+        }
+
+        return data;
+    }
+}
diff --git a/src/main/java/cuchaz/enigma/utils/Utils.java b/src/main/java/cuchaz/enigma/utils/Utils.java
index bd09c64..6788042 100644
--- a/src/main/java/cuchaz/enigma/utils/Utils.java
+++ b/src/main/java/cuchaz/enigma/utils/Utils.java
@@ -21,8 +21,11 @@ import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.net.URI;
 import java.net.URISyntaxException;
-import java.util.Arrays;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Comparator;
 import java.util.List;
+import java.util.stream.Collectors;
 
 public class Utils {
 
@@ -94,4 +97,12 @@ public class Utils {
                 String value = System.getProperty(property);
                 return value == null ? defValue : Boolean.parseBoolean(value);
         }
+
+        public static void delete(Path path) throws IOException {
+                if (path.toFile().exists()) {
+                        for (Path p : Files.walk(path).sorted(Comparator.reverseOrder()).collect(Collectors.toList())) {
+                                Files.delete(p);
+                        }
+                }
+        }
 }