Tweak variable name generation (#86)

* Don't apply offset to all methods in abstract class

* Tweak local variable naming

5 files changed, 78 insertions, 42 deletions

diff --git a/src/main/java/cuchaz/enigma/bytecode/translators/TranslationMethodVisitor.java b/src/main/java/cuchaz/enigma/bytecode/translators/TranslationMethodVisitor.java
index 52b40b05..6d0d550b 100644
--- a/src/main/java/cuchaz/enigma/bytecode/translators/TranslationMethodVisitor.java
+++ b/src/main/java/cuchaz/enigma/bytecode/translators/TranslationMethodVisitor.java
@@ -1,14 +1,13 @@
 package cuchaz.enigma.bytecode.translators;
 
-import cuchaz.enigma.mapping.MethodDescriptor;
-import cuchaz.enigma.mapping.Signature;
-import cuchaz.enigma.mapping.Translator;
-import cuchaz.enigma.mapping.TypeDescriptor;
+import cuchaz.enigma.mapping.*;
 import cuchaz.enigma.mapping.entry.*;
 import org.objectweb.asm.*;
 
+import java.util.Collection;
 import java.util.List;
 import java.util.Locale;
+import java.util.stream.Collectors;
 
 public class TranslationMethodVisitor extends MethodVisitor {
         private final ClassDefEntry ownerEntry;
@@ -85,17 +84,26 @@ public class TranslationMethodVisitor extends MethodVisitor {
                 hasParameterMeta = true;
 
                 String translatedSignature = translator.getTranslatedSignature(Signature.createTypedSignature(signature)).toString();
-                int offsetIndex = index;
+                int argumentIndex = methodEntry.getArgumentIndex(ownerEntry, index);
 
-                if (offsetIndex >= 0) {
-                        LocalVariableDefEntry entry = new LocalVariableDefEntry(methodEntry, offsetIndex, name, new TypeDescriptor(desc));
+                if (argumentIndex >= 0) {
+                        LocalVariableDefEntry entry = new LocalVariableDefEntry(methodEntry, index, name, new TypeDescriptor(desc));
                         LocalVariableDefEntry translatedEntry = translator.getTranslatedVariableDef(entry);
                         String translatedName = translatedEntry.getName();
 
                         // TODO: Better name inference
                         if (translatedName.equals(entry.getName())) {
-                                boolean argument = offsetIndex < methodEntry.getDesc().getArgumentDescs().size();
-                                translatedName = inferName(argument, offsetIndex, translatedEntry.getDesc());
+                                List<TypeDescriptor> arguments = methodEntry.getDesc().getArgumentDescs();
+                                List<TypeDescriptor> translatedArguments = arguments.stream()
+                                                .map(translator::getTranslatedTypeDesc)
+                                                .collect(Collectors.toList());
+
+                                boolean argument = argumentIndex < arguments.size();
+                                if (argument) {
+                                        translatedName = inferArgumentName(argumentIndex, translatedEntry.getDesc(), translatedArguments);
+                                } else {
+                                        translatedName = inferLocalVariableName(argumentIndex, translatedEntry.getDesc());
+                                }
                         }
 
                         super.visitLocalVariable(translatedName, translatedEntry.getDesc().toString(), translatedSignature, start, end, index);
@@ -147,48 +155,56 @@ public class TranslationMethodVisitor extends MethodVisitor {
                 // If we didn't receive any parameter metadata, generate it
                 if (!hasParameterMeta) {
                         List<TypeDescriptor> arguments = methodEntry.getDesc().getArgumentDescs();
-                        int flags = ownerEntry.getAccess().getFlags();
-                        int offset = ((flags & Opcodes.ACC_INTERFACE) != 0 || (flags & Opcodes.ACC_ABSTRACT) != 0) ? 1 : 0;
+                        int offset = ((methodEntry.getAccess().getFlags() & Opcodes.ACC_ABSTRACT) != 0) ? 1 : 0;
 
-                        for (int index = 0; index < arguments.size(); index++) {
+                        for (int argumentIndex = 0; argumentIndex < arguments.size(); argumentIndex++) {
                                 LocalVariableEntry entry = new LocalVariableEntry(methodEntry, offset, "", true);
                                 LocalVariableEntry translatedEntry = translator.getTranslatedVariable(entry);
                                 String translatedName = translatedEntry.getName();
                                 if (translatedName.equals(entry.getName())) {
-                                        super.visitParameter(inferName(true, index, arguments.get(index)), 0);
+                                        super.visitParameter(inferArgumentName(argumentIndex, arguments.get(argumentIndex), arguments), 0);
                                 } else {
                                         super.visitParameter(translatedName, 0);
                                 }
 
-                                offset += arguments.get(index).getSize();
+                                offset += arguments.get(argumentIndex).getSize();
                         }
                 }
                 super.visitEnd();
         }
 
-        private String inferName(boolean argument, int argumentIndex, TypeDescriptor desc) {
+        private String inferArgumentName(int index, TypeDescriptor desc, Collection<TypeDescriptor> arguments) {
+                boolean uniqueType = arguments.stream().filter(desc::equals).count() <= 1;
                 String translatedName;
-                int nameIndex = argumentIndex + 1;
-                StringBuilder nameBuilder = new StringBuilder(argument ? "a" : "v");
+                int nameIndex = index + 1;
+                StringBuilder nameBuilder = new StringBuilder(getTypeName(desc));
+                if (!uniqueType || NameValidator.isReserved(nameBuilder.toString())) {
+                        nameBuilder.append(nameIndex);
+                }
+                translatedName = nameBuilder.toString();
+                return translatedName;
+        }
+
+        private String inferLocalVariableName(int index, TypeDescriptor desc) {
+                int nameIndex = index + 1;
+                return getTypeName(desc) + nameIndex;
+        }
+
+        private String getTypeName(TypeDescriptor desc) {
                 // Unfortunately each of these have different name getters, so they have different code paths
                 if (desc.isPrimitive()) {
                         TypeDescriptor.Primitive argCls = desc.getPrimitive();
-                        nameBuilder.append(argCls.name());
+                        return argCls.name().toLowerCase(Locale.ROOT);
                 } else if (desc.isArray()) {
                         // List types would require this whole block again, so just go with aListx
-                        nameBuilder.append("Arr");
+                        return "arr";
                 } else if (desc.isType()) {
                         String typeName = desc.getTypeEntry().getSimpleName().replace("$", "");
-                        typeName = typeName.substring(0, 1).toUpperCase(Locale.ROOT) + typeName.substring(1);
-                        nameBuilder.append(typeName);
+                        typeName = typeName.substring(0, 1).toLowerCase(Locale.ROOT) + typeName.substring(1);
+                        return typeName;
                 } else {
                         System.err.println("Encountered invalid argument type descriptor " + desc.toString());
-                        nameBuilder.append("Unk");
+                        return "var";
                 }
-                if (!argument || methodEntry.getDesc().getArgumentDescs().size() > 1) {
-                        nameBuilder.append(nameIndex);
-                }
-                translatedName = nameBuilder.toString();
-                return translatedName;
         }
 }
diff --git a/src/main/java/cuchaz/enigma/mapping/NameValidator.java b/src/main/java/cuchaz/enigma/mapping/NameValidator.java
index 9273c9bc..fca8cfcb 100644
--- a/src/main/java/cuchaz/enigma/mapping/NameValidator.java
+++ b/src/main/java/cuchaz/enigma/mapping/NameValidator.java
@@ -66,4 +66,8 @@ public class NameValidator {
         public static String validateArgumentName(String name) {
                 return validateFieldName(name);
         }
+
+        public static boolean isReserved(String name) {
+                return ReservedWords.contains(name);
+        }
 }
diff --git a/src/main/java/cuchaz/enigma/mapping/TypeDescriptor.java b/src/main/java/cuchaz/enigma/mapping/TypeDescriptor.java
index 3c8cc746..6e58aa0f 100644
--- a/src/main/java/cuchaz/enigma/mapping/TypeDescriptor.java
+++ b/src/main/java/cuchaz/enigma/mapping/TypeDescriptor.java
@@ -223,14 +223,14 @@ public class TypeDescriptor {
         }
 
         public enum Primitive {
-                Byte('B'),
-                Character('C'),
-                Short('S'),
-                Integer('I'),
-                Long('J'),
-                Float('F'),
-                Double('D'),
-                Boolean('Z');
+                BYTE('B'),
+                CHARACTER('C'),
+                SHORT('S'),
+                INTEGER('I'),
+                LONG('J'),
+                FLOAT('F'),
+                DOUBLE('D'),
+                BOOLEAN('Z');
 
                 private static final Map<Character, Primitive> lookup;
 
diff --git a/src/main/java/cuchaz/enigma/mapping/entry/MethodDefEntry.java b/src/main/java/cuchaz/enigma/mapping/entry/MethodDefEntry.java
index 960b08d1..fa9e6685 100644
--- a/src/main/java/cuchaz/enigma/mapping/entry/MethodDefEntry.java
+++ b/src/main/java/cuchaz/enigma/mapping/entry/MethodDefEntry.java
@@ -42,4 +42,20 @@ public class MethodDefEntry extends MethodEntry implements DefEntry {
         public MethodDefEntry updateOwnership(ClassEntry classEntry) {
                 return new MethodDefEntry(new ClassEntry(classEntry.getName()), 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;
+        }
 }
diff --git a/src/test/java/cuchaz/enigma/TestTypeDescriptor.java b/src/test/java/cuchaz/enigma/TestTypeDescriptor.java
index 90fd6351..8172848c 100644
--- a/src/test/java/cuchaz/enigma/TestTypeDescriptor.java
+++ b/src/test/java/cuchaz/enigma/TestTypeDescriptor.java
@@ -51,13 +51,13 @@ public class TestTypeDescriptor {
 
         @Test
         public void getPrimitive() {
-                assertThat(new TypeDescriptor("Z").getPrimitive(), is(TypeDescriptor.Primitive.Boolean));
-                assertThat(new TypeDescriptor("B").getPrimitive(), is(TypeDescriptor.Primitive.Byte));
-                assertThat(new TypeDescriptor("C").getPrimitive(), is(TypeDescriptor.Primitive.Character));
-                assertThat(new TypeDescriptor("I").getPrimitive(), is(TypeDescriptor.Primitive.Integer));
-                assertThat(new TypeDescriptor("J").getPrimitive(), is(TypeDescriptor.Primitive.Long));
-                assertThat(new TypeDescriptor("F").getPrimitive(), is(TypeDescriptor.Primitive.Float));
-                assertThat(new TypeDescriptor("D").getPrimitive(), is(TypeDescriptor.Primitive.Double));
+                assertThat(new TypeDescriptor("Z").getPrimitive(), is(TypeDescriptor.Primitive.BOOLEAN));
+                assertThat(new TypeDescriptor("B").getPrimitive(), is(TypeDescriptor.Primitive.BYTE));
+                assertThat(new TypeDescriptor("C").getPrimitive(), is(TypeDescriptor.Primitive.CHARACTER));
+                assertThat(new TypeDescriptor("I").getPrimitive(), is(TypeDescriptor.Primitive.INTEGER));
+                assertThat(new TypeDescriptor("J").getPrimitive(), is(TypeDescriptor.Primitive.LONG));
+                assertThat(new TypeDescriptor("F").getPrimitive(), is(TypeDescriptor.Primitive.FLOAT));
+                assertThat(new TypeDescriptor("D").getPrimitive(), is(TypeDescriptor.Primitive.DOUBLE));
         }
 
         @Test