5 files changed, 46 insertions, 34 deletions

diff --git a/src/cuchaz/enigma/Deobfuscator.java b/src/cuchaz/enigma/Deobfuscator.java
index 1d6f02c..9c84532 100644
--- a/src/cuchaz/enigma/Deobfuscator.java
+++ b/src/cuchaz/enigma/Deobfuscator.java
@@ -232,7 +232,7 @@ public class Deobfuscator
                 sourceTree.acceptVisitor( new SourceIndexVisitor(), index );
                 
                 // DEBUG
-                //root.acceptVisitor( new TreeDumpVisitor( new File( "tree.txt" ) ), null );
+                //sourceTree.acceptVisitor( new TreeDumpVisitor( new File( "tree.txt" ) ), null );
 
                 /* DEBUG
                 for( Token token : index.referenceTokens() )
@@ -420,19 +420,7 @@ public class Deobfuscator
         {
                 if( obfEntry instanceof ClassEntry )
                 {
-                        ClassEntry obfClassEntry = (ClassEntry)obfEntry;
-                        if( obfClassEntry.isInnerClass() )
-                        {
-                                // both classes must be in the list
-                                return m_jarIndex.getObfClassEntries().contains( obfClassEntry.getOuterClassEntry() )
-                                        && m_jarIndex.getObfClassEntries().contains( obfClassEntry.getInnerClassName() );
-                                // TODO: make sure this works for the inner class!!
-                        }
-                        else
-                        {
-                                // class must be in the list
-                                return m_jarIndex.getObfClassEntries().contains( obfEntry );
-                        }
+                        return m_jarIndex.getObfClassEntries().contains( obfEntry );
                 }
                 else
                 {
diff --git a/src/cuchaz/enigma/analysis/JarIndex.java b/src/cuchaz/enigma/analysis/JarIndex.java
index c36e9cb..8ebce35 100644
--- a/src/cuchaz/enigma/analysis/JarIndex.java
+++ b/src/cuchaz/enigma/analysis/JarIndex.java
@@ -63,7 +63,7 @@ public class JarIndex
         private Multimap<FieldEntry,EntryReference<FieldEntry,BehaviorEntry>> m_fieldReferences;
         private Multimap<String,String> m_innerClasses;
         private Map<String,String> m_outerClasses;
-        private Set<String> m_anonymousClasses;
+        private Map<String,BehaviorEntry> m_anonymousClasses;
         private Map<MethodEntry,MethodEntry> m_bridgeMethods;
         
         public JarIndex( )
@@ -78,7 +78,7 @@ public class JarIndex
                 m_fieldReferences = HashMultimap.create();
                 m_innerClasses = HashMultimap.create();
                 m_outerClasses = Maps.newHashMap();
-                m_anonymousClasses = Sets.newHashSet();
+                m_anonymousClasses = Maps.newHashMap();
                 m_bridgeMethods = Maps.newHashMap();
         }
         
@@ -161,9 +161,10 @@ public class JarIndex
                                         m_innerClasses.put( outerClassName, innerClassName );
                                         m_outerClasses.put( innerClassName, outerClassName );
                                         
-                                        if( isAnonymousClass( c, outerClassName ) )
+                                        BehaviorEntry enclosingBehavior = isAnonymousClass( c, outerClassName );
+                                        if( enclosingBehavior != null )
                                         {
-                                                m_anonymousClasses.add( innerClassName );
+                                                m_anonymousClasses.put( innerClassName, enclosingBehavior );
                                                 
                                                 // DEBUG
                                                 //System.out.println( "ANONYMOUS: " + outerClassName + "$" + innerClassName );
@@ -188,6 +189,7 @@ public class JarIndex
                         EntryRenamer.renameClassesInMultimap( renames, m_methodImplementations );
                         EntryRenamer.renameClassesInMultimap( renames, m_behaviorReferences );
                         EntryRenamer.renameClassesInMultimap( renames, m_fieldReferences );
+                        EntryRenamer.renameClassesInMap( renames, m_bridgeMethods );
                 }
                 
                 // step 6: update other indices with bridge method info
@@ -618,7 +620,7 @@ public class JarIndex
                 return true;
         }
 
-        private boolean isAnonymousClass( CtClass c, String outerClassName )
+        private BehaviorEntry isAnonymousClass( CtClass c, String outerClassName )
         {
                 ClassEntry innerClassEntry = new ClassEntry( Descriptor.toJvmName( c.getName() ) );
                 
@@ -631,13 +633,13 @@ public class JarIndex
                 // is abstract?
                 if( Modifier.isAbstract( c.getModifiers() ) )
                 {
-                        return false;
+                        return null;
                 }
                 
                 // is there exactly one constructor?
                 if( c.getDeclaredConstructors().length != 1 )
                 {
-                        return false;
+                        return null;
                 }
                 CtConstructor constructor = c.getDeclaredConstructors()[0];
                 
@@ -646,7 +648,7 @@ public class JarIndex
                 Collection<EntryReference<BehaviorEntry,BehaviorEntry>> references = getBehaviorReferences( constructorEntry );
                 if( references.size() != 1 )
                 {
-                        return false;
+                        return null;
                 }
                 
                 // does the caller use this type?
@@ -657,7 +659,7 @@ public class JarIndex
                         if( fieldClass != null && fieldClass.equals( innerClassEntry ) )
                         {
                                 // caller references this type, so it can't be anonymous
-                                return false;
+                                return null;
                         }
                 }
                 for( BehaviorEntry behaviorEntry : getReferencedBehaviors( caller ) )
@@ -668,12 +670,12 @@ public class JarIndex
                                 if( className.equals( innerClassEntry.getName() ) )
                                 {
                                         // caller references this type, so it can't be anonymous
-                                        return false;
+                                        return null;
                                 }
                         }
                 }
                 
-                return true;
+                return caller;
         }
 
         public Set<ClassEntry> getObfClassEntries( )
@@ -905,7 +907,12 @@ public class JarIndex
         
         public boolean isAnonymousClass( String obfInnerClassName )
         {
-                return m_anonymousClasses.contains( obfInnerClassName );
+                return m_anonymousClasses.containsKey( obfInnerClassName );
+        }
+        
+        public BehaviorEntry getAnonymousClassCaller( String obfInnerClassName )
+        {
+                return m_anonymousClasses.get( obfInnerClassName );
         }
         
         public Set<String> getInterfaces( String className )
diff --git a/src/cuchaz/enigma/analysis/SourceIndexBehaviorVisitor.java b/src/cuchaz/enigma/analysis/SourceIndexBehaviorVisitor.java
index a1dac4b..6238b1e 100644
--- a/src/cuchaz/enigma/analysis/SourceIndexBehaviorVisitor.java
+++ b/src/cuchaz/enigma/analysis/SourceIndexBehaviorVisitor.java
@@ -177,14 +177,17 @@ public class SourceIndexBehaviorVisitor extends SourceIndexVisitor
         public Void visitObjectCreationExpression( ObjectCreationExpression node, SourceIndex index )
         {
                 MemberReference ref = node.getUserData( Keys.MEMBER_REFERENCE );
-                ClassEntry classEntry = new ClassEntry( ref.getDeclaringType().getInternalName() );
-                ConstructorEntry constructorEntry = new ConstructorEntry( classEntry, ref.getSignature() );
-                if( node.getType() instanceof SimpleType )
+                if( ref != null )
                 {
-                        index.addReference(
-                                ((SimpleType)node.getType()).getIdentifierToken(),
-                                new EntryReference<Entry,Entry>( constructorEntry, m_behaviorEntry )
-                        );
+                        ClassEntry classEntry = new ClassEntry( ref.getDeclaringType().getInternalName() );
+                        ConstructorEntry constructorEntry = new ConstructorEntry( classEntry, ref.getSignature() );
+                        if( node.getType() instanceof SimpleType )
+                        {
+                                index.addReference(
+                                        ((SimpleType)node.getType()).getIdentifierToken(),
+                                        new EntryReference<Entry,Entry>( constructorEntry, m_behaviorEntry )
+                                );
+                        }
                 }
                 
                 return recurse( node, index );
diff --git a/src/cuchaz/enigma/bytecode/InnerClassWriter.java b/src/cuchaz/enigma/bytecode/InnerClassWriter.java
index 5044e06..f94a85d 100644
--- a/src/cuchaz/enigma/bytecode/InnerClassWriter.java
+++ b/src/cuchaz/enigma/bytecode/InnerClassWriter.java
@@ -16,8 +16,10 @@ import javassist.CtClass;
 import javassist.bytecode.AccessFlag;
 import javassist.bytecode.ConstPool;
 import javassist.bytecode.Descriptor;
+import javassist.bytecode.EnclosingMethodAttribute;
 import javassist.bytecode.InnerClassesAttribute;
 import cuchaz.enigma.analysis.JarIndex;
+import cuchaz.enigma.mapping.BehaviorEntry;
 import cuchaz.enigma.mapping.ClassEntry;
 
 public class InnerClassWriter
@@ -44,6 +46,18 @@ public class InnerClassWriter
                         // this is an inner class, rename it to outer$inner
                         ClassEntry obfClassEntry = new ClassEntry( obfOuterClassName + "$" + new ClassEntry( obfClassName ).getSimpleName() );
                         c.setName( obfClassEntry.getName() );
+                        
+                        BehaviorEntry caller = m_jarIndex.getAnonymousClassCaller( obfClassName );
+                        if( caller != null )
+                        {
+                                // write the enclosing method attribute
+                                c.getClassFile().addAttribute( new EnclosingMethodAttribute(
+                                        c.getClassFile().getConstPool(),
+                                        caller.getClassName(),
+                                        caller.getName(),
+                                        caller.getSignature()
+                                ) );
+                        }
                 }
                 
                 // write the inner classes if needed
diff --git a/src/cuchaz/enigma/mapping/SignatureUpdater.java b/src/cuchaz/enigma/mapping/SignatureUpdater.java
index 528a743..d1216bd 100644
--- a/src/cuchaz/enigma/mapping/SignatureUpdater.java
+++ b/src/cuchaz/enigma/mapping/SignatureUpdater.java
@@ -14,7 +14,7 @@ import java.io.IOException;
 import java.io.StringReader;
 import java.util.List;
 
-import com.beust.jcommander.internal.Lists;
+import com.google.common.collect.Lists;
 
 public class SignatureUpdater
 {