Annotation + inner class translation

10 files changed, 133 insertions, 23 deletions

diff --git a/src/main/java/cuchaz/enigma/bytecode/AccessFlags.java b/src/main/java/cuchaz/enigma/bytecode/AccessFlags.java
index 0999abf0..152f4629 100644
--- a/src/main/java/cuchaz/enigma/bytecode/AccessFlags.java
+++ b/src/main/java/cuchaz/enigma/bytecode/AccessFlags.java
@@ -5,6 +5,9 @@ import org.objectweb.asm.Opcodes;
 import java.lang.reflect.Modifier;
 
 public class AccessFlags {
+        public static final AccessFlags PUBLIC = new AccessFlags(Modifier.PUBLIC);
+        public static final AccessFlags PUBLIC_STATIC_FINAL = new AccessFlags(Modifier.PUBLIC | Modifier.STATIC | Modifier.FINAL);
+
         private static final int SYNTHETIC_FLAG = 0x00001000;
         private static final int BRIDGED_FLAG = 0x00000040;
 
diff --git a/src/main/java/cuchaz/enigma/bytecode/translators/TranslationAnnotationVisitor.java b/src/main/java/cuchaz/enigma/bytecode/translators/TranslationAnnotationVisitor.java
new file mode 100644
index 00000000..177691b4
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/bytecode/translators/TranslationAnnotationVisitor.java
@@ -0,0 +1,45 @@
+package cuchaz.enigma.bytecode.translators;
+
+import cuchaz.enigma.bytecode.AccessFlags;
+import cuchaz.enigma.mapping.Translator;
+import cuchaz.enigma.mapping.TypeDescriptor;
+import cuchaz.enigma.mapping.entry.ClassEntry;
+import cuchaz.enigma.mapping.entry.FieldDefEntry;
+import org.objectweb.asm.AnnotationVisitor;
+import org.objectweb.asm.Type;
+
+public class TranslationAnnotationVisitor extends AnnotationVisitor {
+        private final Translator translator;
+        private final ClassEntry annotationEntry;
+
+        public TranslationAnnotationVisitor(Translator translator, ClassEntry annotationEntry, int api, AnnotationVisitor av) {
+                super(api, av);
+                this.translator = translator;
+                this.annotationEntry = annotationEntry;
+        }
+
+        @Override
+        public void visit(String name, Object value) {
+                if (value instanceof Type) {
+                        Type type = (Type) value;
+                        super.visit(name, translator.getTranslatedType(type));
+                } else {
+                        super.visit(name, value);
+                }
+        }
+
+        @Override
+        public AnnotationVisitor visitAnnotation(String name, String desc) {
+                TypeDescriptor type = new TypeDescriptor(desc);
+                FieldDefEntry annotationField = translator.getTranslatedFieldDef(new FieldDefEntry(annotationEntry, name, type, AccessFlags.PUBLIC));
+                return super.visitAnnotation(annotationField.getName(), annotationField.getDesc().toString());
+        }
+
+        @Override
+        public void visitEnum(String name, String desc, String value) {
+                TypeDescriptor type = new TypeDescriptor(desc);
+                FieldDefEntry annotationField = translator.getTranslatedFieldDef(new FieldDefEntry(annotationEntry, name, type, AccessFlags.PUBLIC));
+                FieldDefEntry enumField = translator.getTranslatedFieldDef(new FieldDefEntry(type.getTypeEntry(), value, type, AccessFlags.PUBLIC_STATIC_FINAL));
+                super.visitEnum(annotationField.getName(), annotationField.getDesc().toString(), enumField.getName());
+        }
+}
diff --git a/src/main/java/cuchaz/enigma/bytecode/translators/TranslationClassVisitor.java b/src/main/java/cuchaz/enigma/bytecode/translators/TranslationClassVisitor.java
index 3f61c291..5fdfaf11 100644
--- a/src/main/java/cuchaz/enigma/bytecode/translators/TranslationClassVisitor.java
+++ b/src/main/java/cuchaz/enigma/bytecode/translators/TranslationClassVisitor.java
@@ -54,7 +54,8 @@ public class TranslationClassVisitor extends ClassVisitor {
         public FieldVisitor visitField(int access, String name, String desc, String signature, Object value) {
                 FieldDefEntry entry = new FieldDefEntry(obfClassEntry, name, new TypeDescriptor(desc), new AccessFlags(access));
                 FieldDefEntry translatedEntry = translator.getTranslatedFieldDef(entry);
-                return super.visitField(translatedEntry.getAccess().getFlags(), translatedEntry.getName(), translatedEntry.getDesc().toString(), signature, value);
+                FieldVisitor fv = super.visitField(translatedEntry.getAccess().getFlags(), translatedEntry.getName(), translatedEntry.getDesc().toString(), signature, value);
+                return new TranslationFieldVisitor(translator, translatedEntry, api, fv);
         }
 
         @Override
@@ -64,32 +65,38 @@ public class TranslationClassVisitor extends ClassVisitor {
                 if (jarIndex.getBridgedMethod(entry) != null) {
                         translatedEntry.getAccess().setBridged();
                 }
-                MethodVisitor mv = super.visitMethod(translatedEntry.getAccess().getFlags(), name, desc, signature, exceptions);
+                String[] translatedExceptions = new String[exceptions.length];
+                for (int i = 0; i < exceptions.length; i++) {
+                        translatedExceptions[i] = translator.getTranslatedClass(entryPool.getClass(exceptions[i])).getName();
+                }
+                MethodVisitor mv = super.visitMethod(translatedEntry.getAccess().getFlags(), translatedEntry.getName(), translatedEntry.getDesc().toString(), signature, translatedExceptions);
                 return new TranslationMethodVisitor(translator, translatedEntry, api, mv);
         }
 
         @Override
         public void visitOuterClass(String owner, String name, String desc) {
+                ClassEntry ownerEntry = translator.getTranslatedClass(entryPool.getClass(owner));
+                String translatedDesc = desc != null ? translator.getTranslatedTypeDesc(new TypeDescriptor(desc)).toString() : desc;
                 if (name != null) {
-                        ClassEntry ownerEntry = translator.getTranslatedClass(entryPool.getClass(owner));
                         ClassEntry entry = translator.getTranslatedClass(entryPool.getClass(name));
-                        String translatedDesc = translator.getTranslatedTypeDesc(new TypeDescriptor(desc)).toString();
                         super.visitOuterClass(ownerEntry.getName(), entry.getName(), translatedDesc);
                 } else {
-                        super.visitOuterClass(owner, name, desc);
+                        super.visitOuterClass(ownerEntry.getName(), name, translatedDesc);
                 }
         }
 
         @Override
         public AnnotationVisitor visitAnnotation(String desc, boolean visible) {
-                // TODO: Implement
-                return super.visitAnnotation(desc, visible);
+                TypeDescriptor translatedDesc = translator.getTranslatedTypeDesc(new TypeDescriptor(desc));
+                AnnotationVisitor av = super.visitAnnotation(translatedDesc.toString(), visible);
+                return new TranslationAnnotationVisitor(translator, translatedDesc.getTypeEntry(), api, av);
         }
 
         @Override
         public AnnotationVisitor visitTypeAnnotation(int typeRef, TypePath typePath, String desc, boolean visible) {
-                // TODO: Implement
-                return super.visitTypeAnnotation(typeRef, typePath, desc, visible);
+                TypeDescriptor translatedDesc = translator.getTranslatedTypeDesc(new TypeDescriptor(desc));
+                AnnotationVisitor av = super.visitTypeAnnotation(typeRef, typePath, translatedDesc.toString(), visible);
+                return new TranslationAnnotationVisitor(translator, translatedDesc.getTypeEntry(), api, av);
         }
 
         @Override
@@ -101,7 +108,11 @@ public class TranslationClassVisitor extends ClassVisitor {
                         ClassEntry innerEntry = translator.getTranslatedClass(entryPool.getClass(innerName));
                         super.visitInnerClass(translatedEntry.getName(), outerEntry.getName(), innerEntry.getName(), translatedEntry.getAccess().getFlags());
                 } else {
-                        super.visitInnerClass(name, outerName, innerName, access);
+                        int separatorIndex = name.lastIndexOf("$");
+                        String parentName = name.substring(0, separatorIndex);
+                        String childName = name.substring(separatorIndex + 1);
+                        ClassEntry outerEntry = translator.getTranslatedClass(entryPool.getClass(parentName));
+                        super.visitInnerClass(outerEntry.getName() + "$" + childName, outerName, innerName, access);
                 }
         }
 }
diff --git a/src/main/java/cuchaz/enigma/bytecode/translators/TranslationFieldVisitor.java b/src/main/java/cuchaz/enigma/bytecode/translators/TranslationFieldVisitor.java
new file mode 100644
index 00000000..e4695fb6
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/bytecode/translators/TranslationFieldVisitor.java
@@ -0,0 +1,33 @@
+package cuchaz.enigma.bytecode.translators;
+
+import cuchaz.enigma.mapping.Translator;
+import cuchaz.enigma.mapping.TypeDescriptor;
+import cuchaz.enigma.mapping.entry.FieldDefEntry;
+import org.objectweb.asm.AnnotationVisitor;
+import org.objectweb.asm.FieldVisitor;
+import org.objectweb.asm.TypePath;
+
+public class TranslationFieldVisitor extends FieldVisitor {
+        private final FieldDefEntry fieldEntry;
+        private final Translator translator;
+
+        public TranslationFieldVisitor(Translator translator, FieldDefEntry fieldEntry, int api, FieldVisitor fv) {
+                super(api, fv);
+                this.translator = translator;
+                this.fieldEntry = fieldEntry;
+        }
+
+        @Override
+        public AnnotationVisitor visitAnnotation(String desc, boolean visible) {
+                TypeDescriptor typeDesc = translator.getTranslatedTypeDesc(new TypeDescriptor(desc));
+                AnnotationVisitor av = super.visitAnnotation(typeDesc.toString(), visible);
+                return new TranslationAnnotationVisitor(translator, typeDesc.getTypeEntry(), api, av);
+        }
+
+        @Override
+        public AnnotationVisitor visitTypeAnnotation(int typeRef, TypePath typePath, String desc, boolean visible) {
+                TypeDescriptor typeDesc = translator.getTranslatedTypeDesc(new TypeDescriptor(desc));
+                AnnotationVisitor av = super.visitAnnotation(typeDesc.toString(), visible);
+                return new TranslationAnnotationVisitor(translator, typeDesc.getTypeEntry(), api, av);
+        }
+}
diff --git a/src/main/java/cuchaz/enigma/bytecode/translators/TranslationMethodVisitor.java b/src/main/java/cuchaz/enigma/bytecode/translators/TranslationMethodVisitor.java
index b2192e47..c7299628 100644
--- a/src/main/java/cuchaz/enigma/bytecode/translators/TranslationMethodVisitor.java
+++ b/src/main/java/cuchaz/enigma/bytecode/translators/TranslationMethodVisitor.java
@@ -39,15 +39,17 @@ public class TranslationMethodVisitor extends MethodVisitor {
         }
 
         @Override
-        public void visitAttribute(Attribute attr) {
-                // TODO: Implement
-                super.visitAttribute(attr);
+        public AnnotationVisitor visitAnnotation(String desc, boolean visible) {
+                TypeDescriptor typeDesc = translator.getTranslatedTypeDesc(new TypeDescriptor(desc));
+                AnnotationVisitor av = super.visitAnnotation(typeDesc.toString(), visible);
+                return new TranslationAnnotationVisitor(translator, typeDesc.getTypeEntry(), api, av);
         }
 
         @Override
-        public AnnotationVisitor visitAnnotation(String desc, boolean visible) {
-                // TODO: Implement
-                return super.visitAnnotation(desc, visible);
+        public AnnotationVisitor visitTypeAnnotation(int typeRef, TypePath typePath, String desc, boolean visible) {
+                TypeDescriptor typeDesc = translator.getTranslatedTypeDesc(new TypeDescriptor(desc));
+                AnnotationVisitor av = super.visitAnnotation(typeDesc.toString(), visible);
+                return new TranslationAnnotationVisitor(translator, typeDesc.getTypeEntry(), api, av);
         }
 
         @Override
@@ -69,7 +71,7 @@ public class TranslationMethodVisitor extends MethodVisitor {
                                 // List types would require this whole block again, so just go with aListx
                                 nameBuilder.append(nameIndex);
                         } else if (argDesc.isType()) {
-                                String typeName = argDesc.getOwnerEntry().getSimpleName().replace("$", "");
+                                String typeName = argDesc.getTypeEntry().getSimpleName().replace("$", "");
                                 typeName = typeName.substring(0, 1).toUpperCase(Locale.ROOT) + typeName.substring(1);
                                 nameBuilder.append(typeName);
                         }
diff --git a/src/main/java/cuchaz/enigma/mapping/DirectionalTranslator.java b/src/main/java/cuchaz/enigma/mapping/DirectionalTranslator.java
index c9e2ab49..196bbd69 100644
--- a/src/main/java/cuchaz/enigma/mapping/DirectionalTranslator.java
+++ b/src/main/java/cuchaz/enigma/mapping/DirectionalTranslator.java
@@ -185,7 +185,7 @@ public class DirectionalTranslator implements Translator {
                 }
                 // TODO: Translating arguments calls method translation.. Can we refactor the code in such a way that we don't need this?
                 MethodEntry translatedOwner = getTranslatedMethod(entry.getOwnerEntry());
-                return new LocalVariableEntry(translatedOwner, entry.getIndex(), translatedArgumentName);
+                return new LocalVariableEntry(translatedOwner != null ? translatedOwner : entry.getOwnerEntry(), entry.getIndex(), translatedArgumentName);
         }
 
         @Override
diff --git a/src/main/java/cuchaz/enigma/mapping/Mappings.java b/src/main/java/cuchaz/enigma/mapping/Mappings.java
index 9b907a9c..79d9f104 100644
--- a/src/main/java/cuchaz/enigma/mapping/Mappings.java
+++ b/src/main/java/cuchaz/enigma/mapping/Mappings.java
@@ -156,7 +156,7 @@ public class Mappings {
                         for (MethodMapping methodMapping : classMapping.methods()) {
                                 for (TypeDescriptor desc : methodMapping.getObfDesc().types()) {
                                         if (desc.containsType()) {
-                                                classNames.add(desc.getOwnerEntry().getClassName());
+                                                classNames.add(desc.getTypeEntry().getClassName());
                                         }
                                 }
                         }
diff --git a/src/main/java/cuchaz/enigma/mapping/MethodDescriptor.java b/src/main/java/cuchaz/enigma/mapping/MethodDescriptor.java
index b56d1d48..0fc03517 100644
--- a/src/main/java/cuchaz/enigma/mapping/MethodDescriptor.java
+++ b/src/main/java/cuchaz/enigma/mapping/MethodDescriptor.java
@@ -97,7 +97,7 @@ public class MethodDescriptor {
 
         public boolean hasClass(ClassEntry classEntry) {
                 for (TypeDescriptor desc : types()) {
-                        if (desc.containsType() && desc.getOwnerEntry().equals(classEntry)) {
+                        if (desc.containsType() && desc.getTypeEntry().equals(classEntry)) {
                                 return true;
                         }
                 }
diff --git a/src/main/java/cuchaz/enigma/mapping/Translator.java b/src/main/java/cuchaz/enigma/mapping/Translator.java
index ab55a02d..6373da7b 100644
--- a/src/main/java/cuchaz/enigma/mapping/Translator.java
+++ b/src/main/java/cuchaz/enigma/mapping/Translator.java
@@ -12,6 +12,7 @@
 package cuchaz.enigma.mapping;
 
 import cuchaz.enigma.mapping.entry.*;
+import org.objectweb.asm.Type;
 
 public interface Translator {
         ClassEntry getTranslatedClass(ClassEntry entry);
@@ -34,6 +35,18 @@ public interface Translator {
 
         MethodDescriptor getTranslatedMethodDesc(MethodDescriptor descriptor);
 
+        default Type getTranslatedType(Type type) {
+                String descString = type.getDescriptor();
+                // If this is a method
+                if (descString.contains("(")) {
+                        MethodDescriptor descriptor = new MethodDescriptor(descString);
+                        return Type.getMethodType(getTranslatedMethodDesc(descriptor).toString());
+                } else {
+                        TypeDescriptor descriptor = new TypeDescriptor(descString);
+                        return Type.getType(getTranslatedTypeDesc(descriptor).toString());
+                }
+        }
+
         @SuppressWarnings("unchecked")
         default <T extends Entry> T getTranslatedEntry(T entry) {
                 if (entry instanceof ClassDefEntry) {
diff --git a/src/main/java/cuchaz/enigma/mapping/TypeDescriptor.java b/src/main/java/cuchaz/enigma/mapping/TypeDescriptor.java
index 58fa6513..b7b1255a 100644
--- a/src/main/java/cuchaz/enigma/mapping/TypeDescriptor.java
+++ b/src/main/java/cuchaz/enigma/mapping/TypeDescriptor.java
@@ -11,6 +11,7 @@
 
 package cuchaz.enigma.mapping;
 
+import com.google.common.base.Preconditions;
 import com.google.common.collect.Maps;
 import cuchaz.enigma.mapping.entry.ClassEntry;
 
@@ -22,6 +23,8 @@ public class TypeDescriptor {
         protected final String desc;
 
         public TypeDescriptor(String desc) {
+                Preconditions.checkNotNull(desc, "Desc cannot be null");
+
                 // don't deal with generics
                 // this is just for raw jvm types
                 if (desc.charAt(0) == 'T' || desc.indexOf('<') >= 0 || desc.indexOf('>') >= 0) {
@@ -127,7 +130,7 @@ public class TypeDescriptor {
                 return this.desc.charAt(0) == 'L' && this.desc.charAt(this.desc.length() - 1) == ';';
         }
 
-        public ClassEntry getOwnerEntry() {
+        public ClassEntry getTypeEntry() {
                 if (isType()) {
                         String name = this.desc.substring(1, this.desc.length() - 1);
 
@@ -140,7 +143,7 @@ public class TypeDescriptor {
                         return new ClassEntry(name);
 
                 } else if (isArray() && getArrayType().isType()) {
-                        return getArrayType().getOwnerEntry();
+                        return getArrayType().getTypeEntry();
                 } else {
                         throw new IllegalStateException("desc doesn't have a class");
                 }
@@ -185,7 +188,7 @@ public class TypeDescriptor {
         public TypeDescriptor remap(Function<String, String> remapper) {
                 String desc = this.desc;
                 if (isType() || (isArray() && containsType())) {
-                        String replacedName = remapper.apply(this.getOwnerEntry().getName());
+                        String replacedName = remapper.apply(this.getTypeEntry().getName());
                         if (replacedName != null) {
                                 if (this.isType()) {
                                         desc = "L" + replacedName + ";";