Add Mapping-IO import support; small cleanup

2 files changed, 60 insertions, 0 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 4790fee..5bad929 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 3a86476..d3fd341 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)));
+        }
 }