From 6c7738ca744e55e8e2ea03518a5b93000fe9a864 Mon Sep 17 00:00:00 2001
From: NebelNidas
Date: Thu, 6 Oct 2022 20:03:14 +0200
Subject: Add Mapping-IO import support; small cleanup

---
 .../translation/mapping/serde/MappingFormat.java   |  2 +
 .../mapping/serde/MappingIoConverter.java          | 58 ++++++++++++++++++++++
 2 files changed, 60 insertions(+)

(limited to 'enigma/src')

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)));
+	}
 }
-- 
cgit v1.2.3