Index lambda local variables to correct declaring method

2 files changed, 12 insertions, 9 deletions

diff --git a/src/main/java/cuchaz/enigma/analysis/SourceIndexClassVisitor.java b/src/main/java/cuchaz/enigma/analysis/SourceIndexClassVisitor.java
index a4fe9ee..c977703 100644
--- a/src/main/java/cuchaz/enigma/analysis/SourceIndexClassVisitor.java
+++ b/src/main/java/cuchaz/enigma/analysis/SourceIndexClassVisitor.java
@@ -17,18 +17,15 @@ 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.translation.representation.ProcyonEntryFactory;
 import cuchaz.enigma.translation.representation.entry.ClassDefEntry;
 import cuchaz.enigma.translation.representation.entry.ClassEntry;
 import cuchaz.enigma.translation.representation.entry.FieldDefEntry;
 import cuchaz.enigma.translation.representation.entry.MethodDefEntry;
 
 public class SourceIndexClassVisitor extends SourceIndexVisitor {
-        private final ProcyonEntryFactory entryFactory;
         private ClassDefEntry classEntry;
 
         public SourceIndexClassVisitor(ClassDefEntry classEntry) {
-                this.entryFactory = new ProcyonEntryFactory();
                 this.classEntry = classEntry;
         }
 
@@ -60,7 +57,7 @@ public class SourceIndexClassVisitor extends SourceIndexVisitor {
         @Override
         public Void visitMethodDeclaration(MethodDeclaration node, SourceIndex index) {
                 MethodDefinition def = node.getUserData(Keys.METHOD_DEFINITION);
-                MethodDefEntry methodEntry = entryFactory.getMethodDefEntry(def);
+                MethodDefEntry methodEntry = MethodDefEntry.parse(def);
                 AstNode tokenNode = node.getNameToken();
                 if (methodEntry.isConstructor() && methodEntry.getName().equals("<clinit>")) {
                         // for static initializers, check elsewhere for the token node
@@ -73,7 +70,7 @@ public class SourceIndexClassVisitor extends SourceIndexVisitor {
         @Override
         public Void visitConstructorDeclaration(ConstructorDeclaration node, SourceIndex index) {
                 MethodDefinition def = node.getUserData(Keys.METHOD_DEFINITION);
-                MethodDefEntry methodEntry = entryFactory.getMethodDefEntry(def);
+                MethodDefEntry methodEntry = MethodDefEntry.parse(def);
                 index.addDeclaration(node.getNameToken(), methodEntry);
                 return node.acceptVisitor(new SourceIndexMethodVisitor(methodEntry), index);
         }
@@ -81,7 +78,7 @@ public class SourceIndexClassVisitor extends SourceIndexVisitor {
         @Override
         public Void visitFieldDeclaration(FieldDeclaration node, SourceIndex index) {
                 FieldDefinition def = node.getUserData(Keys.FIELD_DEFINITION);
-                FieldDefEntry fieldEntry = entryFactory.getFieldDefEntry(def);
+                FieldDefEntry fieldEntry = FieldDefEntry.parse(def);
                 assert (node.getVariables().size() == 1);
                 VariableInitializer variable = node.getVariables().firstOrNullObject();
                 index.addDeclaration(variable.getNameToken(), fieldEntry);
@@ -92,7 +89,7 @@ public class SourceIndexClassVisitor extends SourceIndexVisitor {
         public Void visitEnumValueDeclaration(EnumValueDeclaration node, SourceIndex index) {
                 // treat enum declarations as field declarations
                 FieldDefinition def = node.getUserData(Keys.FIELD_DEFINITION);
-                FieldDefEntry fieldEntry = entryFactory.getFieldDefEntry(def);
+                FieldDefEntry fieldEntry = FieldDefEntry.parse(def);
                 index.addDeclaration(node.getNameToken(), fieldEntry);
                 return recurse(node, index);
         }
diff --git a/src/main/java/cuchaz/enigma/analysis/SourceIndexMethodVisitor.java b/src/main/java/cuchaz/enigma/analysis/SourceIndexMethodVisitor.java
index c4785b6..fde6edf 100644
--- a/src/main/java/cuchaz/enigma/analysis/SourceIndexMethodVisitor.java
+++ b/src/main/java/cuchaz/enigma/analysis/SourceIndexMethodVisitor.java
@@ -101,8 +101,13 @@ public class SourceIndexMethodVisitor extends SourceIndexVisitor {
                 int parameterIndex = def.getSlot();
 
                 if (parameterIndex >= 0) {
+                        MethodDefEntry ownerMethod = methodEntry;
+                        if (def.getMethod() instanceof MethodDefinition) {
+                                ownerMethod = MethodDefEntry.parse((MethodDefinition) def.getMethod());
+                        }
+
                         TypeDescriptor parameterType = TypeDescriptor.parse(def.getParameterType());
-                        LocalVariableDefEntry localVariableEntry = new LocalVariableDefEntry(methodEntry, parameterIndex, node.getName(), true, parameterType);
+                        LocalVariableDefEntry localVariableEntry = new LocalVariableDefEntry(ownerMethod, parameterIndex, node.getName(), true, parameterType);
                         Identifier identifier = node.getNameToken();
                         // cache the argument entry and the identifier
                         identifierEntryCache.put(identifier.getName(), localVariableEntry);
@@ -170,8 +175,9 @@ public class SourceIndexMethodVisitor extends SourceIndexVisitor {
                                         if (originalVariable != null) {
                                                 int variableIndex = originalVariable.getSlot();
                                                 if (variableIndex >= 0) {
+                                                        MethodDefEntry ownerMethod = MethodDefEntry.parse(originalVariable.getDeclaringMethod());
                                                         TypeDescriptor variableType = TypeDescriptor.parse(originalVariable.getVariableType());
-                                                        LocalVariableDefEntry localVariableEntry = new LocalVariableDefEntry(methodEntry, variableIndex, initializer.getName(), false, variableType);
+                                                        LocalVariableDefEntry localVariableEntry = new LocalVariableDefEntry(ownerMethod, variableIndex, initializer.getName(), false, variableType);
                                                         identifierEntryCache.put(identifier.getName(), localVariableEntry);
                                                         addDeclarationToUnmatched(identifier.getName(), index);
                                                         index.addDeclaration(identifier, localVariableEntry);