3 files changed, 103 insertions, 20 deletions

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 f402c19..4d34720 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
@@ -1,10 +1,19 @@
 package cuchaz.enigma.translation.mapping.serde;
 
 import java.io.IOException;
+import java.io.UncheckedIOException;
 import java.nio.file.Path;
+import java.util.Arrays;
+import java.util.List;
 
 import javax.annotation.Nullable;
 
+import net.fabricmc.mappingio.MappingReader;
+import net.fabricmc.mappingio.MappingWriter;
+import net.fabricmc.mappingio.tree.MemoryMappingTree;
+import net.fabricmc.mappingio.tree.VisitOrder;
+import net.fabricmc.mappingio.tree.VisitableMappingTree;
+
 import cuchaz.enigma.ProgressListener;
 import cuchaz.enigma.translation.mapping.EntryMapping;
 import cuchaz.enigma.translation.mapping.MappingDelta;
@@ -19,27 +28,30 @@ import cuchaz.enigma.translation.mapping.serde.tiny.TinyMappingsWriter;
 import cuchaz.enigma.translation.mapping.serde.tinyv2.TinyV2Reader;
 import cuchaz.enigma.translation.mapping.serde.tinyv2.TinyV2Writer;
 import cuchaz.enigma.translation.mapping.tree.EntryTree;
+import cuchaz.enigma.utils.I18n;
 
 public enum MappingFormat {
-        ENIGMA_FILE(EnigmaMappingsWriter.FILE, EnigmaMappingsReader.FILE, net.fabricmc.mappingio.format.MappingFormat.ENIGMA_FILE),
-        ENIGMA_DIRECTORY(EnigmaMappingsWriter.DIRECTORY, EnigmaMappingsReader.DIRECTORY, net.fabricmc.mappingio.format.MappingFormat.ENIGMA_DIR),
-        ENIGMA_ZIP(EnigmaMappingsWriter.ZIP, EnigmaMappingsReader.ZIP, null),
-        TINY_V2(new TinyV2Writer("intermediary", "named"), new TinyV2Reader(), net.fabricmc.mappingio.format.MappingFormat.TINY_2_FILE),
-        TINY_FILE(TinyMappingsWriter.INSTANCE, TinyMappingsReader.INSTANCE, net.fabricmc.mappingio.format.MappingFormat.TINY_FILE),
-        SRG_FILE(SrgMappingsWriter.INSTANCE, null, net.fabricmc.mappingio.format.MappingFormat.SRG_FILE),
-        TSRG_FILE(null, null, net.fabricmc.mappingio.format.MappingFormat.TSRG_FILE),
-        TSRG_2_FILE(null, null, net.fabricmc.mappingio.format.MappingFormat.TSRG_2_FILE),
-        PROGUARD(null, ProguardMappingsReader.INSTANCE, net.fabricmc.mappingio.format.MappingFormat.PROGUARD_FILE),
-        RECAF(RecafMappingsWriter.INSTANCE, RecafMappingsReader.INSTANCE, null);
+        ENIGMA_FILE(EnigmaMappingsWriter.FILE, EnigmaMappingsReader.FILE, net.fabricmc.mappingio.format.MappingFormat.ENIGMA_FILE, true),
+        ENIGMA_DIRECTORY(EnigmaMappingsWriter.DIRECTORY, EnigmaMappingsReader.DIRECTORY, net.fabricmc.mappingio.format.MappingFormat.ENIGMA_DIR, true),
+        ENIGMA_ZIP(EnigmaMappingsWriter.ZIP, EnigmaMappingsReader.ZIP, null, false),
+        TINY_V2(new TinyV2Writer("intermediary", "named"), new TinyV2Reader(), net.fabricmc.mappingio.format.MappingFormat.TINY_2_FILE, true),
+        TINY_FILE(TinyMappingsWriter.INSTANCE, TinyMappingsReader.INSTANCE, net.fabricmc.mappingio.format.MappingFormat.TINY_FILE, true),
+        SRG_FILE(SrgMappingsWriter.INSTANCE, null, net.fabricmc.mappingio.format.MappingFormat.SRG_FILE, false),
+        TSRG_FILE(null, null, net.fabricmc.mappingio.format.MappingFormat.TSRG_FILE, false),
+        TSRG_2_FILE(null, null, net.fabricmc.mappingio.format.MappingFormat.TSRG_2_FILE, false),
+        PROGUARD(null, ProguardMappingsReader.INSTANCE, net.fabricmc.mappingio.format.MappingFormat.PROGUARD_FILE, true),
+        RECAF(RecafMappingsWriter.INSTANCE, RecafMappingsReader.INSTANCE, null, false);
 
         private final MappingsWriter writer;
         private final MappingsReader reader;
         private final net.fabricmc.mappingio.format.MappingFormat mappingIoCounterpart;
+        private final boolean hasMappingIoWriter;
 
-        MappingFormat(MappingsWriter writer, MappingsReader reader, net.fabricmc.mappingio.format.MappingFormat mappingIoCounterpart) {
+        MappingFormat(MappingsWriter writer, MappingsReader reader, net.fabricmc.mappingio.format.MappingFormat mappingIoCounterpart, boolean hasMappingIoWriter) {
                 this.writer = writer;
                 this.reader = reader;
                 this.mappingIoCounterpart = mappingIoCounterpart;
+                this.hasMappingIoWriter = hasMappingIoWriter;
         }
 
         public void write(EntryTree<EntryMapping> mappings, Path path, ProgressListener progressListener, MappingSaveParameters saveParameters) {
@@ -47,19 +59,49 @@ public enum MappingFormat {
         }
 
         public void write(EntryTree<EntryMapping> mappings, MappingDelta<EntryMapping> delta, Path path, ProgressListener progressListener, MappingSaveParameters saveParameters) {
-                if (writer == null) {
-                        throw new IllegalStateException(name() + " does not support writing");
+                if (!hasMappingIoWriter || !useMappingIo()) {
+                        if (writer == null) {
+                                throw new IllegalStateException(name() + " does not support writing");
+                        }
+
+                        writer.write(mappings, delta, path, progressListener, saveParameters);
+                        return;
                 }
 
-                writer.write(mappings, delta, path, progressListener, saveParameters);
+                try {
+                        VisitableMappingTree tree = MappingIoConverter.toMappingIo(mappings, progressListener);
+                        progressListener.init(1, I18n.translate("progress.mappings.writing"));
+                        progressListener.step(1, null); // Reset message
+
+                        tree.accept(MappingWriter.create(path, mappingIoCounterpart), VisitOrder.createByName());
+                        progressListener.step(1, I18n.translate("progress.done"));
+                } catch (IOException e) {
+                        throw new UncheckedIOException(e);
+                }
         }
 
         public EntryTree<EntryMapping> read(Path path, ProgressListener progressListener, MappingSaveParameters saveParameters) throws IOException, MappingParseException {
-                if (reader == null) {
-                        throw new IllegalStateException(name() + " does not support reading");
+                if (!useMappingIo()) {
+                        if (reader == null) {
+                                throw new IllegalStateException(name() + " does not support reading");
+                        }
+
+                        return reader.read(path, progressListener, saveParameters);
                 }
 
-                return reader.read(path, progressListener, saveParameters);
+                String loadingMessage;
+
+                if (mappingIoCounterpart.hasSingleFile()) {
+                        loadingMessage = I18n.translate("progress.mappings.loading_file");
+                } else {
+                        loadingMessage = I18n.translate("progress.mappings.loading_directory");
+                }
+
+                progressListener.init(1, loadingMessage);
+
+                VisitableMappingTree mappingTree = new MemoryMappingTree();
+                MappingReader.read(path, mappingIoCounterpart, mappingTree);
+                return MappingIoConverter.fromMappingIo(mappingTree, progressListener);
         }
 
         @Nullable
@@ -76,4 +118,35 @@ public enum MappingFormat {
         public net.fabricmc.mappingio.format.MappingFormat getMappingIoCounterpart() {
                 return mappingIoCounterpart;
         }
+
+        public boolean hasMappingIoWriter() {
+                return hasMappingIoWriter;
+        }
+
+        public boolean isReadable() {
+                return reader != null || mappingIoCounterpart != null;
+        }
+
+        public boolean isWritable() {
+                return writer != null || hasMappingIoWriter;
+        }
+
+        private boolean useMappingIo() {
+                if (mappingIoCounterpart == null) return false;
+                return System.getProperty("enigma.use_mappingio", "true").equals("true");
+        }
+
+        public static List<MappingFormat> getReadableFormats() {
+                return Arrays.asList(values())
+                                .stream()
+                                .filter(MappingFormat::isReadable)
+                                .toList();
+        }
+
+        public static List<MappingFormat> getWritableFormats() {
+                return Arrays.asList(values())
+                                .stream()
+                                .filter(MappingFormat::isWritable)
+                                .toList();
+        }
 }
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 a0912fa..4372585 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
@@ -33,19 +33,23 @@ import cuchaz.enigma.utils.I18n;
 
 public class MappingIoConverter {
         public static VisitableMappingTree toMappingIo(EntryTree<EntryMapping> mappings, ProgressListener progress) {
+                return toMappingIo(mappings, progress, "intermediary", "named");
+        }
+
+        public static VisitableMappingTree toMappingIo(EntryTree<EntryMapping> mappings, ProgressListener progress, String fromNs, String toNs) {
                 try {
                         List<EntryTreeNode<EntryMapping>> classes = StreamSupport.stream(mappings.spliterator(), false)
                                         .filter(node -> node.getEntry() instanceof ClassEntry)
                                         .toList();
 
                         progress.init(classes.size(), I18n.translate("progress.mappings.converting.to_mappingio"));
-                        int steps = 0;
+                        int stepsDone = 0;
 
                         MemoryMappingTree mappingTree = new MemoryMappingTree();
-                        mappingTree.visitNamespaces("intermediary", List.of("named"));
+                        mappingTree.visitNamespaces(fromNs, List.of(toNs));
 
                         for (EntryTreeNode<EntryMapping> classNode : classes) {
-                                progress.step(steps++, classNode.getEntry().getFullName());
+                                progress.step(++stepsDone, classNode.getEntry().getFullName());
                                 writeClass(classNode, mappings, mappingTree);
                         }
 
diff --git a/enigma/src/test/java/cuchaz/enigma/translation/mapping/TestReadWriteCycle.java b/enigma/src/test/java/cuchaz/enigma/translation/mapping/TestReadWriteCycle.java
index 681fd3f..34d5a6e 100644
--- a/enigma/src/test/java/cuchaz/enigma/translation/mapping/TestReadWriteCycle.java
+++ b/enigma/src/test/java/cuchaz/enigma/translation/mapping/TestReadWriteCycle.java
@@ -4,6 +4,7 @@ import java.io.File;
 import java.io.IOException;
 
 import org.junit.Assert;
+import org.junit.BeforeClass;
 import org.junit.Test;
 
 import cuchaz.enigma.ProgressListener;
@@ -87,6 +88,11 @@ public class TestReadWriteCycle {
                 tempFile.delete();
         }
 
+        @BeforeClass
+        public static void setup() {
+                System.getProperties().setProperty("enigma.use_mappingio", "false");
+        }
+
         @Test
         public void testEnigmaFile() throws IOException, MappingParseException {
                 testReadWriteCycle(MappingFormat.ENIGMA_FILE, true, ".enigma");