Fix local variable fixer (#172)

* Fix local variable fixer

* LVT index -> LV index

* Small fix

* Use LocalNameGenerator

3 files changed, 92 insertions, 93 deletions

diff --git a/src/main/java/cuchaz/enigma/bytecode/translators/LocalVariableFixVisitor.java b/src/main/java/cuchaz/enigma/bytecode/translators/LocalVariableFixVisitor.java
index 16dbba1..cfd8fbe 100644
--- a/src/main/java/cuchaz/enigma/bytecode/translators/LocalVariableFixVisitor.java
+++ b/src/main/java/cuchaz/enigma/bytecode/translators/LocalVariableFixVisitor.java
@@ -1,14 +1,18 @@
 package cuchaz.enigma.bytecode.translators;
 
 import com.google.common.base.CharMatcher;
+import cuchaz.enigma.translation.LocalNameGenerator;
 import cuchaz.enigma.translation.representation.TypeDescriptor;
 import cuchaz.enigma.translation.representation.entry.ClassDefEntry;
 import cuchaz.enigma.translation.representation.entry.MethodDefEntry;
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.Label;
 import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
 
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 public class LocalVariableFixVisitor extends ClassVisitor {
         private ClassDefEntry ownerEntry;
@@ -31,50 +35,92 @@ public class LocalVariableFixVisitor extends ClassVisitor {
 
         private class Method extends MethodVisitor {
                 private final MethodDefEntry methodEntry;
-                private boolean hasLvt;
+                private final Map<Integer, String> parameterNames = new HashMap<>();
+                private final Map<Integer, Integer> parameterIndices = new HashMap<>();
+                private boolean hasParameterTable;
+                private int parameterIndex = 0;
 
                 Method(int api, MethodDefEntry methodEntry, MethodVisitor visitor) {
                         super(api, visitor);
                         this.methodEntry = methodEntry;
+
+                        int lvIndex = methodEntry.getAccess().isStatic() ? 0 : 1;
+                        List<TypeDescriptor> parameters = methodEntry.getDesc().getArgumentDescs();
+                        for (int parameterIndex = 0; parameterIndex < parameters.size(); parameterIndex++) {
+                                TypeDescriptor param = parameters.get(parameterIndex);
+                                parameterIndices.put(lvIndex, parameterIndex);
+                                lvIndex += param.getSize();
+                        }
                 }
 
                 @Override
-                public void visitLocalVariable(String name, String desc, String signature, Label start, Label end, int index) {
-                        hasLvt = true;
-
-                        String translatedName = name;
+                public void visitParameter(String name, int access) {
+                        hasParameterTable = true;
+                        super.visitParameter(fixParameterName(parameterIndex, name), fixParameterAccess(parameterIndex, access));
+                        parameterIndex++;
+                }
 
-                        if (isInvalidName(name)) {
-                                int argumentIndex = methodEntry.getArgumentIndex(ownerEntry, index);
-
-                                if (argumentIndex >= 0) {
-                                        List<TypeDescriptor> arguments = methodEntry.getDesc().getArgumentDescs();
-                                        boolean argument = argumentIndex < arguments.size();
-                                        if (argument) {
-                                                translatedName = "arg" + (argumentIndex + 1);
-                                        } else {
-                                                translatedName = "var" + (argumentIndex + 1);
-                                        }
-                                }
+                @Override
+                public void visitLocalVariable(String name, String desc, String signature, Label start, Label end, int index) {
+                        if (index == 0 && !methodEntry.getAccess().isStatic()) {
+                                name = "this";
+                        } else if (parameterIndices.containsKey(index)) {
+                                name = fixParameterName(parameterIndices.get(index), name);
+                        } else if (isInvalidName(name)) {
+                                name = LocalNameGenerator.generateLocalVariableName(index, new TypeDescriptor(desc));
                         }
 
-                        super.visitLocalVariable(translatedName, desc, signature, start, end, index);
+                        super.visitLocalVariable(name, desc, signature, start, end, index);
                 }
 
                 private boolean isInvalidName(String name) {
-                        return !CharMatcher.ascii().matchesAllOf(name);
+                        return name == null || name.isEmpty() || !CharMatcher.ascii().matchesAllOf(name);
                 }
 
                 @Override
                 public void visitEnd() {
-                        if (!hasLvt) {
+                        if (!hasParameterTable) {
                                 List<TypeDescriptor> arguments = methodEntry.getDesc().getArgumentDescs();
                                 for (int argumentIndex = 0; argumentIndex < arguments.size(); argumentIndex++) {
-                                        super.visitParameter("arg" + (argumentIndex + 1), 0);
+                                        super.visitParameter(fixParameterName(argumentIndex, null), fixParameterAccess(argumentIndex, 0));
                                 }
                         }
 
                         super.visitEnd();
                 }
+
+                private String fixParameterName(int index, String name) {
+                        if (parameterNames.get(index) != null) {
+                                return parameterNames.get(index); // to make sure that LVT names are consistent with parameter table names
+                        }
+
+                        if (isInvalidName(name)) {
+                                List<TypeDescriptor> arguments = methodEntry.getDesc().getArgumentDescs();
+                                name = LocalNameGenerator.generateArgumentName(index, arguments.get(index), arguments);
+                        }
+
+                        if (index == 0 && ownerEntry.getAccess().isEnum() && methodEntry.getName().equals("<init>")) {
+                                name = "name";
+                        }
+
+                        if (index == 1 && ownerEntry.getAccess().isEnum() && methodEntry.getName().equals("<init>")) {
+                                name = "ordinal";
+                        }
+
+                        parameterNames.put(index, name);
+                        return name;
+                }
+
+                private int fixParameterAccess(int index, int access) {
+                        if (index == 0 && ownerEntry.getAccess().isEnum() && methodEntry.getName().equals("<init>")) {
+                                access |= Opcodes.ACC_SYNTHETIC;
+                        }
+
+                        if (index == 1 && ownerEntry.getAccess().isEnum() && methodEntry.getName().equals("<init>")) {
+                                access |= Opcodes.ACC_SYNTHETIC;
+                        }
+
+                        return access;
+                }
         }
 }
diff --git a/src/main/java/cuchaz/enigma/bytecode/translators/TranslationMethodVisitor.java b/src/main/java/cuchaz/enigma/bytecode/translators/TranslationMethodVisitor.java
index c824265..4bce533 100644
--- a/src/main/java/cuchaz/enigma/bytecode/translators/TranslationMethodVisitor.java
+++ b/src/main/java/cuchaz/enigma/bytecode/translators/TranslationMethodVisitor.java
@@ -1,6 +1,5 @@
 package cuchaz.enigma.bytecode.translators;
 
-import cuchaz.enigma.translation.LocalNameGenerator;
 import cuchaz.enigma.translation.Translator;
 import cuchaz.enigma.translation.representation.MethodDescriptor;
 import cuchaz.enigma.translation.representation.Signature;
@@ -8,21 +7,27 @@ import cuchaz.enigma.translation.representation.TypeDescriptor;
 import cuchaz.enigma.translation.representation.entry.*;
 import org.objectweb.asm.*;
 
-import java.util.List;
-import java.util.stream.Collectors;
-
 public class TranslationMethodVisitor extends MethodVisitor {
-        private final ClassDefEntry ownerEntry;
         private final MethodDefEntry methodEntry;
         private final Translator translator;
 
-        private boolean hasParameterMeta;
+        private int parameterIndex = 0;
+        private int parameterLvIndex;
 
         public TranslationMethodVisitor(Translator translator, ClassDefEntry ownerEntry, MethodDefEntry methodEntry, int api, MethodVisitor mv) {
                 super(api, mv);
                 this.translator = translator;
-                this.ownerEntry = ownerEntry;
                 this.methodEntry = methodEntry;
+
+                parameterLvIndex = methodEntry.getAccess().isStatic() ? 0 : 1;
+        }
+
+        @Override
+        public void visitParameter(String name, int access) {
+                name = translateVariableName(parameterLvIndex, name);
+                parameterLvIndex += methodEntry.getDesc().getArgumentDescs().get(parameterIndex++).getSize();
+
+                super.visitParameter(name, access);
         }
 
         @Override
@@ -119,58 +124,22 @@ public class TranslationMethodVisitor extends MethodVisitor {
 
         @Override
         public void visitLocalVariable(String name, String desc, String signature, Label start, Label end, int index) {
-                hasParameterMeta = true;
-
-                String translatedSignature = translator.translate(Signature.createTypedSignature(signature)).toString();
-                int argumentIndex = methodEntry.getArgumentIndex(ownerEntry, index);
-
-                if (argumentIndex >= 0) {
-                        LocalVariableDefEntry entry = new LocalVariableDefEntry(methodEntry, index, name, true, new TypeDescriptor(desc));
-                        LocalVariableDefEntry translatedEntry = translator.translate(entry);
-                        String translatedName = translatedEntry.getName();
-
-                        if (translatedName.equals(entry.getName())) {
-                                List<TypeDescriptor> arguments = methodEntry.getDesc().getArgumentDescs();
-                                List<TypeDescriptor> translatedArguments = arguments.stream()
-                                                .map(translator::translate)
-                                                .collect(Collectors.toList());
-
-                                boolean argument = argumentIndex < arguments.size();
-                                if (argument) {
-                                        translatedName = LocalNameGenerator.generateArgumentName(argumentIndex, translatedEntry.getDesc(), translatedArguments);
-                                } else {
-                                        translatedName = LocalNameGenerator.generateLocalVariableName(argumentIndex, translatedEntry.getDesc());
-                                }
-                        }
+                signature = translator.translate(Signature.createTypedSignature(signature)).toString();
+                name = translateVariableName(index, name);
+                desc = translator.translate(new TypeDescriptor(desc)).toString();
 
-                        super.visitLocalVariable(translatedName, translatedEntry.getDesc().toString(), translatedSignature, start, end, index);
-                } else {
-                        // Handle "this" variable
-                        TypeDescriptor translatedDesc = translator.translate(new TypeDescriptor(desc));
-                        super.visitLocalVariable(name, translatedDesc.toString(), translatedSignature, start, end, index);
-                }
+                super.visitLocalVariable(name, desc, signature, start, end, index);
         }
 
-        @Override
-        public void visitEnd() {
-                // If we didn't receive any parameter metadata, generate it
-                if (!hasParameterMeta) {
-                        List<TypeDescriptor> arguments = translator.translate(methodEntry.getDesc()).getArgumentDescs();
-                        int offset = ((methodEntry.getAccess().getFlags() & Opcodes.ACC_ABSTRACT) != 0) ? 1 : 0;
-
-                        for (int argumentIndex = 0; argumentIndex < arguments.size(); argumentIndex++) {
-                                LocalVariableEntry entry = new LocalVariableEntry(methodEntry, offset, "", true);
-                                LocalVariableEntry translatedEntry = translator.translate(entry);
-                                String translatedName = translatedEntry.getName();
-                                if (translatedName.equals(entry.getName())) {
-                                        super.visitParameter(LocalNameGenerator.generateArgumentName(argumentIndex, arguments.get(argumentIndex), arguments), 0);
-                                } else {
-                                        super.visitParameter(translatedName, 0);
-                                }
-
-                                offset += arguments.get(argumentIndex).getSize();
-                        }
+        private String translateVariableName(int index, String name) {
+                LocalVariableEntry entry = new LocalVariableEntry(methodEntry, index, "", true);
+                LocalVariableEntry translatedEntry = translator.translate(entry);
+                String translatedName = translatedEntry.getName();
+
+                if (!translatedName.isEmpty()) {
+                        return translatedName;
                 }
-                super.visitEnd();
+
+                return name;
         }
 }
diff --git a/src/main/java/cuchaz/enigma/translation/representation/entry/MethodDefEntry.java b/src/main/java/cuchaz/enigma/translation/representation/entry/MethodDefEntry.java
index bbcaf23..7e89f6a 100644
--- a/src/main/java/cuchaz/enigma/translation/representation/entry/MethodDefEntry.java
+++ b/src/main/java/cuchaz/enigma/translation/representation/entry/MethodDefEntry.java
@@ -72,20 +72,4 @@ public class MethodDefEntry extends MethodEntry implements DefEntry<ClassEntry>
         public MethodDefEntry withParent(ClassEntry parent) {
                 return new MethodDefEntry(new ClassEntry(parent.getFullName()), name, descriptor, signature, access);
         }
-
-        public int getArgumentIndex(ClassDefEntry ownerEntry, int localVariableIndex) {
-                int argumentIndex = localVariableIndex;
-
-                // Enum constructors have an implicit "name" and "ordinal" parameter as well as "this"
-                if (ownerEntry.getAccess().isEnum() && getName().startsWith("<")) {
-                        argumentIndex -= 2;
-                }
-
-                // If we're not static, "this" is bound to index 0
-                if (!getAccess().isStatic()) {
-                        argumentIndex -= 1;
-                }
-
-                return argumentIndex;
-        }
 }