added tests to check constructor tokens

fixed a bug with constructor tokens too

3 files changed, 50 insertions, 12 deletions

diff --git a/src/cuchaz/enigma/analysis/SourceIndex.java b/src/cuchaz/enigma/analysis/SourceIndex.java
index 1a5a80d..b777f9f 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 ab50552..a1dac4b 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 a1c8271..b789726 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 );
         }