4 files changed, 35 insertions, 10 deletions

diff --git a/src/cuchaz/enigma/Deobfuscator.java b/src/cuchaz/enigma/Deobfuscator.java
index 323aa2e..293e1c2 100644
--- a/src/cuchaz/enigma/Deobfuscator.java
+++ b/src/cuchaz/enigma/Deobfuscator.java
@@ -29,6 +29,7 @@ import com.strobel.decompiler.languages.java.ast.InsertParenthesesVisitor;
 import cuchaz.enigma.analysis.JarIndex;
 import cuchaz.enigma.analysis.SourceIndex;
 import cuchaz.enigma.analysis.SourceIndexVisitor;
+import cuchaz.enigma.analysis.TreeDumpVisitor;
 import cuchaz.enigma.mapping.ArgumentEntry;
 import cuchaz.enigma.mapping.ClassEntry;
 import cuchaz.enigma.mapping.ClassMapping;
@@ -49,6 +50,7 @@ public class Deobfuscator
         private JarIndex m_jarIndex;
         private Mappings m_mappings;
         private Renamer m_renamer;
+        private TranslatingTypeLoader m_typeLoader;
         
         public Deobfuscator( File file )
         throws IOException
@@ -93,12 +95,13 @@ public class Deobfuscator
                 m_renamer = new Renamer( m_jarIndex, m_mappings );
                 
                 // update decompiler options
-                m_settings.setTypeLoader( new TranslatingTypeLoader(
+                m_typeLoader = new TranslatingTypeLoader(
                         m_jar,
                         m_jarIndex,
                         getTranslator( TranslationDirection.Obfuscating ),
                         getTranslator( TranslationDirection.Deobfuscating )
-                ) );
+                );
+                m_settings.setTypeLoader( m_typeLoader );
         }
         
         public Translator getTranslator( TranslationDirection direction )
@@ -158,7 +161,8 @@ public class Deobfuscator
                 // render the AST into source
                 StringWriter buf = new StringWriter();
                 root.acceptVisitor( new InsertParenthesesVisitor(), null );
-                //root.acceptVisitor( new TreeDumpVisitor( new File( "tree.txt" ) ), null );
+                // DEBUG
+                root.acceptVisitor( new TreeDumpVisitor( new File( "tree.txt" ) ), null );
                 root.acceptVisitor( new JavaOutputVisitor( new PlainTextOutput( buf ), m_settings ), null );
                 
                 // build the source index
@@ -196,6 +200,9 @@ public class Deobfuscator
                 {
                         throw new Error( "Unknown entry type: " + obfEntry.getClass().getName() );
                 }
+                
+                // clear the type loader cache
+                m_typeLoader.clearCache();
         }
         
         public Entry obfuscateEntry( Entry deobfEntry )
diff --git a/src/cuchaz/enigma/TranslatingTypeLoader.java b/src/cuchaz/enigma/TranslatingTypeLoader.java
index ae27f37..cc86364 100644
--- a/src/cuchaz/enigma/TranslatingTypeLoader.java
+++ b/src/cuchaz/enigma/TranslatingTypeLoader.java
@@ -53,6 +53,11 @@ public class TranslatingTypeLoader implements ITypeLoader
                 m_cache = Maps.newHashMap();
         }
         
+        public void clearCache( )
+        {
+                m_cache.clear();
+        }
+        
         @Override
         public boolean tryLoadType( String deobfClassName, Buffer out )
         {
diff --git a/src/cuchaz/enigma/analysis/SourceIndexVisitor.java b/src/cuchaz/enigma/analysis/SourceIndexVisitor.java
index 6c14ee9..f31eb1a 100644
--- a/src/cuchaz/enigma/analysis/SourceIndexVisitor.java
+++ b/src/cuchaz/enigma/analysis/SourceIndexVisitor.java
@@ -201,6 +201,18 @@ public class SourceIndexVisitor implements IAstVisitor<SourceIndex, Void>
                 return recurse( node, index );
         }
         
+        @Override
+        public Void visitEnumValueDeclaration( EnumValueDeclaration node, SourceIndex index )
+        {
+                // treat enum declarations as field declarations
+                FieldDefinition def = node.getUserData( Keys.FIELD_DEFINITION );
+                ClassEntry classEntry = new ClassEntry( def.getDeclaringType().getInternalName() );
+                FieldEntry fieldEntry = new FieldEntry( classEntry, def.getName() );
+                index.addDeclaration( node.getNameToken(), fieldEntry );
+                
+                return recurse( node, index );
+        }
+        
         private Void recurse( AstNode node, SourceIndex index )
         {
                 for( final AstNode child : node.getChildren() )
@@ -561,12 +573,6 @@ public class SourceIndexVisitor implements IAstVisitor<SourceIndex, Void>
         }
         
         @Override
-        public Void visitEnumValueDeclaration( EnumValueDeclaration node, SourceIndex index )
-        {
-                return recurse( node, index );
-        }
-        
-        @Override
         public Void visitAssertStatement( AssertStatement node, SourceIndex index )
         {
                 return recurse( node, index );
diff --git a/src/cuchaz/enigma/gui/GuiController.java b/src/cuchaz/enigma/gui/GuiController.java
index ffeb99a..f305e34 100644
--- a/src/cuchaz/enigma/gui/GuiController.java
+++ b/src/cuchaz/enigma/gui/GuiController.java
@@ -272,7 +272,14 @@ public class GuiController
                                 m_gui.setSource( m_index.getSource() );
                                 if( obfEntryToShow != null )
                                 {
-                                        m_gui.showToken( m_index.getDeclarationToken( m_deobfuscator.deobfuscateEntry( obfEntryToShow ) ) );
+                                        Entry deobfEntryToShow = m_deobfuscator.deobfuscateEntry( obfEntryToShow );
+                                        Token token = m_index.getDeclarationToken( deobfEntryToShow );
+                                        if( token == null )
+                                        {
+                                                // TEMP
+                                                System.out.println( "WARNING: can't find token for " + obfEntryToShow + " -> " + deobfEntryToShow );
+                                        }
+                                        m_gui.showToken( token );
                                 }
                                 
                                 // set the highlighted tokens