added tests to check constructor tokens

fixed a bug with constructor tokens too

10 files changed, 316 insertions, 24 deletions

diff --git a/src/cuchaz/enigma/analysis/SourceIndex.java b/src/cuchaz/enigma/analysis/SourceIndex.java
index 1a5a80d6..b777f9fd 100644
--- a/src/cuchaz/enigma/analysis/SourceIndex.java
+++ b/src/cuchaz/enigma/analysis/SourceIndex.java
@@ -20,7 +20,7 @@ import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Multimap;
 import com.strobel.decompiler.languages.Region;
-import com.strobel.decompiler.languages.java.ast.Identifier;
+import com.strobel.decompiler.languages.java.ast.AstNode;
 
 import cuchaz.enigma.mapping.Entry;
 
@@ -56,7 +56,7 @@ public class SourceIndex
                 return m_source;
         }
         
-        public Token getToken( Identifier node )
+        public Token getToken( AstNode node )
         {
                 // get a token for this node's region
                 Region region = node.getRegion();
@@ -71,7 +71,7 @@ public class SourceIndex
                 );
                 
                 // for tokens representing inner classes, make sure we only get the simple name
-                int pos = node.getName().lastIndexOf( '$' );
+                int pos = node.toString().lastIndexOf( '$' );
                 if( pos >= 0 )
                 {
                         token.end -= pos + 1;
@@ -92,7 +92,7 @@ public class SourceIndex
                 return token;
         }
         
-        public void addReference( Identifier node, EntryReference<Entry,Entry> deobfReference )
+        public void addReference( AstNode node, EntryReference<Entry,Entry> deobfReference )
         {
                 Token token = getToken( node );
                 if( token != null )
@@ -102,7 +102,7 @@ public class SourceIndex
                 }
         }
         
-        public void addDeclaration( Identifier node, Entry deobfEntry )
+        public void addDeclaration( AstNode node, Entry deobfEntry )
         {
                 Token token = getToken( node );
                 if( token != null )
diff --git a/src/cuchaz/enigma/analysis/SourceIndexBehaviorVisitor.java b/src/cuchaz/enigma/analysis/SourceIndexBehaviorVisitor.java
index ab505528..a1dac4b5 100644
--- a/src/cuchaz/enigma/analysis/SourceIndexBehaviorVisitor.java
+++ b/src/cuchaz/enigma/analysis/SourceIndexBehaviorVisitor.java
@@ -12,9 +12,11 @@ package cuchaz.enigma.analysis;
 
 import com.strobel.assembler.metadata.MemberReference;
 import com.strobel.assembler.metadata.MethodDefinition;
+import com.strobel.assembler.metadata.MethodReference;
 import com.strobel.assembler.metadata.ParameterDefinition;
 import com.strobel.assembler.metadata.TypeReference;
 import com.strobel.decompiler.languages.TextLocation;
+import com.strobel.decompiler.languages.java.ast.AstNode;
 import com.strobel.decompiler.languages.java.ast.ConstructorDeclaration;
 import com.strobel.decompiler.languages.java.ast.IdentifierExpression;
 import com.strobel.decompiler.languages.java.ast.InvocationExpression;
@@ -24,6 +26,8 @@ import com.strobel.decompiler.languages.java.ast.MethodDeclaration;
 import com.strobel.decompiler.languages.java.ast.ObjectCreationExpression;
 import com.strobel.decompiler.languages.java.ast.ParameterDeclaration;
 import com.strobel.decompiler.languages.java.ast.SimpleType;
+import com.strobel.decompiler.languages.java.ast.SuperReferenceExpression;
+import com.strobel.decompiler.languages.java.ast.ThisReferenceExpression;
 
 import cuchaz.enigma.mapping.ArgumentEntry;
 import cuchaz.enigma.mapping.BehaviorEntry;
@@ -58,14 +62,49 @@ public class SourceIndexBehaviorVisitor extends SourceIndexVisitor
         public Void visitInvocationExpression( InvocationExpression node, SourceIndex index )
         {
                 MemberReference ref = node.getUserData( Keys.MEMBER_REFERENCE );
+                
+                // get the behavior entry
                 ClassEntry classEntry = new ClassEntry( ref.getDeclaringType().getInternalName() );
-                if( node.getTarget() instanceof MemberReferenceExpression )
+                BehaviorEntry behaviorEntry = null;
+                if( ref instanceof MethodReference )
                 {
-                        MethodEntry methodEntry = new MethodEntry( classEntry, ref.getName(), ref.getSignature() );
-                        index.addReference(
-                                ((MemberReferenceExpression)node.getTarget()).getMemberNameToken(),
-                                new EntryReference<Entry,Entry>( methodEntry, m_behaviorEntry )
-                        );
+                        MethodReference methodRef = (MethodReference)ref;
+                        if( methodRef.isConstructor() )
+                        {
+                                behaviorEntry = new ConstructorEntry( classEntry, ref.getSignature() );
+                        }
+                        else if( methodRef.isTypeInitializer() )
+                        {
+                                behaviorEntry = new ConstructorEntry( classEntry );
+                        }
+                        else
+                        {
+                                behaviorEntry = new MethodEntry( classEntry, ref.getName(), ref.getSignature() );
+                        }
+                }
+                if( behaviorEntry != null )
+                {
+                        // get the node for the token
+                        AstNode tokenNode = null;
+                        if( node.getTarget() instanceof MemberReferenceExpression )
+                        {
+                                tokenNode = ((MemberReferenceExpression)node.getTarget()).getMemberNameToken();
+                        }
+                        else if( node.getTarget() instanceof SuperReferenceExpression )
+                        {
+                                tokenNode = node.getTarget();
+                        }
+                        else if( node.getTarget() instanceof ThisReferenceExpression )
+                        {
+                                tokenNode = node.getTarget();
+                        }
+                        if( tokenNode != null )
+                        {
+                                index.addReference(
+                                        tokenNode,
+                                        new EntryReference<Entry,Entry>( behaviorEntry, m_behaviorEntry )
+                                );
+                        }
                 }
                 
                 return recurse( node, index );
diff --git a/src/cuchaz/enigma/analysis/SourceIndexClassVisitor.java b/src/cuchaz/enigma/analysis/SourceIndexClassVisitor.java
index a1c82711..b7897268 100644
--- a/src/cuchaz/enigma/analysis/SourceIndexClassVisitor.java
+++ b/src/cuchaz/enigma/analysis/SourceIndexClassVisitor.java
@@ -97,7 +97,6 @@ public class SourceIndexClassVisitor extends SourceIndexVisitor
                 ClassEntry classEntry = new ClassEntry( def.getDeclaringType().getInternalName() );
                 ConstructorEntry constructorEntry = new ConstructorEntry( classEntry, def.getSignature() );
                 index.addDeclaration( node.getNameToken(), constructorEntry );
-                
                 return node.acceptVisitor( new SourceIndexBehaviorVisitor( constructorEntry ), index );
         }
         
diff --git a/test/cuchaz/enigma/TestJarIndexConstructorReferences.java b/test/cuchaz/enigma/TestJarIndexConstructorReferences.java
index 38882a0c..c0691888 100644
--- a/test/cuchaz/enigma/TestJarIndexConstructorReferences.java
+++ b/test/cuchaz/enigma/TestJarIndexConstructorReferences.java
@@ -18,10 +18,13 @@ import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.is;
 
+import java.io.File;
+import java.util.Collection;
 import java.util.jar.JarFile;
 
 import org.junit.Test;
 
+import cuchaz.enigma.analysis.EntryReference;
 import cuchaz.enigma.analysis.JarIndex;
 import cuchaz.enigma.mapping.BehaviorEntry;
 import cuchaz.enigma.mapping.ClassEntry;
@@ -30,17 +33,19 @@ import cuchaz.enigma.mapping.ConstructorEntry;
 public class TestJarIndexConstructorReferences
 {
         private JarIndex m_index;
-        
+                
         private ClassEntry m_baseClass = new ClassEntry( "none/a" );
-        private ClassEntry m_subClass = new ClassEntry( "none/c" );
-        private ClassEntry m_subsubClass = new ClassEntry( "none/d" );
+        private ClassEntry m_subClass = new ClassEntry( "none/d" );
+        private ClassEntry m_subsubClass = new ClassEntry( "none/e" );
+        private ClassEntry m_defaultClass = new ClassEntry( "none/c" );
         private ClassEntry m_callerClass = new ClassEntry( "none/b" );
-        
+
         public TestJarIndexConstructorReferences( )
         throws Exception
         {
+                File jarFile = new File( "build/libs/testConstructors.obf.jar" );
                 m_index = new JarIndex();
-                m_index.indexJar( new JarFile( "build/libs/testConstructors.obf.jar" ), false );
+                m_index.indexJar( new JarFile( jarFile ), false );
         }
         
         @Test
@@ -51,6 +56,7 @@ public class TestJarIndexConstructorReferences
                         m_baseClass,
                         m_subClass,
                         m_subsubClass,
+                        m_defaultClass,
                         m_callerClass
                 ) );
         }
@@ -60,7 +66,8 @@ public class TestJarIndexConstructorReferences
         public void baseDefault( )
         {
                 BehaviorEntry source = new ConstructorEntry( m_baseClass, "()V" );
-                assertThat( m_index.getBehaviorReferences( source ), containsInAnyOrder(
+                Collection<EntryReference<BehaviorEntry,BehaviorEntry>> references = m_index.getBehaviorReferences( source );
+                assertThat( references, containsInAnyOrder(
                         newBehaviorReferenceByMethod( source, m_callerClass.getName(), "a", "()V" ),
                         newBehaviorReferenceByConstructor( source, m_subClass.getName(), "()V" ),
                         newBehaviorReferenceByConstructor( source, m_subClass.getName(), "(III)V" )
@@ -126,4 +133,14 @@ public class TestJarIndexConstructorReferences
                         newBehaviorReferenceByMethod( source, m_callerClass.getName(), "f", "()V" )
                 ) );
         }
+
+        @Test
+        @SuppressWarnings( "unchecked" )
+        public void defaultConstructable( )
+        {
+                BehaviorEntry source = new ConstructorEntry( m_defaultClass, "()V" );
+                assertThat( m_index.getBehaviorReferences( source ), containsInAnyOrder(
+                        newBehaviorReferenceByMethod( source, m_callerClass.getName(), "g", "()V" )
+                ) );
+        }
 }
diff --git a/test/cuchaz/enigma/TestTokensConstructors.java b/test/cuchaz/enigma/TestTokensConstructors.java
new file mode 100644
index 00000000..24091532
--- /dev/null
+++ b/test/cuchaz/enigma/TestTokensConstructors.java
@@ -0,0 +1,152 @@
+/*******************************************************************************
+ * 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;
+
+import static cuchaz.enigma.EntryFactory.newBehaviorReferenceByConstructor;
+import static cuchaz.enigma.EntryFactory.newBehaviorReferenceByMethod;
+import static cuchaz.enigma.EntryFactory.newConstructor;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.nullValue;
+
+import java.io.File;
+
+import org.junit.Test;
+
+import cuchaz.enigma.mapping.BehaviorEntry;
+
+public class TestTokensConstructors extends TokenChecker
+{
+        public TestTokensConstructors( )
+        throws Exception
+        {
+                super( new File( "build/libs/testConstructors.obf.jar" ) );
+        }
+        
+        @Test
+        public void baseDeclarations( )
+        {
+                assertThat( getDeclarationToken( newConstructor( "none/a", "()V" ) ), is( "a" ) );
+                assertThat( getDeclarationToken( newConstructor( "none/a", "(I)V" ) ), is( "a" ) );
+        }
+        
+        @Test
+        public void subDeclarations( )
+        {
+                assertThat( getDeclarationToken( newConstructor( "none/d", "()V" ) ), is( "d" ) );
+                assertThat( getDeclarationToken( newConstructor( "none/d", "(I)V" ) ), is( "d" ) );
+                assertThat( getDeclarationToken( newConstructor( "none/d", "(II)V" ) ), is( "d" ) );
+                assertThat( getDeclarationToken( newConstructor( "none/d", "(III)V" ) ), is( "d" ) );
+        }
+        
+        @Test
+        public void subsubDeclarations( )
+        {
+                assertThat( getDeclarationToken( newConstructor( "none/e", "(I)V" ) ), is( "e" ) );
+        }
+        
+        @Test
+        public void defaultDeclarations( )
+        {
+                assertThat( getDeclarationToken( newConstructor( "none/c", "()V" ) ), nullValue() );
+        }
+        
+        @Test
+        public void baseDefaultReferences( )
+        {
+                BehaviorEntry source = newConstructor( "none/a", "()V" );
+                assertThat(
+                        getReferenceTokens( newBehaviorReferenceByMethod( source, "none/b", "a", "()V" ) ),
+                        containsInAnyOrder( "a" )
+                );
+                assertThat(
+                        getReferenceTokens( newBehaviorReferenceByConstructor( source, "none/d", "()V" ) ),
+                        containsInAnyOrder( "super" ) // implicit call, decompiled to "super"
+                );
+                assertThat(
+                        getReferenceTokens( newBehaviorReferenceByConstructor( source, "none/d", "(III)V" ) ),
+                        containsInAnyOrder( "super" ) // implicit call, decompiled to "super"
+                );
+        }
+        
+        @Test
+        public void baseIntReferences( )
+        {
+                BehaviorEntry source = newConstructor( "none/a", "(I)V" );
+                assertThat(
+                        getReferenceTokens( newBehaviorReferenceByMethod( source, "none/b", "b", "()V" ) ),
+                        containsInAnyOrder( "a" )
+                );
+        }
+
+        @Test
+        public void subDefaultReferences( )
+        {
+                BehaviorEntry source = newConstructor( "none/d", "()V" );
+                assertThat(
+                        getReferenceTokens( newBehaviorReferenceByMethod( source, "none/b", "c", "()V" ) ),
+                        containsInAnyOrder( "d" )
+                );
+                assertThat(
+                        getReferenceTokens( newBehaviorReferenceByConstructor( source, "none/d", "(I)V" ) ),
+                        containsInAnyOrder( "this" )
+                );
+        }
+
+        @Test
+        public void subIntReferences( )
+        {
+                BehaviorEntry source = newConstructor( "none/d", "(I)V" );
+                assertThat(
+                        getReferenceTokens( newBehaviorReferenceByMethod( source, "none/b", "d", "()V" ) ),
+                        containsInAnyOrder( "d" )
+                );
+                assertThat(
+                        getReferenceTokens( newBehaviorReferenceByConstructor( source, "none/d", "(II)V" ) ),
+                        containsInAnyOrder( "this" )
+                );
+                assertThat(
+                        getReferenceTokens( newBehaviorReferenceByConstructor( source, "none/e", "(I)V" ) ),
+                        containsInAnyOrder( "super" )
+                );
+        }
+
+        @Test
+        public void subIntIntReferences( )
+        {
+                BehaviorEntry source = newConstructor( "none/d", "(II)V" );
+                assertThat(
+                        getReferenceTokens( newBehaviorReferenceByMethod( source, "none/b", "e", "()V" ) ),
+                        containsInAnyOrder( "d" )
+                );
+        }
+
+        @Test
+        public void subsubIntReferences( )
+        {
+                BehaviorEntry source = newConstructor( "none/e", "(I)V" );
+                assertThat(
+                        getReferenceTokens( newBehaviorReferenceByMethod( source, "none/b", "f", "()V" ) ),
+                        containsInAnyOrder( "e" )
+                );
+        }
+
+        @Test
+        public void defaultConstructableReferences( )
+        {
+                BehaviorEntry source = newConstructor( "none/c", "()V" );
+                assertThat(
+                        getReferenceTokens( newBehaviorReferenceByMethod( source, "none/b", "g", "()V" ) ),
+                        containsInAnyOrder( "c" )
+                );
+        }
+}
diff --git a/test/cuchaz/enigma/TokenChecker.java b/test/cuchaz/enigma/TokenChecker.java
new file mode 100644
index 00000000..07ab9eca
--- /dev/null
+++ b/test/cuchaz/enigma/TokenChecker.java
@@ -0,0 +1,71 @@
+/*******************************************************************************
+ * 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;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.Collection;
+import java.util.List;
+
+import com.beust.jcommander.internal.Lists;
+import com.strobel.decompiler.languages.java.ast.CompilationUnit;
+
+import cuchaz.enigma.analysis.EntryReference;
+import cuchaz.enigma.analysis.SourceIndex;
+import cuchaz.enigma.analysis.Token;
+import cuchaz.enigma.analysis.TreeDumpVisitor;
+import cuchaz.enigma.mapping.Entry;
+
+public class TokenChecker
+{
+        private Deobfuscator m_deobfuscator;
+                
+        protected TokenChecker( File jarFile )
+        throws IOException
+        {
+                m_deobfuscator = new Deobfuscator( jarFile );
+        }
+        
+        protected String getDeclarationToken( Entry entry )
+        {
+                // decompile the class
+                CompilationUnit tree = m_deobfuscator.getSourceTree( entry.getClassName() );
+                // DEBUG
+                //tree.acceptVisitor( new TreeDumpVisitor( new File( "tree." + entry.getClassName().replace( '/', '.' ) + ".txt" ) ), null );
+                String source = m_deobfuscator.getSource( tree );
+                SourceIndex index = m_deobfuscator.getSourceIndex( tree, source );
+                
+                // get the token value
+                Token token = index.getDeclarationToken( entry );
+                if( token == null )
+                {
+                        return null;
+                }
+                return source.substring( token.start, token.end );
+        }
+        
+        @SuppressWarnings( "unchecked" )
+        protected Collection<String> getReferenceTokens( EntryReference<? extends Entry,? extends Entry> reference )
+        {
+                // decompile the class
+                CompilationUnit tree = m_deobfuscator.getSourceTree( reference.context.getClassName() );
+                String source = m_deobfuscator.getSource( tree );
+                SourceIndex index = m_deobfuscator.getSourceIndex( tree, source );
+                
+                // get the token values
+                List<String> values = Lists.newArrayList();
+                for( Token token : index.getReferenceTokens( (EntryReference<Entry,Entry>)reference ) )
+                {
+                        values.add( source.substring( token.start, token.end ) );
+                }
+                return values;
+        }
+}
diff --git a/test/cuchaz/enigma/inputs/constructors/Caller.java b/test/cuchaz/enigma/inputs/constructors/Caller.java
index f356b1b3..b2186198 100644
--- a/test/cuchaz/enigma/inputs/constructors/Caller.java
+++ b/test/cuchaz/enigma/inputs/constructors/Caller.java
@@ -20,28 +20,35 @@ public class Caller
         // c()V
         public void callSubDefault( )
         {
-                // none/c.<init>()V
+                // none/d.<init>()V
                 System.out.println( new SubClass() );
         }
 
         // d()V
         public void callSubInt( )
         {
-                // none/c.<init>(I)V
+                // none/d.<init>(I)V
                 System.out.println( new SubClass( 6 ) );
         }
 
         // e()V
         public void callSubIntInt( )
         {
-                // none/c.<init>(II)V
+                // none/d.<init>(II)V
                 System.out.println( new SubClass( 4, 2 ) );
         }
 
         // f()V
         public void callSubSubInt( )
         {
-                // none/d.<init>(I)V
+                // none/e.<init>(I)V
                 System.out.println( new SubSubClass( 3 ) );
         }
+        
+        // g()V
+        public void callDefaultConstructable()
+        {
+                // none/c.<init>()V
+                System.out.println( new DefaultConstructable() );
+        }
 }
diff --git a/test/cuchaz/enigma/inputs/constructors/DefaultConstructable.java b/test/cuchaz/enigma/inputs/constructors/DefaultConstructable.java
new file mode 100644
index 00000000..6cfd35e6
--- /dev/null
+++ b/test/cuchaz/enigma/inputs/constructors/DefaultConstructable.java
@@ -0,0 +1,6 @@
+package cuchaz.enigma.inputs.constructors;
+
+public class DefaultConstructable
+{
+        // only default constructor
+}
diff --git a/test/cuchaz/enigma/inputs/constructors/SubClass.java b/test/cuchaz/enigma/inputs/constructors/SubClass.java
index 2235de37..6ef77323 100644
--- a/test/cuchaz/enigma/inputs/constructors/SubClass.java
+++ b/test/cuchaz/enigma/inputs/constructors/SubClass.java
@@ -1,6 +1,6 @@
 package cuchaz.enigma.inputs.constructors;
 
-// none/c extends none/a
+// none/d extends none/a
 public class SubClass extends BaseClass
 {
         // <init>()V
@@ -24,6 +24,7 @@ public class SubClass extends BaseClass
                 this( a + b );
         }
         
+        // <init>(III)V
         public SubClass( int a, int b, int c )
         {
                 // none/a.<init>()V
diff --git a/test/cuchaz/enigma/inputs/constructors/SubSubClass.java b/test/cuchaz/enigma/inputs/constructors/SubSubClass.java
index a9445d8b..76a0f1f3 100644
--- a/test/cuchaz/enigma/inputs/constructors/SubSubClass.java
+++ b/test/cuchaz/enigma/inputs/constructors/SubSubClass.java
@@ -1,6 +1,6 @@
 package cuchaz.enigma.inputs.constructors;
 
-// none/d extends none/c
+// none/e extends none/d
 public class SubSubClass extends SubClass
 {
         // <init>(I)V