more tweaks, improvements, and bug fixes

2 files changed, 77 insertions, 8 deletions

diff --git a/src/cuchaz/enigma/convert/ClassIdentity.java b/src/cuchaz/enigma/convert/ClassIdentity.java
index d07e0a4..35667b0 100644
--- a/src/cuchaz/enigma/convert/ClassIdentity.java
+++ b/src/cuchaz/enigma/convert/ClassIdentity.java
@@ -16,6 +16,7 @@ import java.security.NoSuchAlgorithmException;
 import java.util.Enumeration;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import javassist.CannotCompileException;
 import javassist.CtBehavior;
@@ -38,6 +39,7 @@ import com.google.common.collect.HashMultiset;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Multiset;
+import com.google.common.collect.Sets;
 
 import cuchaz.enigma.Constants;
 import cuchaz.enigma.Util;
@@ -67,6 +69,7 @@ public class ClassIdentity {
         private String m_staticInitializer;
         private String m_extends;
         private Multiset<String> m_implements;
+        private Set<String> m_stringLiterals;
         private Multiset<String> m_implementations;
         private Multiset<String> m_references;
         private String m_outer;
@@ -140,6 +143,14 @@ public class ClassIdentity {
                         m_implements.add(scrubClassName(Descriptor.toJvmName(interfaceName)));
                 }
                 
+                m_stringLiterals = Sets.newHashSet();
+                ConstPool constants = c.getClassFile().getConstPool();
+                for (int i=1; i<constants.getSize(); i++) {
+                        if (constants.getTag(i) == ConstPool.CONST_String) {
+                                m_stringLiterals.add(constants.getStringInfo(i));
+                        }
+                }
+                
                 // stuff from the jar index
                 
                 m_implementations = HashMultiset.create();
@@ -410,13 +421,15 @@ public class ClassIdentity {
         public int getMatchScore(ClassIdentity other) {
                 return 2*getNumMatches(m_extends, other.m_extends)
                         + 2*getNumMatches(m_outer, other.m_outer)
+                        + 2*getNumMatches(m_implements, other.m_implements)
+                         + getNumMatches(m_stringLiterals, other.m_stringLiterals)
                         + getNumMatches(m_fields, other.m_fields)
                         + getNumMatches(m_methods, other.m_methods)
                         + getNumMatches(m_constructors, other.m_constructors);
         }
         
         public int getMaxMatchScore() {
-                return 2 + 2 + m_fields.size() + m_methods.size() + m_constructors.size();
+                return 2 + 2 + 2*m_implements.size() + m_stringLiterals.size() + m_fields.size() + m_methods.size() + m_constructors.size();
         }
         
         public boolean matches(CtClass c) {
@@ -426,6 +439,16 @@ public class ClassIdentity {
                         && m_constructors.size() == c.getDeclaredConstructors().length;
         }
         
+        private int getNumMatches(Set<String> a, Set<String> b) {
+                int numMatches = 0;
+                for (String val : a) {
+                        if (b.contains(val)) {
+                                numMatches++;
+                        }
+                }
+                return numMatches;
+        }
+        
         private int getNumMatches(Multiset<String> a, Multiset<String> b) {
                 int numMatches = 0;
                 for (String val : a) {
diff --git a/src/cuchaz/enigma/convert/MappingsConverter.java b/src/cuchaz/enigma/convert/MappingsConverter.java
index f38723f..5883878 100644
--- a/src/cuchaz/enigma/convert/MappingsConverter.java
+++ b/src/cuchaz/enigma/convert/MappingsConverter.java
@@ -30,7 +30,10 @@ import cuchaz.enigma.analysis.JarIndex;
 import cuchaz.enigma.convert.ClassNamer.SidedClassNamer;
 import cuchaz.enigma.mapping.ClassEntry;
 import cuchaz.enigma.mapping.ClassMapping;
+import cuchaz.enigma.mapping.ClassNameReplacer;
+import cuchaz.enigma.mapping.FieldMapping;
 import cuchaz.enigma.mapping.Mappings;
+import cuchaz.enigma.mapping.MethodMapping;
 
 public class MappingsConverter {
         
@@ -129,15 +132,20 @@ public class MappingsConverter {
                         for (Entry<ClassEntry,ClassEntry> match : matchesByDestChainSize.get(chainSize)) {
                                 
                                 // get class info
-                                ClassEntry sourceClassEntry = match.getKey();
-                                ClassEntry deobfClassEntry = sourceDeobfuscator.deobfuscateEntry(sourceClassEntry);
-                                ClassEntry destClassEntry = match.getValue();
-                                List<ClassEntry> destClassChain = destDeobfuscator.getJarIndex().getObfClassChain(destClassEntry);
+                                ClassEntry obfSourceClassEntry = match.getKey();
+                                ClassEntry obfDestClassEntry = match.getValue();
+                                List<ClassEntry> destClassChain = destDeobfuscator.getJarIndex().getObfClassChain(obfDestClassEntry);
+                                
+                                ClassMapping sourceMapping = sourceDeobfuscator.getMappings().getClassByObf(obfSourceClassEntry);
+                                if (sourceMapping == null) {
+                                        // if this class was never deobfuscated, don't try to match it
+                                        continue;
+                                }
                                 
                                 // find out where to make the dest class mapping
                                 if (destClassChain.size() == 1) {
                                         // not an inner class, add directly to mappings
-                                        newMappings.addClassMapping(new ClassMapping(destClassEntry.getName(), deobfClassEntry.getName()));
+                                        newMappings.addClassMapping(migrateClassMapping(obfDestClassEntry, sourceMapping, matches, false));
                                 } else {
                                         // inner class, find the outer class mapping
                                         ClassMapping destMapping = null;
@@ -157,14 +165,52 @@ public class MappingsConverter {
                                                         }
                                                 }
                                         }
-                                        String deobfName = deobfClassEntry.isInnerClass() ? deobfClassEntry.getInnerClassName() : deobfClassEntry.getSimpleName();
-                                        destMapping.addInnerClassMapping(new ClassMapping(destClassEntry.getName(), deobfName));
+                                        destMapping.addInnerClassMapping(migrateClassMapping(obfDestClassEntry, sourceMapping, matches, true));
                                 }
                         }
                 }
                 return newMappings;
         }
         
+        private static ClassMapping migrateClassMapping(ClassEntry newObfClass, ClassMapping mapping, final Matches matches, boolean useSimpleName) {
+                
+                ClassNameReplacer replacer = new ClassNameReplacer() {
+                        @Override
+                        public String replace(String className) {
+                                ClassEntry newClassEntry = matches.getUniqueMatches().get(new ClassEntry(className));
+                                if (newClassEntry != null) {
+                                        return newClassEntry.getName();
+                                }
+                                return null;
+                        }
+                };
+                
+                ClassMapping newMapping;
+                String deobfName = mapping.getDeobfName();
+                if (deobfName != null) {
+                        if (useSimpleName) {
+                                deobfName = new ClassEntry(deobfName).getSimpleName();
+                        }
+                        newMapping = new ClassMapping(newObfClass.getName(), deobfName);
+                } else {
+                        newMapping = new ClassMapping(newObfClass.getName());
+                }
+                
+                // copy fields
+                for (FieldMapping fieldMapping : mapping.fields()) {
+                        // TODO: map field obf names too...
+                        newMapping.addFieldMapping(new FieldMapping(fieldMapping, replacer));
+                }
+                
+                // copy methods
+                for (MethodMapping methodMapping : mapping.methods()) {
+                        // TODO: map method obf names too...
+                        newMapping.addMethodMapping(new MethodMapping(methodMapping, replacer));
+                }
+                
+                return newMapping;
+        }
+
         public static void convertMappings(Mappings mappings, BiMap<ClassEntry,ClassEntry> changes) {
                 
                 // sort the changes so classes are renamed in the correct order