fix method/argument renaming edge cases, add method/argument name rebuilder, add Package to Access enum

8 files changed, 174 insertions, 19 deletions

diff --git a/src/main/java/cuchaz/enigma/Deobfuscator.java b/src/main/java/cuchaz/enigma/Deobfuscator.java
index b3fb3158..2207999d 100644
--- a/src/main/java/cuchaz/enigma/Deobfuscator.java
+++ b/src/main/java/cuchaz/enigma/Deobfuscator.java
@@ -11,6 +11,7 @@
 package cuchaz.enigma;
 
 import com.google.common.base.Charsets;
+import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
 import com.strobel.assembler.metadata.MetadataSystem;
@@ -32,9 +33,7 @@ import javassist.CtClass;
 import javassist.bytecode.Descriptor;
 
 import java.io.*;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 import java.util.jar.JarEntry;
 import java.util.jar.JarFile;
 import java.util.jar.JarOutputStream;
@@ -297,6 +296,92 @@ public class Deobfuscator {
         }
     }
 
+    private void addAllPotentialAncestors(Set<ClassEntry> classEntries, ClassEntry classObfEntry) {
+        for (ClassEntry interfaceEntry : jarIndex.getTranslationIndex().getInterfaces(classObfEntry)) {
+            if (classEntries.add(interfaceEntry)) {
+                addAllPotentialAncestors(classEntries, interfaceEntry);
+            }
+        }
+
+        ClassEntry superClassEntry = jarIndex.getTranslationIndex().getSuperclass(classObfEntry);
+        if (superClassEntry != null && classEntries.add(superClassEntry)) {
+            addAllPotentialAncestors(classEntries, superClassEntry);
+        }
+    }
+
+    private boolean isBehaviorProvider(ClassEntry classObfEntry, BehaviorEntry behaviorEntry) {
+        if (behaviorEntry instanceof MethodEntry) {
+            MethodEntry methodEntry = (MethodEntry) behaviorEntry;
+
+            Set<ClassEntry> classEntries = new HashSet<>();
+            addAllPotentialAncestors(classEntries, classObfEntry);
+
+            for (ClassEntry parentEntry : classEntries) {
+                MethodEntry ancestorMethodEntry = new MethodEntry(parentEntry, methodEntry.getName(), methodEntry.getSignature());
+                if (jarIndex.containsObfBehavior(ancestorMethodEntry)) {
+                    return false;
+                }
+            }
+        }
+
+        return true;
+    }
+
+    public void rebuildMethodNames(ProgressListener progress) {
+        int i = 0;
+        Map<ClassMapping, Map<Entry, String>> renameClassMap = new HashMap<>();
+
+        progress.init(getMappings().classes().size() * 3, "Rebuilding method names");
+
+        for (ClassMapping classMapping : Lists.newArrayList(getMappings().classes())) {
+            Map<Entry, String> renameEntries = new HashMap<>();
+
+            progress.onProgress(i++, classMapping.getDeobfName());
+
+            for (MethodMapping methodMapping : Lists.newArrayList(classMapping.methods())) {
+                ClassEntry classObfEntry = classMapping.getObfEntry();
+                BehaviorEntry obfEntry = methodMapping.getObfEntry(classObfEntry);
+
+                if (isBehaviorProvider(classObfEntry, obfEntry)) {
+                    if (hasDeobfuscatedName(obfEntry) && !(obfEntry instanceof ConstructorEntry)
+                            && !(methodMapping.getDeobfName().equals(methodMapping.getObfName()))) {
+                        renameEntries.put(obfEntry, methodMapping.getDeobfName());
+                    }
+
+                    for (ArgumentMapping argumentMapping : Lists.newArrayList(methodMapping.arguments())) {
+                        Entry argObfEntry = argumentMapping.getObfEntry(obfEntry);
+                        if (hasDeobfuscatedName(argObfEntry)) {
+                            renameEntries.put(argObfEntry, deobfuscateEntry(argObfEntry).getName());
+                        }
+                    }
+                }
+            }
+
+            renameClassMap.put(classMapping, renameEntries);
+        }
+
+        for (Map.Entry<ClassMapping, Map<Entry, String>> renameClassMapEntry : renameClassMap.entrySet()) {
+            progress.onProgress(i++, renameClassMapEntry.getKey().getDeobfName());
+
+            for (Map.Entry<Entry, String> entry : renameClassMapEntry.getValue().entrySet()) {
+                Entry obfEntry = entry.getKey();
+
+                removeMapping(obfEntry);
+            }
+        }
+
+        for (Map.Entry<ClassMapping, Map<Entry, String>> renameClassMapEntry : renameClassMap.entrySet()) {
+            progress.onProgress(i++, renameClassMapEntry.getKey().getDeobfName());
+
+            for (Map.Entry<Entry, String> entry : renameClassMapEntry.getValue().entrySet()) {
+                Entry obfEntry = entry.getKey();
+                String name = entry.getValue();
+
+                rename(obfEntry, name);
+            }
+        }
+    }
+
     public void writeJar(File out, ProgressListener progress) {
         final TranslatingTypeLoader loader = new TranslatingTypeLoader(
                 this.jar,
diff --git a/src/main/java/cuchaz/enigma/analysis/Access.java b/src/main/java/cuchaz/enigma/analysis/Access.java
index 877327f1..b8aafaa0 100644
--- a/src/main/java/cuchaz/enigma/analysis/Access.java
+++ b/src/main/java/cuchaz/enigma/analysis/Access.java
@@ -19,6 +19,7 @@ public enum Access {
 
     Public,
     Protected,
+    Package,
     Private;
 
     public static Access get(CtBehavior behavior) {
@@ -30,12 +31,18 @@ public enum Access {
     }
 
     public static Access get(int modifiers) {
-        if (Modifier.isPublic(modifiers)) {
+        boolean isPublic = Modifier.isPublic(modifiers);
+        boolean isProtected = Modifier.isProtected(modifiers);
+        boolean isPrivate = Modifier.isPrivate(modifiers);
+
+        if (isPublic && !isProtected && !isPrivate) {
             return Public;
-        } else if (Modifier.isProtected(modifiers)) {
+        } else if (!isPublic && isProtected && !isPrivate) {
             return Protected;
-        } else if (Modifier.isPrivate(modifiers)) {
+        } else if (!isPublic && !isProtected && isPrivate) {
             return Private;
+        } else if (!isPublic && !isProtected && !isPrivate) {
+            return Package;
         }
         // assume public by default
         return Public;
diff --git a/src/main/java/cuchaz/enigma/analysis/JarIndex.java b/src/main/java/cuchaz/enigma/analysis/JarIndex.java
index 851f3faa..e501540d 100644
--- a/src/main/java/cuchaz/enigma/analysis/JarIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/JarIndex.java
@@ -38,6 +38,7 @@ public class JarIndex {
     private Map<ClassEntry, ClassEntry> outerClassesByInner;
     private Map<ClassEntry, BehaviorEntry> anonymousClasses;
     private Map<MethodEntry, MethodEntry> bridgedMethods;
+    private Set<MethodEntry> syntheticMethods;
 
     public JarIndex() {
         this.obfClassEntries = Sets.newHashSet();
@@ -52,6 +53,7 @@ public class JarIndex {
         this.outerClassesByInner = Maps.newHashMap();
         this.anonymousClasses = Maps.newHashMap();
         this.bridgedMethods = Maps.newHashMap();
+        this.syntheticMethods = Sets.newHashSet();
     }
 
     public void indexJar(JarFile jar, boolean buildInnerClasses) {
@@ -155,6 +157,11 @@ public class JarIndex {
         if (behaviorEntry instanceof MethodEntry) {
             MethodEntry methodEntry = (MethodEntry) behaviorEntry;
 
+            // is synthetic
+            if ((behavior.getModifiers() & AccessFlag.SYNTHETIC) != 0) {
+                syntheticMethods.add(methodEntry);
+            }
+
             // index implementation
             this.methodImplementations.put(behaviorEntry.getClassName(), methodEntry);
 
@@ -720,6 +727,10 @@ public class JarIndex {
         return this.anonymousClasses.containsKey(obfInnerClassEntry);
     }
 
+    public boolean isSyntheticMethod(MethodEntry methodEntry) {
+        return this.syntheticMethods.contains(methodEntry);
+    }
+
     public BehaviorEntry getAnonymousClassCaller(ClassEntry obfInnerClassName) {
         return this.anonymousClasses.get(obfInnerClassName);
     }
diff --git a/src/main/java/cuchaz/enigma/gui/GuiController.java b/src/main/java/cuchaz/enigma/gui/GuiController.java
index 1deb7d2a..3188ff02 100644
--- a/src/main/java/cuchaz/enigma/gui/GuiController.java
+++ b/src/main/java/cuchaz/enigma/gui/GuiController.java
@@ -100,6 +100,10 @@ public class GuiController {
         refreshCurrentClass();
     }
 
+    public void rebuildMethodNames() {
+        ProgressDialog.runInThread(this.gui.getFrame(), progress -> this.deobfuscator.rebuildMethodNames(progress));
+    }
+
     public void exportSource(final File dirOut) {
         ProgressDialog.runInThread(this.gui.getFrame(), progress -> this.deobfuscator.writeSources(dirOut, progress));
     }
diff --git a/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java b/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java
index befe1291..dcd7c93a 100644
--- a/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java
+++ b/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java
@@ -25,6 +25,7 @@ public class MenuBar extends JMenuBar {
     public final JMenuItem saveMappingsSrgMenu;
     public final JMenuItem closeMappingsMenu;
 
+    public final JMenuItem rebuildMethodNamesMenu;
 
     public final JMenuItem exportSourceMenu;
     public final JMenuItem exportJarMenu;
@@ -174,6 +175,15 @@ public class MenuBar extends JMenuBar {
             }
             menu.addSeparator();
             {
+                JMenuItem item = new JMenuItem("Rebuild Method Names");
+                menu.add(item);
+                item.addActionListener(event -> {
+                    this.gui.getController().rebuildMethodNames();
+                });
+                this.rebuildMethodNamesMenu = item;
+            }
+            menu.addSeparator();
+            {
                 JMenuItem item = new JMenuItem("Export Source...");
                 menu.add(item);
                 item.addActionListener(event -> {
diff --git a/src/main/java/cuchaz/enigma/mapping/ArgumentMapping.java b/src/main/java/cuchaz/enigma/mapping/ArgumentMapping.java
index d117de02..e3f89272 100644
--- a/src/main/java/cuchaz/enigma/mapping/ArgumentMapping.java
+++ b/src/main/java/cuchaz/enigma/mapping/ArgumentMapping.java
@@ -38,6 +38,10 @@ public class ArgumentMapping implements Comparable<ArgumentMapping> {
         this.name = NameValidator.validateArgumentName(val);
     }
 
+    public ArgumentEntry getObfEntry(BehaviorEntry behaviorEntry) {
+        return new ArgumentEntry(behaviorEntry, index, name);
+    }
+
     @Override
     public int compareTo(ArgumentMapping other) {
         return Integer.compare(this.index, other.index);
diff --git a/src/main/java/cuchaz/enigma/mapping/Mappings.java b/src/main/java/cuchaz/enigma/mapping/Mappings.java
index 47c42324..7061be79 100644
--- a/src/main/java/cuchaz/enigma/mapping/Mappings.java
+++ b/src/main/java/cuchaz/enigma/mapping/Mappings.java
@@ -175,9 +175,9 @@ public class Mappings {
         return classMapping != null && classMapping.containsDeobfField(deobfName, obfType);
     }
 
-    public boolean containsDeobfMethod(ClassEntry obfClassEntry, String deobfName, Signature deobfSignature) {
+    public boolean containsDeobfMethod(ClassEntry obfClassEntry, String deobfName, Signature obfSignature) {
         ClassMapping classMapping = this.classesByObf.get(obfClassEntry.getName());
-        return classMapping != null && classMapping.containsDeobfMethod(deobfName, deobfSignature);
+        return classMapping != null && classMapping.containsDeobfMethod(deobfName, obfSignature);
     }
 
     public boolean containsArgument(BehaviorEntry obfBehaviorEntry, String name) {
diff --git a/src/main/java/cuchaz/enigma/mapping/MappingsRenamer.java b/src/main/java/cuchaz/enigma/mapping/MappingsRenamer.java
index f244748f..7c6c8318 100644
--- a/src/main/java/cuchaz/enigma/mapping/MappingsRenamer.java
+++ b/src/main/java/cuchaz/enigma/mapping/MappingsRenamer.java
@@ -17,6 +17,7 @@ import java.util.List;
 import java.util.Set;
 import java.util.zip.GZIPOutputStream;
 
+import com.google.common.collect.Lists;
 import cuchaz.enigma.analysis.JarIndex;
 import cuchaz.enigma.throwables.IllegalNameException;
 import cuchaz.enigma.throwables.MappingConflict;
@@ -110,10 +111,16 @@ public class MappingsRenamer {
 
         deobfName = NameValidator.validateMethodName(deobfName);
         for (MethodEntry entry : implementations) {
-            Signature deobfSignature = m_mappings.getTranslator(TranslationDirection.Deobfuscating, m_index.getTranslationIndex()).translateSignature(obf.getSignature());
-            MethodEntry targetEntry = new MethodEntry(entry.getClassEntry(), deobfName, deobfSignature);
-            if (m_mappings.containsDeobfMethod(entry.getClassEntry(), deobfName, entry.getSignature()) || m_index.containsObfBehavior(targetEntry)) {
+            MethodEntry targetEntry = new MethodEntry(entry.getClassEntry(), deobfName, entry.getSignature());
+
+            // TODO: Verify if I don't break things
+            ClassMapping classMapping = m_mappings.getClassByObf(entry.getClassEntry());
+            if ((classMapping != null && classMapping.containsDeobfMethod(deobfName, entry.getSignature()) && classMapping.getMethodByObf(entry.getName(), entry.getSignature()) != classMapping.getMethodByDeobf(deobfName, entry.getSignature()))
+                    || m_index.containsObfBehavior(targetEntry)) {
                 String deobfClassName = m_mappings.getTranslator(TranslationDirection.Deobfuscating, m_index.getTranslationIndex()).translateClass(entry.getClassName());
+                if (deobfClassName == null) {
+                    deobfClassName = entry.getClassName();
+                }
                 throw new IllegalNameException(deobfName, "There is already a method with that name and signature in class " + deobfClassName);
             }
         }
@@ -126,12 +133,18 @@ public class MappingsRenamer {
     public void setMethodName(MethodEntry obf, String deobfName) {
         deobfName = NameValidator.validateMethodName(deobfName);
         MethodEntry targetEntry = new MethodEntry(obf.getClassEntry(), deobfName, obf.getSignature());
-        if (m_mappings.containsDeobfMethod(obf.getClassEntry(), deobfName, obf.getSignature()) || m_index.containsObfBehavior(targetEntry)) {
+        ClassMapping classMapping = getOrCreateClassMapping(obf.getClassEntry());
+
+        // TODO: Verify if I don't break things
+        if ((m_mappings.containsDeobfMethod(obf.getClassEntry(), deobfName, obf.getSignature()) && classMapping.getMethodByObf(obf.getName(), obf.getSignature()) != classMapping.getMethodByDeobf(deobfName, obf.getSignature()))
+                || m_index.containsObfBehavior(targetEntry)) {
             String deobfClassName = m_mappings.getTranslator(TranslationDirection.Deobfuscating, m_index.getTranslationIndex()).translateClass(obf.getClassName());
+            if (deobfClassName == null) {
+                deobfClassName = obf.getClassName();
+            }
             throw new IllegalNameException(deobfName, "There is already a method with that name and signature in class " + deobfClassName);
         }
 
-        ClassMapping classMapping = getOrCreateClassMapping(obf.getClassEntry());
         classMapping.setMethodName(obf.getName(), obf.getSignature(), deobfName);
     }
 
@@ -163,9 +176,21 @@ public class MappingsRenamer {
 
         Set<MethodEntry> implementations = m_index.getRelatedMethodImplementations(obfMethod);
         for (MethodEntry entry : implementations) {
-            // NOTE: don't need to check arguments for name collisions with names determined by Procyon
-            if (m_mappings.containsArgument(entry, deobfName)) {
-                throw new IllegalNameException(deobfName, "There is already an argument with that name");
+            ClassMapping classMapping = m_mappings.getClassByObf(entry.getClassEntry());
+            if (classMapping != null) {
+                MethodMapping mapping = classMapping.getMethodByObf(entry.getName(), entry.getSignature());
+                // NOTE: don't need to check arguments for name collisions with names determined by Procyon
+                // TODO: Verify if I don't break things
+                if (mapping != null) {
+                    for (ArgumentMapping argumentMapping : Lists.newArrayList(mapping.arguments())) {
+                        if (argumentMapping.getIndex() != obf.getIndex()) {
+                            if (mapping.getDeobfArgumentName(argumentMapping.getIndex()).equals(deobfName)
+                                    || argumentMapping.getName().equals(deobfName)) {
+                                throw new IllegalNameException(deobfName, "There is already an argument with that name");
+                            }
+                        }
+                    }
+                }
             }
         }
 
@@ -176,12 +201,21 @@ public class MappingsRenamer {
 
     public void setArgumentName(ArgumentEntry obf, String deobfName) {
         deobfName = NameValidator.validateArgumentName(deobfName);
+        ClassMapping classMapping = getOrCreateClassMapping(obf.getClassEntry());
+        MethodMapping mapping = classMapping.getMethodByObf(obf.getMethodName(), obf.getMethodSignature());
         // NOTE: don't need to check arguments for name collisions with names determined by Procyon
-        if (m_mappings.containsArgument(obf.getBehaviorEntry(), deobfName)) {
-            throw new IllegalNameException(deobfName, "There is already an argument with that name");
+        // TODO: Verify if I don't break things
+        if (mapping != null) {
+            for (ArgumentMapping argumentMapping : Lists.newArrayList(mapping.arguments())) {
+                if (argumentMapping.getIndex() != obf.getIndex()) {
+                    if (mapping.getDeobfArgumentName(argumentMapping.getIndex()).equals(deobfName)
+                            || argumentMapping.getName().equals(deobfName)) {
+                        throw new IllegalNameException(deobfName, "There is already an argument with that name");
+                    }
+                }
+            }
         }
 
-        ClassMapping classMapping = getOrCreateClassMapping(obf.getClassEntry());
         classMapping.setArgumentName(obf.getMethodName(), obf.getMethodSignature(), obf.getIndex(), deobfName);
     }