2 files changed, 64 insertions, 42 deletions

diff --git a/src/cuchaz/enigma/analysis/JarClassIterator.java b/src/cuchaz/enigma/analysis/JarClassIterator.java
index 6c9f124..10ae805 100644
--- a/src/cuchaz/enigma/analysis/JarClassIterator.java
+++ b/src/cuchaz/enigma/analysis/JarClassIterator.java
@@ -67,33 +67,7 @@ public class JarClassIterator implements Iterator<CtClass>
                 JarEntry entry = m_iter.next();
                 try
                 {
-                        // read the class into a buffer
-                        ByteArrayOutputStream bos = new ByteArrayOutputStream();
-                        byte[] buf = new byte[Constants.KiB];
-                        int totalNumBytesRead = 0;
-                        InputStream in = m_jar.getInputStream( entry );
-                        while( in.available() > 0 )
-                        {
-                                int numBytesRead = in.read( buf );
-                                if( numBytesRead < 0 )
-                                {
-                                        break;
-                                }
-                                bos.write( buf, 0, numBytesRead );
-                                
-                                // sanity checking
-                                totalNumBytesRead += numBytesRead;
-                                if( totalNumBytesRead > Constants.MiB )
-                                {
-                                        throw new Error( "Class file " + entry.getName() + " larger than 1 MiB! Something is wrong!" );
-                                }
-                        }
-                        
-                        // get a javassist handle for the class
-                        String className = Descriptor.toJavaName( getClassEntry( entry ).getName() );
-                        ClassPool classPool = new ClassPool();
-                        classPool.insertClassPath( new ByteArrayClassPath( className, bos.toByteArray() ) );
-                        return classPool.get( className );
+                        return getClass( m_jar, entry );
                 }
                 catch( IOException | NotFoundException ex )
                 {
@@ -136,6 +110,50 @@ public class JarClassIterator implements Iterator<CtClass>
                 };
         }
         
+        public static CtClass getClass( JarFile jar, ClassEntry classEntry )
+        {
+                try
+                {
+                        return getClass( jar, new JarEntry( classEntry.getName() + ".class" ) );
+                }
+                catch( IOException | NotFoundException ex )
+                {
+                        throw new Error( "Unable to load class: " + classEntry.getName() );
+                }
+        }
+        
+        private static CtClass getClass( JarFile jar, JarEntry entry )
+        throws IOException, NotFoundException
+        {
+                // read the class into a buffer
+                ByteArrayOutputStream bos = new ByteArrayOutputStream();
+                byte[] buf = new byte[Constants.KiB];
+                int totalNumBytesRead = 0;
+                InputStream in = jar.getInputStream( entry );
+                while( in.available() > 0 )
+                {
+                        int numBytesRead = in.read( buf );
+                        if( numBytesRead < 0 )
+                        {
+                                break;
+                        }
+                        bos.write( buf, 0, numBytesRead );
+                        
+                        // sanity checking
+                        totalNumBytesRead += numBytesRead;
+                        if( totalNumBytesRead > Constants.MiB )
+                        {
+                                throw new Error( "Class file " + entry.getName() + " larger than 1 MiB! Something is wrong!" );
+                        }
+                }
+                
+                // get a javassist handle for the class
+                String className = Descriptor.toJavaName( getClassEntry( entry ).getName() );
+                ClassPool classPool = new ClassPool();
+                classPool.insertClassPath( new ByteArrayClassPath( className, bos.toByteArray() ) );
+                return classPool.get( className );
+        }
+        
         private static ClassEntry getClassEntry( JarEntry entry )
         {
                 return new ClassEntry( entry.getName().substring( 0, entry.getName().length() - ".class".length() ) );
diff --git a/src/cuchaz/enigma/analysis/JarIndex.java b/src/cuchaz/enigma/analysis/JarIndex.java
index deacf16..b479b69 100644
--- a/src/cuchaz/enigma/analysis/JarIndex.java
+++ b/src/cuchaz/enigma/analysis/JarIndex.java
@@ -95,7 +95,7 @@ public class JarIndex
                 // step 2: index method/field access
                 for( CtClass c : JarClassIterator.classes( jar ) )
                 {
-                        fixClass( c );
+                        ClassRenamer.moveAllClassesOutOfDefaultPackage( c, Constants.NonePackage );
                         ClassEntry classEntry = new ClassEntry( Descriptor.toJvmName( c.getName() ) );
                         for( CtField field : c.getDeclaredFields() )
                         {
@@ -112,7 +112,7 @@ public class JarIndex
                 // step 3: index the types, methods
                 for( CtClass c : JarClassIterator.classes( jar ) )
                 {
-                        fixClass( c );
+                        ClassRenamer.moveAllClassesOutOfDefaultPackage( c, Constants.NonePackage );
                         String className = Descriptor.toJvmName( c.getName() );
                         m_ancestries.addSuperclass( className, Descriptor.toJvmName( c.getClassFile().getSuperclass() ) );
                         for( String interfaceName : c.getClassFile().getInterfaces() )
@@ -128,8 +128,7 @@ public class JarIndex
                 // step 4: index inner classes and anonymous classes
                 for( CtClass c : JarClassIterator.classes( jar ) )
                 {
-                        fixClass( c );
-                        
+                        ClassRenamer.moveAllClassesOutOfDefaultPackage( c, Constants.NonePackage );
                         String outerClassName = findOuterClass( c );
                         if( outerClassName != null )
                         {
@@ -164,17 +163,6 @@ public class JarIndex
                 renameMethods( m_bridgeMethods );
         }
 
-        private void fixClass( CtClass c )
-        {
-                ClassEntry classEntry = new ClassEntry( Descriptor.toJvmName( c.getName() ) );
-                if( classEntry.isInDefaultPackage() )
-                {
-                        // move class out of default package
-                        classEntry = new ClassEntry( Constants.NonePackage + "/" + classEntry.getName() );
-                        ClassRenamer.moveAllClassesOutOfDefaultPackage( c, Constants.NonePackage );
-                }
-        }
-
         private void indexBehavior( CtBehavior behavior )
         {
                 // get the method entry
@@ -729,6 +717,22 @@ public class JarIndex
         
         private void renameClasses( Map<String,String> renames )
         {
+                // rename class entries
+                Set<ClassEntry> obfClassEntries = Sets.newHashSet();
+                for( ClassEntry classEntry : m_obfClassEntries )
+                {
+                        if( renames.containsKey( classEntry.getName() ) )
+                        {
+                                obfClassEntries.add( new ClassEntry( renames.get( classEntry.getName() ) ) );
+                        }
+                        else
+                        {
+                                obfClassEntries.add( classEntry );
+                        }
+                }
+                m_obfClassEntries = obfClassEntries;
+                
+                // rename others
                 m_ancestries.renameClasses( renames );
                 renameClassesInMultimap( renames, m_methodImplementations );
                 renameClassesInMultimap( renames, m_behaviorReferences );