4 files changed, 147 insertions, 93 deletions

diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
index e7e75669..73472ca9 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
@@ -29,6 +29,7 @@ import javax.swing.JOptionPane;
 import javax.swing.SwingUtilities;
 
 import com.google.common.collect.Lists;
+import net.fabricmc.mappingio.MappingReader;
 import net.fabricmc.mappingio.MappingWriter;
 import net.fabricmc.mappingio.tree.MemoryMappingTree;
 
@@ -102,6 +103,7 @@ public class GuiController implements ClientPacketHandler {
 
         private Path loadedMappingPath;
         private MappingFormat loadedMappingFormat;
+        public boolean useMappingIo;
 
         private ClassHandleProvider chp;
 
@@ -152,8 +154,16 @@ public class GuiController implements ClientPacketHandler {
                 return ProgressDialog.runOffThread(gui.getFrame(), progress -> {
                         try {
                                 MappingSaveParameters saveParameters = enigma.getProfile().getMappingSaveParameters();
+                                EntryTree<EntryMapping> mappings;
+
+                                if (useMappingIo) {
+                                        MemoryMappingTree mappingTree = new MemoryMappingTree();
+                                        MappingReader.read(path, format.getMappingIoCounterpart(), mappingTree);
+                                        mappings = MappingIoConverter.fromMappingIo(mappingTree);
+                                } else {
+                                        mappings = format.read(path, progress, saveParameters);
+                                }
 
-                                EntryTree<EntryMapping> mappings = format.read(path, progress, saveParameters);
                                 project.setMappings(mappings);
 
                                 loadedMappingFormat = format;
@@ -208,7 +218,12 @@ public class GuiController implements ClientPacketHandler {
                         loadedMappingFormat = format;
                         loadedMappingPath = path;
 
-                        if (saveAll) {
+                        if (useMappingIo) {
+                                MemoryMappingTree mappingTree = MappingIoConverter.toMappingIo(mapper.getObfToDeobf());
+                                MappingWriter writer = MappingWriter.create(path, format.getMappingIoCounterpart());
+                                mappingTree.accept(writer);
+                                writer.close();
+                        } else if (saveAll) {
                                 format.write(mapper.getObfToDeobf(), path, progress, saveParameters);
                         } else {
                                 format.write(mapper.getObfToDeobf(), delta, path, progress, saveParameters);
@@ -216,48 +231,6 @@ public class GuiController implements ClientPacketHandler {
                 });
         }
 
-        public CompletableFuture<Void> saveMappings(Path path, net.fabricmc.mappingio.format.MappingFormat format) {
-                if (project == null) {
-                        return CompletableFuture.completedFuture(null);
-                }
-
-                return ProgressDialog.runOffThread(this.gui.getFrame(), progress -> {
-                        EntryRemapper mapper = project.getMapper();
-                        MappingSaveParameters saveParameters = enigma.getProfile().getMappingSaveParameters();
-
-                        MappingDelta<EntryMapping> delta = mapper.takeMappingDelta();
-                        boolean saveAll = !path.equals(loadedMappingPath);
-
-                        switch (format) {
-                        case ENIGMA:
-                                loadedMappingFormat = MappingFormat.ENIGMA_DIRECTORY;
-                                loadedMappingPath = path;
-                                break;
-                        case PROGUARD:
-                                loadedMappingFormat = MappingFormat.PROGUARD;
-                                loadedMappingPath = path;
-                                break;
-                        case SRG:
-                                loadedMappingFormat = MappingFormat.SRG_FILE;
-                                loadedMappingPath = path;
-                                break;
-                        case TINY:
-                                loadedMappingFormat = MappingFormat.TINY_FILE;
-                                loadedMappingPath = path;
-                                break;
-                        case TINY_2:
-                                loadedMappingFormat = MappingFormat.TINY_V2;
-                                loadedMappingPath = path;
-                                break;
-                        }
-
-                        MemoryMappingTree mappingTree = MappingIoConverter.toMappingIo(mapper.getObfToDeobf());
-                        MappingWriter writer = MappingWriter.create(path, format);
-                        mappingTree.accept(writer);
-                        writer.close();
-                });
-        }
-
         public void closeMappings() {
                 if (project == null) {
                         return;
diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java
index 1d4c1ca1..3692eb3b 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java
@@ -46,7 +46,7 @@ public class MenuBar {
         private final JMenu fileMenu = new JMenu();
         private final JMenuItem jarOpenItem = new JMenuItem();
         private final JMenuItem jarCloseItem = new JMenuItem();
-        private final JMenu openMenu = new JMenu();
+        private final JMenu openMappingsMenu = new JMenu();
         private final JMenuItem saveMappingsItem = new JMenuItem();
         private final JMenu saveMappingsAsMenu = new JMenu();
         private final JMenuItem closeMappingsItem = new JMenuItem();
@@ -89,7 +89,7 @@ public class MenuBar {
 
                 this.retranslateUi();
 
-                prepareOpenMenu(this.openMenu, gui);
+                prepareOpenMappingsMenu(this.openMappingsMenu, gui);
                 prepareSaveMappingsAsMenu(this.saveMappingsAsMenu, this.saveMappingsItem, gui);
                 prepareDecompilerMenu(this.decompilerMenu, gui);
                 prepareThemesMenu(this.themesMenu, gui);
@@ -99,7 +99,7 @@ public class MenuBar {
                 this.fileMenu.add(this.jarOpenItem);
                 this.fileMenu.add(this.jarCloseItem);
                 this.fileMenu.addSeparator();
-                this.fileMenu.add(this.openMenu);
+                this.fileMenu.add(this.openMappingsMenu);
                 this.fileMenu.add(this.saveMappingsItem);
                 this.fileMenu.add(this.saveMappingsAsMenu);
                 this.fileMenu.add(this.closeMappingsItem);
@@ -175,7 +175,7 @@ public class MenuBar {
                 this.startServerItem.setText(I18n.translate(connectionState != ConnectionState.HOSTING ? "menu.collab.server.start" : "menu.collab.server.stop"));
 
                 this.jarCloseItem.setEnabled(jarOpen);
-                this.openMenu.setEnabled(jarOpen);
+                this.openMappingsMenu.setEnabled(jarOpen);
                 this.saveMappingsItem.setEnabled(jarOpen && this.gui.enigmaMappingsFileChooser.getSelectedFile() != null && connectionState != ConnectionState.CONNECTED);
                 this.saveMappingsAsMenu.setEnabled(jarOpen);
                 this.closeMappingsItem.setEnabled(jarOpen);
@@ -190,7 +190,7 @@ public class MenuBar {
                 this.fileMenu.setText(I18n.translate("menu.file"));
                 this.jarOpenItem.setText(I18n.translate("menu.file.jar.open"));
                 this.jarCloseItem.setText(I18n.translate("menu.file.jar.close"));
-                this.openMenu.setText(I18n.translate("menu.file.mappings.open"));
+                this.openMappingsMenu.setText(I18n.translate("menu.file.mappings.open"));
                 this.saveMappingsItem.setText(I18n.translate("menu.file.mappings.save"));
                 this.saveMappingsAsMenu.setText(I18n.translate("menu.file.mappings.save_as"));
                 this.closeMappingsItem.setText(I18n.translate("menu.file.mappings.close"));
@@ -403,67 +403,88 @@ public class MenuBar {
                 GuiUtil.openUrl("https://github.com/FabricMC/Enigma");
         }
 
-        private static void prepareOpenMenu(JMenu openMenu, Gui gui) {
+        private static void prepareOpenMappingsMenu(JMenu openMappingsMenu, Gui gui) {
+                List<MappingFormat> readableMappingIoFormats = Arrays.asList(
+                                MappingFormat.ENIGMA_DIRECTORY,
+                                MappingFormat.TINY_FILE,
+                                MappingFormat.TINY_V2,
+                                MappingFormat.SRG_FILE,
+                                MappingFormat.TSRG_FILE,
+                                MappingFormat.TSRG_2_FILE,
+                                MappingFormat.PROGUARD);
+
+                // Enigma's own readers
                 for (MappingFormat format : MappingFormat.values()) {
                         if (format.getReader() != null) {
-                                JMenuItem item = new JMenuItem(I18n.translate("mapping_format." + format.name().toLowerCase(Locale.ROOT)));
-                                item.addActionListener(event -> {
-                                        gui.enigmaMappingsFileChooser.setCurrentDirectory(new File(UiConfig.getLastSelectedDir()));
-
-                                        if (gui.enigmaMappingsFileChooser.showOpenDialog(gui.getFrame()) == JFileChooser.APPROVE_OPTION) {
-                                                File selectedFile = gui.enigmaMappingsFileChooser.getSelectedFile();
-                                                gui.getController().openMappings(format, selectedFile.toPath());
-                                                UiConfig.setLastSelectedDir(gui.enigmaMappingsFileChooser.getCurrentDirectory().toString());
-                                        }
-                                });
-                                openMenu.add(item);
+                                addOpenMappingsMenuEntry(I18n.translate("mapping_format." + format.name().toLowerCase(Locale.ROOT)),
+                                                format, false, openMappingsMenu, gui);
                         }
                 }
+
+                openMappingsMenu.addSeparator();
+
+                // Mapping-IO readers
+                for (MappingFormat format : readableMappingIoFormats) {
+                        addOpenMappingsMenuEntry(I18n.translate(format.getMappingIoCounterpart().name + " (via Mapping-IO, experimental)"),
+                                        format, true, openMappingsMenu, gui);
+                }
+        }
+
+        private static void addOpenMappingsMenuEntry(String text, MappingFormat format, boolean mappingIo, JMenu openMappingsMenu, Gui gui) {
+                JMenuItem item = new JMenuItem(text);
+                item.addActionListener(event -> {
+                        gui.enigmaMappingsFileChooser.setCurrentDirectory(new File(UiConfig.getLastSelectedDir()));
+
+                        if (gui.enigmaMappingsFileChooser.showOpenDialog(gui.getFrame()) == JFileChooser.APPROVE_OPTION) {
+                                File selectedFile = gui.enigmaMappingsFileChooser.getSelectedFile();
+                                gui.getController().useMappingIo = mappingIo;
+                                gui.getController().openMappings(format, selectedFile.toPath());
+                                UiConfig.setLastSelectedDir(gui.enigmaMappingsFileChooser.getCurrentDirectory().toString());
+                        }
+                });
+                openMappingsMenu.add(item);
         }
 
         private static void prepareSaveMappingsAsMenu(JMenu saveMappingsAsMenu, JMenuItem saveMappingsItem, Gui gui) {
+                List<MappingFormat> writableMappingIoFormats = Arrays.asList(
+                                MappingFormat.ENIGMA_DIRECTORY,
+                                MappingFormat.TINY_V2);
+
+                // Enigma's own writers
                 for (MappingFormat format : MappingFormat.values()) {
                         if (format.getWriter() != null) {
-                                JMenuItem item = new JMenuItem(I18n.translate("mapping_format." + format.name().toLowerCase(Locale.ROOT)));
-                                item.addActionListener(event -> {
-                                        // TODO: Use a specific file chooser for it
-                                        if (gui.enigmaMappingsFileChooser.getCurrentDirectory() == null) {
-                                                gui.enigmaMappingsFileChooser.setCurrentDirectory(new File(UiConfig.getLastSelectedDir()));
-                                        }
-
-                                        if (gui.enigmaMappingsFileChooser.showSaveDialog(gui.getFrame()) == JFileChooser.APPROVE_OPTION) {
-                                                gui.getController().saveMappings(gui.enigmaMappingsFileChooser.getSelectedFile().toPath(), format);
-                                                saveMappingsItem.setEnabled(true);
-                                                UiConfig.setLastSelectedDir(gui.enigmaMappingsFileChooser.getCurrentDirectory().toString());
-                                        }
-                                });
-                                saveMappingsAsMenu.add(item);
+                                addSaveMappingsAsMenuEntry(I18n.translate("mapping_format." + format.name().toLowerCase(Locale.ROOT)),
+                                                format, false, saveMappingsAsMenu, saveMappingsItem, gui);
                         }
                 }
 
                 saveMappingsAsMenu.addSeparator();
 
-                List<net.fabricmc.mappingio.format.MappingFormat> writableMappingIoFormats = Arrays.asList(
-                                net.fabricmc.mappingio.format.MappingFormat.ENIGMA,
-                                net.fabricmc.mappingio.format.MappingFormat.TINY_2);
-                for (net.fabricmc.mappingio.format.MappingFormat format : writableMappingIoFormats) {
-                        JMenuItem item = new JMenuItem(format.name + " (via mapping-io, experimental)");
-                        item.addActionListener(event -> {
-                                // TODO: Use a specific file chooser for it
-                                if (gui.enigmaMappingsFileChooser.getCurrentDirectory() == null) {
-                                        gui.enigmaMappingsFileChooser.setCurrentDirectory(new File(UiConfig.getLastSelectedDir()));
-                                }
-
-                                if (gui.enigmaMappingsFileChooser.showSaveDialog(gui.getFrame()) == JFileChooser.APPROVE_OPTION) {
-                                        gui.getController().saveMappings(gui.enigmaMappingsFileChooser.getSelectedFile().toPath(), format);
-                                        saveMappingsItem.setEnabled(true);
-                                        UiConfig.setLastSelectedDir(gui.enigmaMappingsFileChooser.getCurrentDirectory().toString());
-                                }
-                        });
-                        saveMappingsAsMenu.add(item);
+                // Mapping-IO writers
+                for (MappingFormat format : writableMappingIoFormats) {
+                        addSaveMappingsAsMenuEntry(format.getMappingIoCounterpart().name + " (via Mapping-IO, experimental)",
+                                        format, true, saveMappingsAsMenu, saveMappingsItem, gui);
                 }
         }
 
+        private static void addSaveMappingsAsMenuEntry(String text, MappingFormat format, boolean mappingIo, JMenu saveMappingsAsMenu, JMenuItem saveMappingsItem, Gui gui) {
+                JMenuItem item = new JMenuItem(text);
+                item.addActionListener(event -> {
+                        // TODO: Use a specific file chooser for it
+                        if (gui.enigmaMappingsFileChooser.getCurrentDirectory() == null) {
+                                gui.enigmaMappingsFileChooser.setCurrentDirectory(new File(UiConfig.getLastSelectedDir()));
+                        }
+
+                        if (gui.enigmaMappingsFileChooser.showSaveDialog(gui.getFrame()) == JFileChooser.APPROVE_OPTION) {
+                                gui.getController().useMappingIo = mappingIo;
+                                gui.getController().saveMappings(gui.enigmaMappingsFileChooser.getSelectedFile().toPath(), format);
+                                saveMappingsItem.setEnabled(true);
+                                UiConfig.setLastSelectedDir(gui.enigmaMappingsFileChooser.getCurrentDirectory().toString());
+                        }
+                });
+                saveMappingsAsMenu.add(item);
+        }
+
         private static void prepareDecompilerMenu(JMenu decompilerMenu, Gui gui) {
                 ButtonGroup decompilerGroup = new ButtonGroup();
 
diff --git a/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingFormat.java b/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingFormat.java
index 4790fee4..5bad9298 100644
--- a/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingFormat.java
+++ b/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingFormat.java
@@ -27,6 +27,8 @@ public enum MappingFormat {
         TINY_V2(new TinyV2Writer("intermediary", "named"), new TinyV2Reader(), net.fabricmc.mappingio.format.MappingFormat.TINY_2),
         TINY_FILE(TinyMappingsWriter.INSTANCE, TinyMappingsReader.INSTANCE, net.fabricmc.mappingio.format.MappingFormat.TINY),
         SRG_FILE(SrgMappingsWriter.INSTANCE, null, net.fabricmc.mappingio.format.MappingFormat.SRG),
+        TSRG_FILE(null, null, net.fabricmc.mappingio.format.MappingFormat.TSRG),
+        TSRG_2_FILE(null, null, net.fabricmc.mappingio.format.MappingFormat.TSRG2),
         PROGUARD(null, ProguardMappingsReader.INSTANCE, net.fabricmc.mappingio.format.MappingFormat.PROGUARD),
         RECAF(RecafMappingsWriter.INSTANCE, RecafMappingsReader.INSTANCE, null);
 
diff --git a/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingIoConverter.java b/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingIoConverter.java
index 3a864766..d3fd341f 100644
--- a/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingIoConverter.java
+++ b/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingIoConverter.java
@@ -6,11 +6,18 @@ import java.util.List;
 
 import net.fabricmc.mappingio.MappedElementKind;
 import net.fabricmc.mappingio.tree.MemoryMappingTree;
+import net.fabricmc.mappingio.tree.MappingTree.ClassMapping;
+import net.fabricmc.mappingio.tree.MappingTree.FieldMapping;
+import net.fabricmc.mappingio.tree.MappingTree.MethodArgMapping;
+import net.fabricmc.mappingio.tree.MappingTree.MethodMapping;
 
 import cuchaz.enigma.translation.mapping.EntryMap;
 import cuchaz.enigma.translation.mapping.EntryMapping;
 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.MethodDescriptor;
+import cuchaz.enigma.translation.representation.TypeDescriptor;
 import cuchaz.enigma.translation.representation.entry.ClassEntry;
 import cuchaz.enigma.translation.representation.entry.Entry;
 import cuchaz.enigma.translation.representation.entry.FieldEntry;
@@ -123,4 +130,55 @@ public class MappingIoConverter {
                 mappingTree.visitDstName(MappedElementKind.METHOD_ARG, 0, methodArgMapping.targetName());
                 mappingTree.visitComment(MappedElementKind.METHOD_ARG, methodArgMapping.javadoc());
         }
+
+        public static EntryTree<EntryMapping> fromMappingIo(MemoryMappingTree mappingTree) {
+                EntryTree<EntryMapping> dstMappingTree = new HashEntryTree<>();
+
+                for (ClassMapping classMapping : mappingTree.getClasses()) {
+                        readClass(classMapping, dstMappingTree);
+                }
+
+                return dstMappingTree;
+        }
+
+        private static void readClass(ClassMapping classMapping, EntryTree<EntryMapping> mappingTree) {
+                ClassEntry currentClass = new ClassEntry(classMapping.getSrcName());
+                String dstName = classMapping.getDstName(0);
+
+                if (dstName != null) {
+                        dstName = dstName.substring(dstName.lastIndexOf('$') + 1);
+                }
+
+                mappingTree.insert(currentClass, new EntryMapping(dstName));
+
+                for (FieldMapping fieldMapping : classMapping.getFields()) {
+                        readField(fieldMapping, currentClass, mappingTree);
+                }
+
+                for (MethodMapping methodMapping : classMapping.getMethods()) {
+                        readMethod(methodMapping, currentClass, mappingTree);
+                }
+        }
+
+        private static void readField(FieldMapping fieldMapping, ClassEntry parent, EntryTree<EntryMapping> mappingTree) {
+                mappingTree.insert(new FieldEntry(parent, fieldMapping.getSrcName(), new TypeDescriptor(fieldMapping.getSrcDesc())),
+                                new EntryMapping(fieldMapping.getDstName(0)));
+        }
+
+        private static void readMethod(MethodMapping methodMapping, ClassEntry parent, EntryTree<EntryMapping> mappingTree) {
+                MethodEntry currentMethod;
+                mappingTree.insert(currentMethod = new MethodEntry(parent, methodMapping.getSrcName(), new MethodDescriptor(methodMapping.getSrcDesc())),
+                                new EntryMapping(methodMapping.getDstName(0)));
+
+                for (MethodArgMapping methodArgMapping : methodMapping.getArgs()) {
+                        readMethodArg(methodArgMapping, currentMethod, mappingTree);
+                }
+        }
+
+        private static void readMethodArg(MethodArgMapping methodArgMapping, MethodEntry parent, EntryTree<EntryMapping> mappingTree) {
+                String methodArgSrcName = methodArgMapping.getSrcName() != null ? methodArgMapping.getSrcName() : "";
+
+                mappingTree.insert(new LocalVariableEntry(parent, methodArgMapping.getLvIndex(), methodArgSrcName, true, null),
+                                new EntryMapping(methodArgMapping.getDstName(0)));
+        }
 }