7 files changed, 223 insertions, 148 deletions

diff --git a/src/cuchaz/enigma/Deobfuscator.java b/src/cuchaz/enigma/Deobfuscator.java
index e067183e..e6e647e9 100644
--- a/src/cuchaz/enigma/Deobfuscator.java
+++ b/src/cuchaz/enigma/Deobfuscator.java
@@ -10,10 +10,12 @@
  ******************************************************************************/
 package cuchaz.enigma;
 
+import java.io.BufferedReader;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.StringReader;
 import java.io.StringWriter;
 import java.util.Enumeration;
 import java.util.List;
@@ -21,6 +23,7 @@ import java.util.Map;
 import java.util.jar.JarEntry;
 import java.util.jar.JarFile;
 
+import com.beust.jcommander.internal.Lists;
 import com.strobel.decompiler.Decompiler;
 import com.strobel.decompiler.DecompilerSettings;
 import com.strobel.decompiler.PlainTextOutput;
@@ -45,6 +48,7 @@ public class Deobfuscator
         private Ancestries m_ancestries;
         private Mappings m_mappings;
         private Renamer m_renamer;
+        private List<String> m_obfClassNames;
         
         public Deobfuscator( File file )
         throws IOException
@@ -65,6 +69,26 @@ public class Deobfuscator
                         Util.closeQuietly( jarIn );
                 }
                 
+                // get the obf class names
+                m_obfClassNames = Lists.newArrayList();
+                {
+                        Enumeration<JarEntry> entries = m_jar.entries();
+                        while( entries.hasMoreElements() )
+                        {
+                                JarEntry entry = entries.nextElement();
+                                
+                                // skip everything but class files
+                                if( !entry.getName().endsWith( ".class" ) )
+                                {
+                                        continue;
+                                }
+                                
+                                // get the class name from the file
+                                String className = entry.getName().substring( 0, entry.getName().length() - 6 );
+                                m_obfClassNames.add( className );
+                        }
+                }
+                
                 // config the decompiler
                 m_settings = DecompilerSettings.javaDefaults();
                 m_settings.setForceExplicitImports( true );
@@ -112,20 +136,9 @@ public class Deobfuscator
         
         public void getSeparatedClasses( List<ClassFile> obfClasses, Map<ClassFile,String> deobfClasses )
         {
-                Enumeration<JarEntry> entries = m_jar.entries();
-                while( entries.hasMoreElements() )
+                for( String obfClassName : m_obfClassNames )
                 {
-                        JarEntry entry = entries.nextElement();
-                        
-                        // skip everything but class files
-                        if( !entry.getName().endsWith( ".class" ) )
-                        {
-                                continue;
-                        }
-                        
-                        // get the class name from the file
-                        String className = entry.getName().substring( 0, entry.getName().length() - 6 );
-                        ClassFile classFile = new ClassFile( className );
+                        ClassFile classFile = new ClassFile( obfClassName );
                         
                         // separate the classes
                         ClassMapping classMapping = m_mappings.getClassByObf( classFile.getName() );
@@ -159,7 +172,41 @@ public class Deobfuscator
                 // decompile it!
                 StringWriter buf = new StringWriter();
                 Decompiler.decompile( deobfName, new PlainTextOutput( buf ), m_settings );
-                return buf.toString();
+                return fixSource( buf.toString() );
+        }
+        
+        private String fixSource( String source )
+        {
+                // fix the imports from the default package in the source
+                try
+                {
+                        StringBuilder buf = new StringBuilder();
+                        BufferedReader reader = new BufferedReader( new StringReader( source ) );
+                        String line = null;
+                        while( ( line = reader.readLine() ) != null )
+                        {
+                                String[] parts = line.trim().split( " " );
+                                if( parts.length == 2 && parts[0].equals( "import" ) )
+                                {
+                                        // is this an (illegal) import from the default package?
+                                        String className = parts[1];
+                                        if( className.indexOf( '.' ) < 0 )
+                                        {
+                                                // this is an illegal import, replace it
+                                                line = "import __DEFAULT__." + parts[1];
+                                        }
+                                }
+                                
+                                buf.append( line );
+                                buf.append( "\n" );
+                        }
+                        return buf.toString();
+                }
+                catch( IOException ex )
+                {
+                        // dealing with IOExceptions on StringReaders is silly...
+                        throw new Error( ex );
+                }
         }
         
         // NOTE: these methods are a bit messy... oh well
@@ -265,4 +312,16 @@ public class Deobfuscator
                         throw new Error( "Unknown entry type: " + obfEntry.getClass().getName() );
                 }
         }
+
+        public boolean entryIsObfuscatedIdenfitier( Entry obfEntry )
+        {
+                if( obfEntry instanceof ClassEntry )
+                {
+                        // obf classes must be in the list
+                        return m_obfClassNames.contains( obfEntry.getName() );
+                }
+                
+                // assume everything else is an identifier
+                return true;
+        }
 }
diff --git a/src/cuchaz/enigma/analysis/Analyzer.java b/src/cuchaz/enigma/analysis/Analyzer.java
index dad8dc52..2b7e0b0b 100644
--- a/src/cuchaz/enigma/analysis/Analyzer.java
+++ b/src/cuchaz/enigma/analysis/Analyzer.java
@@ -35,6 +35,7 @@ import com.sun.source.util.Trees;
 
 import cuchaz.enigma.mapping.ArgumentEntry;
 import cuchaz.enigma.mapping.ClassEntry;
+import cuchaz.enigma.mapping.Entry;
 import cuchaz.enigma.mapping.FieldEntry;
 import cuchaz.enigma.mapping.MethodEntry;
 
@@ -78,39 +79,31 @@ class TreeVisitor extends TreeScanner<CompilationUnitTree, SourcedAst>
         
         private ClassEntry indexClass( ClassTree classTree, SourcedAst ast )
         {
-                // build the entry
-                ClassEntry entry = new ClassEntry( ast.getFullClassName( classTree.getSimpleName().toString() ) );
+                // index the class name
+                ClassEntry classEntry = indexClassIdentifier( classTree, ast );
+                assert( classEntry != null );
                 
-                // lex the source at this tree node
-                for( Token token : new Lexer( ast.getSource( classTree ).toString() ) )
+                // index the extends clause
+                indexClassIdentifier( classTree.getExtendsClause(), ast );
+                
+                // index the implements clauses
+                for( Tree implementsTree : classTree.getImplementsClause() )
                 {
-                        // scan until we get the first identifier
-                        if( token.type == TokenType.IDENTIFIER )
-                        {
-                                m_index.add( entry, offsetToken( token, ast.getStart( classTree ) ) );
-                                break;
-                        }
+                        indexClassIdentifier( implementsTree, ast );
                 }
                 
-                return entry;
+                return classEntry;
         }
         
         private FieldEntry indexField( VariableTree variableTree, SourcedAst ast, ClassEntry classEntry )
         {
-                // build the entry
+                // index the field name
                 FieldEntry entry = new FieldEntry( classEntry, variableTree.getName().toString() );
+                Token nameToken = new Lexer( ast.getSource( variableTree ) ).getFirstIdentifierMatching( variableTree.getName() );
+                addToken( entry, nameToken, variableTree, ast );
                 
-                // lex the source at this tree node
-                Lexer lexer = new Lexer( ast.getSource( variableTree ).toString() );
-                for( Token token : lexer )
-                {
-                        // scan until we find an identifier that matches the field name
-                        if( token.type == TokenType.IDENTIFIER && lexer.getText( token ).equals( entry.getName() ) )
-                        {
-                                m_index.add( entry, offsetToken( token, ast.getStart( variableTree ) ) );
-                                break;
-                        }
-                }
+                // index the field type
+                indexClassIdentifier( variableTree.getType(), ast );
                 
                 return entry;
         }
@@ -136,13 +129,13 @@ class TreeVisitor extends TreeScanner<CompilationUnitTree, SourcedAst>
                 MethodEntry entry = new MethodEntry( classEntry, methodTree.getName().toString(), signature.toString() );
                 
                 // lex the source at this tree node
-                Lexer lexer = new Lexer( ast.getSource( methodTree ).toString() );
+                Lexer lexer = new Lexer( ast.getSource( methodTree ) );
                 for( Token token : lexer )
                 {
                         // scan until we find an identifier that matches the method name
                         if( token.type == TokenType.IDENTIFIER && lexer.getText( token ).equals( entry.getName() ) )
                         {
-                                m_index.add( entry, offsetToken( token, ast.getStart( methodTree ) ) );
+                                addToken( entry, token, methodTree, ast );
                                 break;
                         }
                 }
@@ -152,27 +145,47 @@ class TreeVisitor extends TreeScanner<CompilationUnitTree, SourcedAst>
         
         private void indexArgument( VariableTree variableTree, SourcedAst ast, MethodEntry methodEntry, int index )
         {
-                // build the entry
+                // index argument name
                 ArgumentEntry entry = new ArgumentEntry( methodEntry, index, variableTree.getName().toString() );
+                Token token = new Lexer( ast.getSource( variableTree ) ).getLastIdentifier();
+                addToken( entry, token, variableTree, ast );
                 
-                // lex the source at this tree node
-                Lexer lexer = new Lexer( ast.getSource( variableTree ).toString() );
-                for( Token token : lexer )
+                // index argument type
+                indexClassIdentifier( variableTree.getType(), ast );
+        }
+        
+        private ClassEntry indexClassIdentifier( Tree tree, SourcedAst ast )
+        {
+                if( tree == null )
                 {
-                        // scan until we find an identifier that matches the variable name
-                        if( token.type == TokenType.IDENTIFIER && lexer.getText( token ).equals( entry.getName() ) )
-                        {
-                                m_index.add( entry, offsetToken( token, ast.getStart( variableTree ) ) );
-                                break;
-                        }
+                        return null;
+                }
+                
+                Lexer lexer = new Lexer( ast.getSource( tree ) );
+                Token token = lexer.getFirstIdentifier();
+                if( token == null )
+                {
+                        return null;
                 }
+                
+                ClassEntry classEntry = new ClassEntry( ast.getFullClassName( lexer.getText( token ) ) );
+                addToken( classEntry, token, tree, ast );
+                return classEntry;
         }
-        
-        private Token offsetToken( Token in, int offset )
+
+        private void addToken( Entry entry, Token token, Tree tree, SourcedAst ast )
         {
-                return new Token( in.type, in.start + offset, in.length );
+                if( token == null )
+                {
+                        throw new IllegalArgumentException( "token cannot be null!" );
+                }
+                
+                // offset the token by the tree
+                Token offsetToken = new Token( token.type, token.start + ast.getStart( tree ), token.length );
+                
+                m_index.add( entry, offsetToken );
         }
-        
+
         private String toJvmType( Tree tree, SourcedAst ast )
         {
                 switch( tree.getKind() )
diff --git a/src/cuchaz/enigma/analysis/Lexer.java b/src/cuchaz/enigma/analysis/Lexer.java
index acb52bf8..602e3a9b 100644
--- a/src/cuchaz/enigma/analysis/Lexer.java
+++ b/src/cuchaz/enigma/analysis/Lexer.java
@@ -14,6 +14,7 @@ import java.util.Iterator;
 
 import jsyntaxpane.SyntaxDocument;
 import jsyntaxpane.Token;
+import jsyntaxpane.TokenType;
 import jsyntaxpane.lexers.JavaLexer;
 
 public class Lexer implements Iterable<Token>
@@ -21,10 +22,10 @@ public class Lexer implements Iterable<Token>
         private SyntaxDocument m_doc;
         private Iterator<Token> m_iter;
         
-        public Lexer( String source )
+        public Lexer( CharSequence source )
         {
                 m_doc = new SyntaxDocument( new JavaLexer() );
-                m_doc.append( source );
+                m_doc.append( source.toString() );
                 m_iter = m_doc.getTokens( 0, m_doc.getLength() );
         }
         
@@ -38,4 +39,41 @@ public class Lexer implements Iterable<Token>
         {
                 return token.getString( m_doc );
         }
+        
+        public Token getFirstIdentifier( )
+        {
+                for( Token token : this )
+                {
+                        if( token.type == TokenType.IDENTIFIER )
+                        {
+                                return token;
+                        }
+                }
+                return null;
+        }
+        
+        public Token getFirstIdentifierMatching( CharSequence val )
+        {
+                for( Token token : this )
+                {
+                        if( token.type == TokenType.IDENTIFIER && getText( token ).equals( val.toString() ) )
+                        {
+                                return token;
+                        }
+                }
+                return null;
+        }
+        
+        public Token getLastIdentifier( )
+        {
+                Token lastToken = null;
+                for( Token token : this )
+                {
+                        if( token.type == TokenType.IDENTIFIER )
+                        {
+                                lastToken = token;
+                        }
+                }
+                return lastToken;
+        }
 }
diff --git a/src/cuchaz/enigma/analysis/SourceIndex.java b/src/cuchaz/enigma/analysis/SourceIndex.java
index 61c833ce..de163087 100644
--- a/src/cuchaz/enigma/analysis/SourceIndex.java
+++ b/src/cuchaz/enigma/analysis/SourceIndex.java
@@ -10,46 +10,52 @@
  ******************************************************************************/
 package cuchaz.enigma.analysis;
 
+import java.util.Collection;
 import java.util.Iterator;
 import java.util.Map;
-import java.util.Set;
 
 import jsyntaxpane.Token;
 
-import com.google.common.collect.BiMap;
-import com.google.common.collect.HashBiMap;
+import com.google.common.collect.HashMultimap;
+import com.google.common.collect.Multimap;
 
 import cuchaz.enigma.mapping.Entry;
 
 public class SourceIndex implements Iterable<Map.Entry<Entry,Token>>
 {
-        private BiMap<Entry,Token> m_entryToToken;
-        private BiMap<Token,Entry> m_tokenToEntry;
+        private Multimap<Entry,Token> m_entryToTokens;
         
         public SourceIndex( )
         {
-                m_entryToToken = HashBiMap.create();
-                m_tokenToEntry = m_entryToToken.inverse();
+                m_entryToTokens = HashMultimap.create();
         }
         
         public void add( Entry entry, Token token )
         {
-                m_entryToToken.put( entry, token );
+                m_entryToTokens.put( entry, token );
         }
         
         public Iterator<Map.Entry<Entry,Token>> iterator( )
         {
-                return m_entryToToken.entrySet().iterator();
+                return m_entryToTokens.entries().iterator();
         }
         
-        public Set<Token> tokens( )
+        public Collection<Token> tokens( )
         {
-                return m_entryToToken.values();
+                return m_entryToTokens.values();
         }
         
         public Entry getEntry( Token token )
         {
-                return m_tokenToEntry.get( token );
+                // linear search is fast enough for now
+                for( Map.Entry<Entry,Token> entry : this )
+                {
+                        if( entry.getValue().equals( token ) )
+                        {
+                                return entry.getKey();
+                        }
+                }
+                return null;
         }
         
         public Map.Entry<Entry,Token> getEntry( int pos )
@@ -66,8 +72,8 @@ public class SourceIndex implements Iterable<Map.Entry<Entry,Token>>
                 return null;
         }
         
-        public Token getToken( Entry entry )
+        public Collection<Token> getTokens( Entry entry )
         {
-                return m_entryToToken.get( entry );
+                return m_entryToTokens.get( entry );
         }
 }
diff --git a/src/cuchaz/enigma/analysis/SourcedAst.java b/src/cuchaz/enigma/analysis/SourcedAst.java
index 968c8804..a88cc754 100644
--- a/src/cuchaz/enigma/analysis/SourcedAst.java
+++ b/src/cuchaz/enigma/analysis/SourcedAst.java
@@ -16,7 +16,6 @@ import java.util.HashMap;
 import javassist.bytecode.Descriptor;
 
 import com.google.common.collect.Maps;
-import com.sun.source.tree.ClassTree;
 import com.sun.source.tree.CompilationUnitTree;
 import com.sun.source.tree.ImportTree;
 import com.sun.source.tree.Tree;
@@ -29,6 +28,7 @@ public class SourcedAst
         private Trees m_trees;
         private SourcePositions m_positions;
         private HashMap<String,String> m_classNameIndex;
+        private String m_packageName;
         
         public SourcedAst( CompilationUnitTree tree, Trees trees )
         {
@@ -49,6 +49,13 @@ public class SourcedAst
                         // get the full and simple class names
                         String fullName = Descriptor.toJvmName( importTree.getQualifiedIdentifier().toString() );
                         String simpleName = fullName;
+                        
+                        if( fullName.startsWith( "__DEFAULT__/" ) )
+                        {
+                                // remove the default package flag
+                                fullName = fullName.substring( 12 );
+                        }
+                        
                         String[] parts = fullName.split( "/" );
                         if( parts.length > 0 )
                         {
@@ -58,30 +65,26 @@ public class SourcedAst
                         m_classNameIndex.put( simpleName, fullName );
                 }
                 
-                // index the self class using the package name
+                // get the package name
+                m_packageName = null;
                 if( m_tree.getPackageName() != null )
                 {
-                        String packageName = Descriptor.toJvmName( m_tree.getPackageName().toString() );
-                        for( Tree typeTree : m_tree.getTypeDecls() )
-                        {
-                                if( typeTree instanceof ClassTree )
-                                {
-                                        ClassTree classTree = (ClassTree)typeTree;
-                                        String className = classTree.getSimpleName().toString();
-                                        m_classNameIndex.put( className, packageName + "/" + className );
-                                }
-                        }
+                        m_packageName = Descriptor.toJvmName( m_tree.getPackageName().toString() );
                 }
         }
 
         public int getStart( Tree node )
         {
-                return (int)m_positions.getStartPosition( m_tree, node );
+                int pos = (int)m_positions.getStartPosition( m_tree, node );
+                assert( pos >= 0 );
+                return pos;
         }
 
         public int getEnd( Tree node )
         {
-                return (int)m_positions.getEndPosition( m_tree, node );
+                int pos = (int)m_positions.getEndPosition( m_tree, node );
+                assert( pos >= 0 );
+                return pos;
         }
         
         public int getLine( Tree node )
@@ -121,8 +124,16 @@ public class SourcedAst
                 String fullClassName = m_classNameIndex.get( simpleClassName );
                 if( fullClassName == null )
                 {
-                        // no mapping was found, the name is probably already fully-qualified
-                        fullClassName = simpleClassName;
+                        if( m_packageName != null )
+                        {
+                                // no mapping was found, assume it's in the package
+                                fullClassName = m_packageName + "/" + simpleClassName;
+                        }
+                        else
+                        {
+                                // this must be in the default package
+                                fullClassName = simpleClassName;
+                        }
                 }
                 return fullClassName;
         }
diff --git a/src/cuchaz/enigma/gui/GuiController.java b/src/cuchaz/enigma/gui/GuiController.java
index 7d37febf..2219e05b 100644
--- a/src/cuchaz/enigma/gui/GuiController.java
+++ b/src/cuchaz/enigma/gui/GuiController.java
@@ -132,6 +132,16 @@ public class GuiController
                 return m_deobfuscator.hasMapping( pair.obf );
         }
         
+        public boolean entryIsObfuscatedIdenfitier( int pos )
+        {
+                EntryPair<Entry> pair = getEntryPair( pos );
+                if( pair == null || pair.obf == null )
+                {
+                        return false;
+                }
+                return m_deobfuscator.entryIsObfuscatedIdenfitier( pair.obf );
+        }
+        
         public ClassInheritanceTreeNode getClassInheritance( ClassEntry classEntry )
         {
                 Translator deobfuscatingTranslator = m_deobfuscator.getTranslator( TranslationDirection.Deobfuscating );
@@ -216,7 +226,7 @@ public class GuiController
                                         {
                                                 deobfuscatedTokens.add( token );
                                         }
-                                        else
+                                        else if( entryIsObfuscatedIdenfitier( token.start ) )
                                         {
                                                 obfuscatedTokens.add( token );
                                         }
diff --git a/src/cuchaz/enigma/gui/SourceFormatter.java b/src/cuchaz/enigma/gui/SourceFormatter.java
deleted file mode 100644
index f3878405..00000000
--- a/src/cuchaz/enigma/gui/SourceFormatter.java
+++ /dev/null
@@ -1,62 +0,0 @@
-/*******************************************************************************
- * Copyright (c) 2014 Jeff Martin.
- * All rights reserved. This program and the accompanying materials
- * are made available under the terms of the GNU Public License v3.0
- * which accompanies this distribution, and is available at
- * http://www.gnu.org/licenses/gpl.html
- * 
- * Contributors:
- *     Jeff Martin - initial API and implementation
- ******************************************************************************/
-package cuchaz.enigma.gui;
-
-import java.io.BufferedReader;
-import java.io.IOException;
-import java.io.StringReader;
-
-public class SourceFormatter
-{
-        public static String format( String in )
-        {
-                return collapseNewlines( in );
-        }
-        
-        private static String collapseNewlines( String in )
-        {
-                StringBuffer buf = new StringBuffer();
-                int numBlankLines = 0;
-                
-                BufferedReader reader = new BufferedReader( new StringReader( in ) );
-                String line = null;
-                try
-                {
-                        while( ( line = reader.readLine() ) != null )
-                        {
-                                // how blank lines is this?
-                                boolean isBlank = line.trim().length() == 0;
-                                if( isBlank )
-                                {
-                                        numBlankLines++;
-                                        
-                                        // stop printing blank lines after the first one
-                                        if( numBlankLines < 2 )
-                                        {
-                                                buf.append( line );
-                                                buf.append( "\n" );
-                                        }
-                                }
-                                else
-                                {
-                                        numBlankLines = 0;
-                                        buf.append( line );
-                                        buf.append( "\n" );
-                                }
-                        }
-                }
-                catch( IOException ex )
-                {
-                        // StringReader will never throw an IOExecption here...
-                }
-                return buf.toString();
-        }
-}