Fix parsed local variable indexing

5 files changed, 54 insertions, 55 deletions

diff --git a/src/main/java/cuchaz/enigma/analysis/SourceIndexClassVisitor.java b/src/main/java/cuchaz/enigma/analysis/SourceIndexClassVisitor.java
index 1e2eed8d..dd5bcef0 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 01e773b3..51195a4d 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 176f2838..e588d24b 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);
diff --git a/src/main/java/cuchaz/enigma/bytecode/translators/TranslationMethodVisitor.java b/src/main/java/cuchaz/enigma/bytecode/translators/TranslationMethodVisitor.java
index 82a1263b..a6c7552a 100644
--- a/src/main/java/cuchaz/enigma/bytecode/translators/TranslationMethodVisitor.java
+++ b/src/main/java/cuchaz/enigma/bytecode/translators/TranslationMethodVisitor.java
@@ -86,12 +86,7 @@ public class TranslationMethodVisitor extends MethodVisitor {
 
                 String translatedSignature = translator.getTranslatedSignature(Signature.createTypedSignature(signature)).toString();
 
-                // If we're not static, "this" is bound to index 0
-                int offset = methodEntry.getAccess().isStatic() ? 0 : 1;
-                if (ownerEntry.getAccess().isEnum() && methodEntry.getName().startsWith("<")) {
-                        // Enum constructors have an implicit "name" and "ordinal" parameter as well as "this"
-                        offset = 3;
-                }
+                int offset = methodEntry.getVariableOffset(ownerEntry);
 
                 int offsetIndex = index - offset;
                 if (offsetIndex >= 0) {
@@ -101,8 +96,8 @@ public class TranslationMethodVisitor extends MethodVisitor {
 
                         // TODO: Better name inference
                         if (translatedName.equals(entry.getName())) {
-                                String prefix = offsetIndex < methodEntry.getDesc().getArgumentDescs().size() ? "a" : "v";
-                                translatedName = inferName(prefix, offsetIndex, translatedEntry.getDesc());
+                                boolean argument = offsetIndex < methodEntry.getDesc().getArgumentDescs().size();
+                                translatedName = inferName(argument, offsetIndex, translatedEntry.getDesc());
                         }
 
                         super.visitLocalVariable(translatedName, translatedEntry.getDesc().toString(), translatedSignature, start, end, index);
@@ -159,7 +154,7 @@ public class TranslationMethodVisitor extends MethodVisitor {
                                 LocalVariableEntry translatedEntry = translator.getTranslatedVariable(entry);
                                 String translatedName = translatedEntry.getName();
                                 if (translatedName.equals(entry.getName())) {
-                                        super.visitParameter(inferName("a", index, arguments.get(index)), 0);
+                                        super.visitParameter(inferName(true, index, arguments.get(index)), 0);
                                 } else {
                                         super.visitParameter(translatedName, 0);
                                 }
@@ -168,10 +163,10 @@ public class TranslationMethodVisitor extends MethodVisitor {
                 super.visitEnd();
         }
 
-        private String inferName(String prefix, int argumentIndex, TypeDescriptor desc) {
+        private String inferName(boolean argument, int argumentIndex, TypeDescriptor desc) {
                 String translatedName;
                 int nameIndex = argumentIndex + 1;
-                StringBuilder nameBuilder = new StringBuilder(prefix);
+                StringBuilder nameBuilder = new StringBuilder(argument ? "a" : "v");
                 // Unfortunately each of these have different name getters, so they have different code paths
                 if (desc.isPrimitive()) {
                         TypeDescriptor.Primitive argCls = desc.getPrimitive();
@@ -184,7 +179,7 @@ public class TranslationMethodVisitor extends MethodVisitor {
                         typeName = typeName.substring(0, 1).toUpperCase(Locale.ROOT) + typeName.substring(1);
                         nameBuilder.append(typeName);
                 }
-                if (methodEntry.getDesc().getArgumentDescs().size() > 1) {
+                if (!argument || methodEntry.getDesc().getArgumentDescs().size() > 1) {
                         nameBuilder.append(nameIndex);
                 }
                 translatedName = nameBuilder.toString();
diff --git a/src/main/java/cuchaz/enigma/mapping/entry/MethodDefEntry.java b/src/main/java/cuchaz/enigma/mapping/entry/MethodDefEntry.java
index ec3af694..bb7c85eb 100644
--- a/src/main/java/cuchaz/enigma/mapping/entry/MethodDefEntry.java
+++ b/src/main/java/cuchaz/enigma/mapping/entry/MethodDefEntry.java
@@ -37,6 +37,16 @@ public class MethodDefEntry extends MethodEntry {
                 return signature;
         }
 
+        public int getVariableOffset(ClassDefEntry ownerEntry) {
+                // Enum constructors have an implicit "name" and "ordinal" parameter as well as "this"
+                if (ownerEntry.getAccess().isEnum() && getName().startsWith("<")) {
+                        return 3;
+                } else {
+                        // If we're not static, "this" is bound to index 0
+                        return getAccess().isStatic() ? 0 : 1;
+                }
+        }
+
         @Override
         public MethodDefEntry updateOwnership(ClassEntry classEntry) {
                 return new MethodDefEntry(new ClassEntry(classEntry.getName()), name, descriptor, signature, access);