refactor Ancestries into Ancestries and JarIndex

6 files changed, 197 insertions, 166 deletions

diff --git a/src/cuchaz/enigma/Deobfuscator.java b/src/cuchaz/enigma/Deobfuscator.java
index 5321d2d..c35a483 100644
--- a/src/cuchaz/enigma/Deobfuscator.java
+++ b/src/cuchaz/enigma/Deobfuscator.java
@@ -31,7 +31,7 @@ import com.strobel.decompiler.languages.java.ast.AstBuilder;
 import com.strobel.decompiler.languages.java.ast.CompilationUnit;
 import com.strobel.decompiler.languages.java.ast.InsertParenthesesVisitor;
 
-import cuchaz.enigma.analysis.Ancestries;
+import cuchaz.enigma.analysis.JarIndex;
 import cuchaz.enigma.analysis.SourceIndex;
 import cuchaz.enigma.analysis.SourceIndexVisitor;
 import cuchaz.enigma.mapping.ArgumentEntry;
@@ -50,7 +50,7 @@ public class Deobfuscator
         private File m_file;
         private JarFile m_jar;
         private DecompilerSettings m_settings;
-        private Ancestries m_ancestries;
+        private JarIndex m_jarIndex;
         private Mappings m_mappings;
         private Renamer m_renamer;
         private List<String> m_obfClassNames;
@@ -65,9 +65,9 @@ public class Deobfuscator
                 InputStream jarIn = null;
                 try
                 {
-                        m_ancestries = new Ancestries();
+                        m_jarIndex = new JarIndex();
                         jarIn = new FileInputStream( m_file );
-                        m_ancestries.readFromJar( jarIn );
+                        m_jarIndex.indexJar( jarIn );
                 }
                 finally
                 {
@@ -107,9 +107,9 @@ public class Deobfuscator
                 return m_file.getName();
         }
         
-        public Ancestries getAncestries( )
+        public JarIndex getJarIndex( )
         {
-                return m_ancestries;
+                return m_jarIndex;
         }
         
         public Mappings getMappings( )
@@ -123,7 +123,7 @@ public class Deobfuscator
                         val = new Mappings();
                 }
                 m_mappings = val;
-                m_renamer = new Renamer( m_ancestries, m_mappings );
+                m_renamer = new Renamer( m_jarIndex, m_mappings );
                 
                 // update decompiler options
                 m_settings.setTypeLoader( new TranslatingTypeLoader(
@@ -135,7 +135,7 @@ public class Deobfuscator
         
         public Translator getTranslator( TranslationDirection direction )
         {
-                return m_mappings.getTranslator( m_ancestries, direction );
+                return m_mappings.getTranslator( m_jarIndex.getAncestries(), direction );
         }
         
         public void getSeparatedClasses( List<String> obfClasses, List<String> deobfClasses )
diff --git a/src/cuchaz/enigma/analysis/Ancestries.java b/src/cuchaz/enigma/analysis/Ancestries.java
index e6c8bbf..83c239c 100644
--- a/src/cuchaz/enigma/analysis/Ancestries.java
+++ b/src/cuchaz/enigma/analysis/Ancestries.java
@@ -10,110 +10,27 @@
  ******************************************************************************/
 package cuchaz.enigma.analysis;
 
-import java.io.ByteArrayOutputStream;
-import java.io.IOException;
-import java.io.InputStream;
 import java.io.Serializable;
 import java.util.ArrayList;
-import java.util.Collection;
 import java.util.List;
 import java.util.Map;
-import java.util.zip.ZipEntry;
-import java.util.zip.ZipInputStream;
 
-import javassist.ByteArrayClassPath;
-import javassist.ClassPool;
-import javassist.CtClass;
-import javassist.NotFoundException;
 import javassist.bytecode.Descriptor;
-import javassist.bytecode.MethodInfo;
 
-import com.google.common.collect.HashMultimap;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
-import com.google.common.collect.Multimap;
-
-import cuchaz.enigma.Constants;
-import cuchaz.enigma.mapping.ClassEntry;
-import cuchaz.enigma.mapping.MethodEntry;
-import cuchaz.enigma.mapping.Translator;
 
 public class Ancestries implements Serializable
 {
         private static final long serialVersionUID = 738687982126844179L;
         
         private Map<String,String> m_superclasses;
-        private Multimap<String,String> m_methodImplementations;
         
         public Ancestries( )
         {
                 m_superclasses = Maps.newHashMap();
-                m_methodImplementations = HashMultimap.create();
         }
         
-        @SuppressWarnings( "unchecked" )
-        public void readFromJar( InputStream in )
-        throws IOException
-        {
-                ClassPool classPool = new ClassPool();
-                
-                ZipInputStream zin = new ZipInputStream( in );
-                ZipEntry entry;
-                while( ( entry = zin.getNextEntry() ) != null )
-                {
-                        // filter out non-classes
-                        if( entry.isDirectory() || !entry.getName().endsWith( ".class" ) )
-                        {
-                                continue;
-                        }
-                        
-                        // read the class into a buffer
-                        ByteArrayOutputStream bos = new ByteArrayOutputStream();
-                        byte[] buf = new byte[Constants.KiB];
-                        int totalNumBytesRead = 0;
-                        while( zin.available() > 0 )
-                        {
-                                int numBytesRead = zin.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!" );
-                                }
-                        }
-                        
-                        // determine the class name (ie chop off the ".class")
-                        String className = Descriptor.toJavaName( entry.getName().substring( 0, entry.getName().length() - ".class".length() ) );
-                        
-                        // get a javassist handle for the class
-                        classPool.insertClassPath( new ByteArrayClassPath( className, bos.toByteArray() ) );
-                        try
-                        {
-                                CtClass c = classPool.get( className );
-                                addSuperclass( c.getName(), c.getClassFile().getSuperclass() );
-                                addMethodImplementations( c.getName(), (List<MethodInfo>)c.getClassFile().getMethods() );
-                        }
-                        catch( NotFoundException ex )
-                        {
-                                throw new Error( "Unable to load class: " + className );
-                        }
-                }
-        }
-        
-        private void addMethodImplementations( String name, List<MethodInfo> methods )
-        {
-                for( MethodInfo method : methods )
-                {
-                        m_methodImplementations.put( name, getMethodKey( method.getName(), method.getDescriptor() ) );
-                }
-        }
-
         public void addSuperclass( String className, String superclassName )
         {
                 className = Descriptor.toJvmName( className );
@@ -165,71 +82,9 @@ public class Ancestries implements Serializable
                 return subclasses;
         }
         
-        public boolean isMethodImplemented( MethodEntry methodEntry )
-        {
-                return isMethodImplemented( methodEntry.getClassName(), methodEntry.getName(), methodEntry.getSignature() );
-        }
-        
-        public boolean isMethodImplemented( String className, String methodName, String methodSignature )
-        {
-                Collection<String> implementations = m_methodImplementations.get( className );
-                if( implementations == null )
-                {
-                        return false;
-                }
-                return implementations.contains( getMethodKey( methodName, methodSignature ) );
-        }
-        
-        public ClassInheritanceTreeNode getClassInheritance( Translator deobfuscatingTranslator, ClassEntry obfClassEntry )
-        {
-                // get the root node
-                List<String> ancestry = getAncestry( obfClassEntry.getName() );
-                ClassInheritanceTreeNode rootNode = new ClassInheritanceTreeNode( deobfuscatingTranslator, ancestry.get( ancestry.size() - 1 ) );
-                
-                // expand all children recursively
-                rootNode.load( this, true );
-                
-                return rootNode;
-        }
-        
-        public MethodInheritanceTreeNode getMethodInheritance( Translator deobfuscatingTranslator, MethodEntry obfMethodEntry )
-        {
-                // travel to the ancestor implementation
-                String baseImplementationClassName = obfMethodEntry.getClassName();
-                for( String ancestorClassName : getAncestry( obfMethodEntry.getClassName() ) )
-                {
-                        if( isMethodImplemented( ancestorClassName, obfMethodEntry.getName(), obfMethodEntry.getSignature() ) )
-                        {
-                                baseImplementationClassName = ancestorClassName;
-                        }
-                }
-                
-                // make a root node at the base
-                MethodEntry methodEntry = new MethodEntry(
-                        new ClassEntry( baseImplementationClassName ),
-                        obfMethodEntry.getName(),
-                        obfMethodEntry.getSignature()
-                );
-                MethodInheritanceTreeNode rootNode = new MethodInheritanceTreeNode(
-                        deobfuscatingTranslator,
-                        methodEntry,
-                        isMethodImplemented( methodEntry )
-                );
-                
-                // expand the full tree
-                rootNode.load( this, true );
-                
-                return rootNode;
-        }
-        
         private boolean isJre( String className )
         {
                 return className.startsWith( "java/" )
                         || className.startsWith( "javax/" );
         }
-        
-        private String getMethodKey( String name, String signature )
-        {
-                return name + signature;
-        }
 }
diff --git a/src/cuchaz/enigma/analysis/JarIndex.java b/src/cuchaz/enigma/analysis/JarIndex.java
new file mode 100644
index 0000000..8a8384c
--- /dev/null
+++ b/src/cuchaz/enigma/analysis/JarIndex.java
@@ -0,0 +1,176 @@
+/*******************************************************************************
+ * 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.analysis;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Collection;
+import java.util.List;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipInputStream;
+
+import com.google.common.collect.HashMultimap;
+import com.google.common.collect.Multimap;
+
+import javassist.ByteArrayClassPath;
+import javassist.ClassPool;
+import javassist.CtClass;
+import javassist.NotFoundException;
+import javassist.bytecode.Descriptor;
+import javassist.bytecode.MethodInfo;
+import cuchaz.enigma.Constants;
+import cuchaz.enigma.mapping.ClassEntry;
+import cuchaz.enigma.mapping.MethodEntry;
+import cuchaz.enigma.mapping.Translator;
+
+public class JarIndex
+{
+        private Ancestries m_ancestries;
+        private Multimap<String,String> m_methodImplementations;
+        
+        public JarIndex( )
+        {
+                m_ancestries = new Ancestries();
+                m_methodImplementations = HashMultimap.create();
+        }
+        
+        @SuppressWarnings( "unchecked" )
+        public void indexJar( InputStream in )
+        throws IOException
+        {
+                ClassPool classPool = new ClassPool();
+                
+                ZipInputStream zin = new ZipInputStream( in );
+                ZipEntry entry;
+                while( ( entry = zin.getNextEntry() ) != null )
+                {
+                        // filter out non-classes
+                        if( entry.isDirectory() || !entry.getName().endsWith( ".class" ) )
+                        {
+                                continue;
+                        }
+                        
+                        // read the class into a buffer
+                        ByteArrayOutputStream bos = new ByteArrayOutputStream();
+                        byte[] buf = new byte[Constants.KiB];
+                        int totalNumBytesRead = 0;
+                        while( zin.available() > 0 )
+                        {
+                                int numBytesRead = zin.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!" );
+                                }
+                        }
+                        
+                        // determine the class name (ie chop off the ".class")
+                        String className = Descriptor.toJavaName( entry.getName().substring( 0, entry.getName().length() - ".class".length() ) );
+                        
+                        // get a javassist handle for the class
+                        classPool.insertClassPath( new ByteArrayClassPath( className, bos.toByteArray() ) );
+                        try
+                        {
+                                CtClass c = classPool.get( className );
+                                m_ancestries.addSuperclass( c.getName(), c.getClassFile().getSuperclass() );
+                                addMethodImplementations( c.getName(), (List<MethodInfo>)c.getClassFile().getMethods() );
+                        }
+                        catch( NotFoundException ex )
+                        {
+                                throw new Error( "Unable to load class: " + className );
+                        }
+                }
+        }
+        
+        private void addMethodImplementations( String name, List<MethodInfo> methods )
+        {
+                for( MethodInfo method : methods )
+                {
+                        m_methodImplementations.put( name, getMethodKey( method.getName(), method.getDescriptor() ) );
+                }
+        }
+        
+        public Ancestries getAncestries( )
+        {
+                return m_ancestries;
+        }
+        
+        public boolean isMethodImplemented( MethodEntry methodEntry )
+        {
+                return isMethodImplemented( methodEntry.getClassName(), methodEntry.getName(), methodEntry.getSignature() );
+        }
+        
+        public boolean isMethodImplemented( String className, String methodName, String methodSignature )
+        {
+                Collection<String> implementations = m_methodImplementations.get( className );
+                if( implementations == null )
+                {
+                        return false;
+                }
+                return implementations.contains( getMethodKey( methodName, methodSignature ) );
+        }
+        
+
+        public ClassInheritanceTreeNode getClassInheritance( Translator deobfuscatingTranslator, ClassEntry obfClassEntry )
+        {
+                // get the root node
+                List<String> ancestry = m_ancestries.getAncestry( obfClassEntry.getName() );
+                ClassInheritanceTreeNode rootNode = new ClassInheritanceTreeNode( deobfuscatingTranslator, ancestry.get( ancestry.size() - 1 ) );
+                
+                // expand all children recursively
+                rootNode.load( m_ancestries, true );
+                
+                return rootNode;
+        }
+        
+        public MethodInheritanceTreeNode getMethodInheritance( Translator deobfuscatingTranslator, MethodEntry obfMethodEntry )
+        {
+                // travel to the ancestor implementation
+                String baseImplementationClassName = obfMethodEntry.getClassName();
+                for( String ancestorClassName : m_ancestries.getAncestry( obfMethodEntry.getClassName() ) )
+                {
+                        if( isMethodImplemented( ancestorClassName, obfMethodEntry.getName(), obfMethodEntry.getSignature() ) )
+                        {
+                                baseImplementationClassName = ancestorClassName;
+                        }
+                }
+                
+                // make a root node at the base
+                MethodEntry methodEntry = new MethodEntry(
+                        new ClassEntry( baseImplementationClassName ),
+                        obfMethodEntry.getName(),
+                        obfMethodEntry.getSignature()
+                );
+                MethodInheritanceTreeNode rootNode = new MethodInheritanceTreeNode(
+                        deobfuscatingTranslator,
+                        methodEntry,
+                        isMethodImplemented( methodEntry )
+                );
+                
+                // expand the full tree
+                rootNode.load( this, true );
+                
+                return rootNode;
+        }
+        
+        private String getMethodKey( String name, String signature )
+        {
+                return name + signature;
+        }
+}
diff --git a/src/cuchaz/enigma/analysis/MethodInheritanceTreeNode.java b/src/cuchaz/enigma/analysis/MethodInheritanceTreeNode.java
index 1fecf48..a28a9f4 100644
--- a/src/cuchaz/enigma/analysis/MethodInheritanceTreeNode.java
+++ b/src/cuchaz/enigma/analysis/MethodInheritanceTreeNode.java
@@ -79,11 +79,11 @@ public class MethodInheritanceTreeNode extends DefaultMutableTreeNode
                 }
         }
         
-        public void load( Ancestries ancestries, boolean recurse )
+        public void load( JarIndex index, boolean recurse )
         {
                 // get all the child nodes
                 List<MethodInheritanceTreeNode> nodes = Lists.newArrayList();
-                for( String subclassName : ancestries.getSubclasses( m_entry.getClassName() ) )
+                for( String subclassName : index.getAncestries().getSubclasses( m_entry.getClassName() ) )
                 {
                         MethodEntry methodEntry = new MethodEntry(
                                 new ClassEntry( subclassName ),
@@ -93,7 +93,7 @@ public class MethodInheritanceTreeNode extends DefaultMutableTreeNode
                         nodes.add( new MethodInheritanceTreeNode(
                                 m_deobfuscatingTranslator,
                                 methodEntry,
-                                ancestries.isMethodImplemented( subclassName, m_entry.getName(), m_entry.getSignature() )
+                                index.isMethodImplemented( subclassName, m_entry.getName(), m_entry.getSignature() )
                         ) );
                 }
                 
@@ -107,7 +107,7 @@ public class MethodInheritanceTreeNode extends DefaultMutableTreeNode
                 {
                         for( MethodInheritanceTreeNode node : nodes )
                         {
-                                node.load( ancestries, true );
+                                node.load( index, true );
                         }
                 }
         }
diff --git a/src/cuchaz/enigma/gui/GuiController.java b/src/cuchaz/enigma/gui/GuiController.java
index 1946b4a..880f001 100644
--- a/src/cuchaz/enigma/gui/GuiController.java
+++ b/src/cuchaz/enigma/gui/GuiController.java
@@ -132,7 +132,7 @@ public class GuiController
         
         public ClassInheritanceTreeNode getClassInheritance( ClassEntry obfClassEntry )
         {
-                ClassInheritanceTreeNode rootNode = m_deobfuscator.getAncestries().getClassInheritance(
+                ClassInheritanceTreeNode rootNode = m_deobfuscator.getJarIndex().getClassInheritance(
                         m_deobfuscator.getTranslator( TranslationDirection.Deobfuscating ),
                         obfClassEntry
                 );
@@ -141,7 +141,7 @@ public class GuiController
         
         public MethodInheritanceTreeNode getMethodInheritance( MethodEntry obfMethodEntry )
         {
-                MethodInheritanceTreeNode rootNode = m_deobfuscator.getAncestries().getMethodInheritance(
+                MethodInheritanceTreeNode rootNode = m_deobfuscator.getJarIndex().getMethodInheritance(
                         m_deobfuscator.getTranslator( TranslationDirection.Deobfuscating ),
                         obfMethodEntry
                 );
diff --git a/src/cuchaz/enigma/mapping/Renamer.java b/src/cuchaz/enigma/mapping/Renamer.java
index 5a75c01..d372575 100644
--- a/src/cuchaz/enigma/mapping/Renamer.java
+++ b/src/cuchaz/enigma/mapping/Renamer.java
@@ -15,17 +15,17 @@ import java.io.ObjectOutputStream;
 import java.io.OutputStream;
 import java.util.zip.GZIPOutputStream;
 
-import cuchaz.enigma.analysis.Ancestries;
+import cuchaz.enigma.analysis.JarIndex;
 import cuchaz.enigma.analysis.MethodInheritanceTreeNode;
 
 public class Renamer
 {
-        private Ancestries m_ancestries;
+        private JarIndex m_index;
         private Mappings m_mappings;
         
-        public Renamer( Ancestries ancestries, Mappings mappings )
+        public Renamer( JarIndex index, Mappings mappings )
         {
-                m_ancestries = ancestries;
+                m_index = index;
                 m_mappings = mappings;
         }
         
@@ -61,7 +61,7 @@ public class Renamer
         {
                 // get the method tree
                 setMethodTreeName(
-                        m_ancestries.getMethodInheritance( m_mappings.getTranslator( m_ancestries, TranslationDirection.Deobfuscating ), obf ),
+                        m_index.getMethodInheritance( m_mappings.getTranslator( m_index.getAncestries(), TranslationDirection.Deobfuscating ), obf ),
                         deobfName
                 );
         }
@@ -90,7 +90,7 @@ public class Renamer
                         classMapping = createClassMapping( obf.getClassEntry() );
                 }
                 
-                String deobfSignature = m_mappings.getTranslator( m_ancestries, TranslationDirection.Deobfuscating ).translateSignature( obf.getSignature() );
+                String deobfSignature = m_mappings.getTranslator( m_index.getAncestries(), TranslationDirection.Deobfuscating ).translateSignature( obf.getSignature() );
                 classMapping.setMethodNameAndSignature( obf.getName(), obf.getSignature(), deobfName, deobfSignature );
                 
                 // TODO: update ancestor/descendant methods in other classes in the inheritance hierarchy too
@@ -129,7 +129,7 @@ public class Renamer
         
         private void updateDeobfMethodSignatures( )
         {
-                Translator translator = m_mappings.getTranslator( m_ancestries, TranslationDirection.Deobfuscating );
+                Translator translator = m_mappings.getTranslator( m_index.getAncestries(), TranslationDirection.Deobfuscating );
                 for( ClassMapping classMapping : m_mappings.m_classesByObf.values() )
                 {
                         classMapping.updateDeobfMethodSignatures( translator );