3 files changed, 37 insertions, 43 deletions

diff --git a/src/main/java/cuchaz/enigma/analysis/SourceIndexClassVisitor.java b/src/main/java/cuchaz/enigma/analysis/SourceIndexClassVisitor.java
index 1e2eed8..dd5bcef 100644
--- a/src/main/java/cuchaz/enigma/analysis/SourceIndexClassVisitor.java
+++ b/src/main/java/cuchaz/enigma/analysis/SourceIndexClassVisitor.java
@@ -17,16 +17,18 @@ import com.strobel.assembler.metadata.TypeDefinition;
 import com.strobel.assembler.metadata.TypeReference;
 import com.strobel.decompiler.languages.TextLocation;
 import com.strobel.decompiler.languages.java.ast.*;
+import cuchaz.enigma.bytecode.AccessFlags;
+import cuchaz.enigma.mapping.Signature;
 import cuchaz.enigma.mapping.entry.*;
 
 public class SourceIndexClassVisitor extends SourceIndexVisitor {
         private final ReferencedEntryPool entryPool;
         private final ProcyonEntryFactory entryFactory;
 
-        private ClassEntry classEntry;
+        private ClassDefEntry classEntry;
         private boolean isEnum;
 
-        public SourceIndexClassVisitor(ReferencedEntryPool entryPool, ClassEntry classEntry) {
+        public SourceIndexClassVisitor(ReferencedEntryPool entryPool, ClassDefEntry classEntry) {
                 super(entryPool);
                 this.entryPool = entryPool;
                 this.entryFactory = new ProcyonEntryFactory(entryPool);
@@ -37,7 +39,7 @@ public class SourceIndexClassVisitor extends SourceIndexVisitor {
         public Void visitTypeDeclaration(TypeDeclaration node, SourceIndex index) {
                 // is this this class, or a subtype?
                 TypeDefinition def = node.getUserData(Keys.TYPE_DEFINITION);
-                ClassEntry classEntry = new ClassEntry(def.getInternalName());
+                ClassDefEntry classEntry = new ClassDefEntry(def.getInternalName(), Signature.createSignature(def.getSignature()), new AccessFlags(def.getModifiers()));
                 if (!classEntry.equals(this.classEntry)) {
                         // it's a subtype, recurse
                         index.addDeclaration(node.getNameToken(), classEntry);
@@ -68,7 +70,7 @@ public class SourceIndexClassVisitor extends SourceIndexVisitor {
                         tokenNode = node.getModifiers().firstOrNullObject();
                 }
                 index.addDeclaration(tokenNode, methodEntry);
-                return node.acceptVisitor(new SourceIndexMethodVisitor(entryPool, methodEntry, false), index);
+                return node.acceptVisitor(new SourceIndexMethodVisitor(entryPool, classEntry, methodEntry), index);
         }
 
         @Override
@@ -76,7 +78,7 @@ public class SourceIndexClassVisitor extends SourceIndexVisitor {
                 MethodDefinition def = node.getUserData(Keys.METHOD_DEFINITION);
                 MethodDefEntry methodEntry = entryFactory.getMethodDefEntry(def);
                 index.addDeclaration(node.getNameToken(), methodEntry);
-                return node.acceptVisitor(new SourceIndexMethodVisitor(entryPool, methodEntry, isEnum), index);
+                return node.acceptVisitor(new SourceIndexMethodVisitor(entryPool, classEntry, methodEntry), index);
         }
 
         @Override
diff --git a/src/main/java/cuchaz/enigma/analysis/SourceIndexMethodVisitor.java b/src/main/java/cuchaz/enigma/analysis/SourceIndexMethodVisitor.java
index 01e773b..51195a4 100644
--- a/src/main/java/cuchaz/enigma/analysis/SourceIndexMethodVisitor.java
+++ b/src/main/java/cuchaz/enigma/analysis/SourceIndexMethodVisitor.java
@@ -13,15 +13,14 @@ package cuchaz.enigma.analysis;
 
 import com.google.common.collect.HashMultimap;
 import com.google.common.collect.Multimap;
-import com.strobel.assembler.metadata.MemberReference;
-import com.strobel.assembler.metadata.MethodReference;
-import com.strobel.assembler.metadata.ParameterDefinition;
-import com.strobel.assembler.metadata.TypeReference;
+import com.strobel.assembler.metadata.*;
+import com.strobel.decompiler.ast.Variable;
 import com.strobel.decompiler.languages.TextLocation;
 import com.strobel.decompiler.languages.java.ast.*;
-import cuchaz.enigma.mapping.*;
+import cuchaz.enigma.mapping.TypeDescriptor;
 import cuchaz.enigma.mapping.entry.*;
 
+import java.lang.Error;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -29,21 +28,18 @@ public class SourceIndexMethodVisitor extends SourceIndexVisitor {
         private final ReferencedEntryPool entryPool;
         private final ProcyonEntryFactory entryFactory;
 
-        private MethodDefEntry methodEntry;
+        private final ClassDefEntry ownerEntry;
+        private final MethodDefEntry methodEntry;
 
-        // TODO: Really fix Procyon index problem with inner classes
-        private int argumentPosition;
-        private int localsPosition;
         private Multimap<String, Identifier> unmatchedIdentifier = HashMultimap.create();
         private Map<String, Entry> identifierEntryCache = new HashMap<>();
 
-        public SourceIndexMethodVisitor(ReferencedEntryPool entryPool, MethodDefEntry methodEntry, boolean isEnum) {
+        public SourceIndexMethodVisitor(ReferencedEntryPool entryPool, ClassDefEntry ownerEntry, MethodDefEntry methodEntry) {
                 super(entryPool);
                 this.entryPool = entryPool;
                 this.entryFactory = new ProcyonEntryFactory(entryPool);
+                this.ownerEntry = ownerEntry;
                 this.methodEntry = methodEntry;
-                this.argumentPosition = isEnum ? 2 : 0;
-                this.localsPosition = 0;
         }
 
         @Override
@@ -112,14 +108,13 @@ public class SourceIndexMethodVisitor extends SourceIndexVisitor {
         @Override
         public Void visitParameterDeclaration(ParameterDeclaration node, SourceIndex index) {
                 ParameterDefinition def = node.getUserData(Keys.PARAMETER_DEFINITION);
-                if (def.getMethod() instanceof MemberReference && def.getMethod() instanceof MethodReference) {
-                        MethodEntry methodEntry = entryFactory.getMethodEntry((MethodReference) def.getMethod());
-                        LocalVariableEntry localVariableEntry = new LocalVariableEntry(methodEntry, argumentPosition++, node.getName());
-                        Identifier identifier = node.getNameToken();
-                        // cache the argument entry and the identifier
-                        identifierEntryCache.put(identifier.getName(), localVariableEntry);
-                        index.addDeclaration(identifier, localVariableEntry);
-                }
+
+                int variableOffset = this.methodEntry.getVariableOffset(ownerEntry);
+                LocalVariableEntry localVariableEntry = new LocalVariableEntry(methodEntry, def.getSlot() - variableOffset, node.getName());
+                Identifier identifier = node.getNameToken();
+                // cache the argument entry and the identifier
+                identifierEntryCache.put(identifier.getName(), localVariableEntry);
+                index.addDeclaration(identifier, localVariableEntry);
 
                 return recurse(node, index);
         }
@@ -173,18 +168,6 @@ public class SourceIndexMethodVisitor extends SourceIndexVisitor {
         }
 
         @Override
-        public Void visitForEachStatement(ForEachStatement node, SourceIndex index) {
-                if (node.getVariableType() instanceof SimpleType) {
-                        Identifier identifier = node.getVariableNameToken();
-                        LocalVariableEntry localVariableEntry = new LocalVariableEntry(methodEntry, argumentPosition + localsPosition++, identifier.getName());
-                        identifierEntryCache.put(identifier.getName(), localVariableEntry);
-                        addDeclarationToUnmatched(identifier.getName(), index);
-                        index.addDeclaration(identifier, localVariableEntry);
-                }
-                return recurse(node, index);
-        }
-
-        @Override
         public Void visitVariableDeclaration(VariableDeclarationStatement node, SourceIndex index) {
                 AstNodeCollection<VariableInitializer> variables = node.getVariables();
 
@@ -193,10 +176,17 @@ public class SourceIndexMethodVisitor extends SourceIndexVisitor {
                         VariableInitializer initializer = variables.firstOrNullObject();
                         if (initializer != null && node.getType() instanceof SimpleType) {
                                 Identifier identifier = initializer.getNameToken();
-                                LocalVariableEntry localVariableEntry = new LocalVariableEntry(methodEntry, argumentPosition + localsPosition++, initializer.getName());
-                                identifierEntryCache.put(identifier.getName(), localVariableEntry);
-                                addDeclarationToUnmatched(identifier.getName(), index);
-                                index.addDeclaration(identifier, localVariableEntry);
+                                Variable variable = initializer.getUserData(Keys.VARIABLE);
+                                if (variable != null) {
+                                        VariableDefinition originalVariable = variable.getOriginalVariable();
+                                        if (originalVariable != null) {
+                                                int variableOffset = methodEntry.getVariableOffset(ownerEntry);
+                                                LocalVariableEntry localVariableEntry = new LocalVariableEntry(methodEntry, originalVariable.getSlot() - variableOffset, initializer.getName());
+                                                identifierEntryCache.put(identifier.getName(), localVariableEntry);
+                                                addDeclarationToUnmatched(identifier.getName(), index);
+                                                index.addDeclaration(identifier, localVariableEntry);
+                                        }
+                                }
                         }
                 }
                 return recurse(node, index);
diff --git a/src/main/java/cuchaz/enigma/analysis/SourceIndexVisitor.java b/src/main/java/cuchaz/enigma/analysis/SourceIndexVisitor.java
index 176f283..e588d24 100644
--- a/src/main/java/cuchaz/enigma/analysis/SourceIndexVisitor.java
+++ b/src/main/java/cuchaz/enigma/analysis/SourceIndexVisitor.java
@@ -14,7 +14,9 @@ package cuchaz.enigma.analysis;
 import com.strobel.assembler.metadata.TypeDefinition;
 import com.strobel.decompiler.languages.java.ast.*;
 import com.strobel.decompiler.patterns.Pattern;
-import cuchaz.enigma.mapping.entry.ClassEntry;
+import cuchaz.enigma.bytecode.AccessFlags;
+import cuchaz.enigma.mapping.Signature;
+import cuchaz.enigma.mapping.entry.ClassDefEntry;
 import cuchaz.enigma.mapping.entry.ReferencedEntryPool;
 
 public class SourceIndexVisitor implements IAstVisitor<SourceIndex, Void> {
@@ -27,7 +29,7 @@ public class SourceIndexVisitor implements IAstVisitor<SourceIndex, Void> {
         @Override
         public Void visitTypeDeclaration(TypeDeclaration node, SourceIndex index) {
                 TypeDefinition def = node.getUserData(Keys.TYPE_DEFINITION);
-                ClassEntry classEntry = new ClassEntry(def.getInternalName());
+                ClassDefEntry classEntry = new ClassDefEntry(def.getInternalName(), Signature.createSignature(def.getSignature()), new AccessFlags(def.getModifiers()));
                 index.addDeclaration(node.getNameToken(), classEntry);
 
                 return node.acceptVisitor(new SourceIndexClassVisitor(entryPool, classEntry), index);