1 files changed, 201 insertions, 0 deletions

diff --git a/src/main/java/cuchaz/enigma/mapping/Mappings.java b/src/main/java/cuchaz/enigma/mapping/Mappings.java
new file mode 100644
index 0000000..a48ec3f
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/mapping/Mappings.java
@@ -0,0 +1,201 @@
+/*******************************************************************************
+ * Copyright (c) 2015 Jeff Martin.
+ * All rights reserved. This program and the accompanying materials
+ * are made available under the terms of the GNU Lesser General Public
+ * License v3.0 which accompanies this distribution, and is available at
+ * http://www.gnu.org/licenses/lgpl.html
+ * <p>
+ * Contributors:
+ * Jeff Martin - initial API and implementation
+ ******************************************************************************/
+package cuchaz.enigma.mapping;
+
+import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
+import com.google.common.collect.Sets;
+
+import java.io.Serializable;
+import java.util.*;
+
+import cuchaz.enigma.analysis.TranslationIndex;
+
+public class Mappings implements Serializable {
+
+    private static final long serialVersionUID = 4649790259460259026L;
+
+    protected Map<String, ClassMapping> m_classesByObf;
+    protected Map<String, ClassMapping> m_classesByDeobf;
+
+    public Mappings() {
+        m_classesByObf = Maps.newHashMap();
+        m_classesByDeobf = Maps.newHashMap();
+    }
+
+    public Mappings(Iterable<ClassMapping> classes) {
+        this();
+
+        for (ClassMapping classMapping : classes) {
+            m_classesByObf.put(classMapping.getObfFullName(), classMapping);
+            if (classMapping.getDeobfName() != null) {
+                m_classesByDeobf.put(classMapping.getDeobfName(), classMapping);
+            }
+        }
+    }
+
+    public Collection<ClassMapping> classes() {
+        assert (m_classesByObf.size() >= m_classesByDeobf.size());
+        return m_classesByObf.values();
+    }
+
+    public void addClassMapping(ClassMapping classMapping) {
+        if (m_classesByObf.containsKey(classMapping.getObfFullName())) {
+            throw new Error("Already have mapping for " + classMapping.getObfFullName());
+        }
+        boolean obfWasAdded = m_classesByObf.put(classMapping.getObfFullName(), classMapping) == null;
+        assert (obfWasAdded);
+        if (classMapping.getDeobfName() != null) {
+            if (m_classesByDeobf.containsKey(classMapping.getDeobfName())) {
+                throw new Error("Already have mapping for " + classMapping.getDeobfName());
+            }
+            boolean deobfWasAdded = m_classesByDeobf.put(classMapping.getDeobfName(), classMapping) == null;
+            assert (deobfWasAdded);
+        }
+    }
+
+    public void removeClassMapping(ClassMapping classMapping) {
+        boolean obfWasRemoved = m_classesByObf.remove(classMapping.getObfFullName()) != null;
+        assert (obfWasRemoved);
+        if (classMapping.getDeobfName() != null) {
+            boolean deobfWasRemoved = m_classesByDeobf.remove(classMapping.getDeobfName()) != null;
+            assert (deobfWasRemoved);
+        }
+    }
+
+    public ClassMapping getClassByObf(ClassEntry entry) {
+        return getClassByObf(entry.getName());
+    }
+
+    public ClassMapping getClassByObf(String obfName) {
+        return m_classesByObf.get(obfName);
+    }
+
+    public ClassMapping getClassByDeobf(ClassEntry entry) {
+        return getClassByDeobf(entry.getName());
+    }
+
+    public ClassMapping getClassByDeobf(String deobfName) {
+        return m_classesByDeobf.get(deobfName);
+    }
+
+    public void setClassDeobfName(ClassMapping classMapping, String deobfName) {
+        if (classMapping.getDeobfName() != null) {
+            boolean wasRemoved = m_classesByDeobf.remove(classMapping.getDeobfName()) != null;
+            assert (wasRemoved);
+        }
+        classMapping.setDeobfName(deobfName);
+        if (deobfName != null) {
+            boolean wasAdded = m_classesByDeobf.put(deobfName, classMapping) == null;
+            assert (wasAdded);
+        }
+    }
+
+    public Translator getTranslator(TranslationDirection direction, TranslationIndex index) {
+        switch (direction) {
+            case Deobfuscating:
+
+                return new Translator(direction, m_classesByObf, index);
+
+            case Obfuscating:
+
+                // fill in the missing deobf class entries with obf entries
+                Map<String, ClassMapping> classes = Maps.newHashMap();
+                for (ClassMapping classMapping : classes()) {
+                    if (classMapping.getDeobfName() != null) {
+                        classes.put(classMapping.getDeobfName(), classMapping);
+                    } else {
+                        classes.put(classMapping.getObfFullName(), classMapping);
+                    }
+                }
+
+                // translate the translation index
+                // NOTE: this isn't actually recursive
+                TranslationIndex deobfIndex = new TranslationIndex(index, getTranslator(TranslationDirection.Deobfuscating, index));
+
+                return new Translator(direction, classes, deobfIndex);
+
+            default:
+                throw new Error("Invalid translation direction!");
+        }
+    }
+
+    @Override
+    public String toString() {
+        StringBuilder buf = new StringBuilder();
+        for (ClassMapping classMapping : m_classesByObf.values()) {
+            buf.append(classMapping.toString());
+            buf.append("\n");
+        }
+        return buf.toString();
+    }
+
+    public void renameObfClass(String oldObfName, String newObfName) {
+        new ArrayList<>(classes()).stream().filter(classMapping -> classMapping.renameObfClass(oldObfName, newObfName)).forEach(classMapping -> {
+            boolean wasRemoved = m_classesByObf.remove(oldObfName) != null;
+            assert (wasRemoved);
+            boolean wasAdded = m_classesByObf.put(newObfName, classMapping) == null;
+            assert (wasAdded);
+        });
+    }
+
+    public Set<String> getAllObfClassNames() {
+        final Set<String> classNames = Sets.newHashSet();
+        for (ClassMapping classMapping : classes()) {
+
+            // add the class name
+            classNames.add(classMapping.getObfFullName());
+
+            // add classes from method signatures
+            for (MethodMapping methodMapping : classMapping.methods()) {
+                for (Type type : methodMapping.getObfSignature().types()) {
+                    if (type.hasClass()) {
+                        classNames.add(type.getClassEntry().getClassName());
+                    }
+                }
+            }
+        }
+        return classNames;
+    }
+
+    public boolean containsDeobfClass(String deobfName) {
+        return m_classesByDeobf.containsKey(deobfName);
+    }
+
+    public boolean containsDeobfField(ClassEntry obfClassEntry, String deobfName, Type obfType) {
+        ClassMapping classMapping = m_classesByObf.get(obfClassEntry.getName());
+        return classMapping != null && classMapping.containsDeobfField(deobfName, obfType);
+    }
+
+    public boolean containsDeobfMethod(ClassEntry obfClassEntry, String deobfName, Signature deobfSignature) {
+        ClassMapping classMapping = m_classesByObf.get(obfClassEntry.getName());
+        return classMapping != null && classMapping.containsDeobfMethod(deobfName, deobfSignature);
+    }
+
+    public boolean containsArgument(BehaviorEntry obfBehaviorEntry, String name) {
+        ClassMapping classMapping = m_classesByObf.get(obfBehaviorEntry.getClassName());
+        return classMapping != null && classMapping.containsArgument(obfBehaviorEntry, name);
+    }
+
+    public List<ClassMapping> getClassMappingChain(ClassEntry obfClass) {
+        List<ClassMapping> mappingChain = Lists.newArrayList();
+        ClassMapping classMapping = null;
+        for (ClassEntry obfClassEntry : obfClass.getClassChain()) {
+            if (mappingChain.isEmpty()) {
+                classMapping = m_classesByObf.get(obfClassEntry.getName());
+            } else if (classMapping != null) {
+                classMapping = classMapping.getInnerClassByObfSimple(obfClassEntry.getInnermostClassName());
+            }
+            mappingChain.add(classMapping);
+        }
+        return mappingChain;
+    }
+}