2 files changed, 93 insertions, 13 deletions

diff --git a/src/main/java/cuchaz/enigma/mapping/DirectionalTranslator.java b/src/main/java/cuchaz/enigma/mapping/DirectionalTranslator.java
index 7af1a52..dec5acf 100644
--- a/src/main/java/cuchaz/enigma/mapping/DirectionalTranslator.java
+++ b/src/main/java/cuchaz/enigma/mapping/DirectionalTranslator.java
@@ -15,13 +15,19 @@ import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import cuchaz.enigma.analysis.TranslationIndex;
 import cuchaz.enigma.bytecode.AccessFlags;
+import cuchaz.enigma.bytecode.translators.TranslationSignatureVisitor;
 import cuchaz.enigma.mapping.entry.*;
+import org.objectweb.asm.signature.SignatureReader;
+import org.objectweb.asm.signature.SignatureVisitor;
+import org.objectweb.asm.signature.SignatureWriter;
 
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.regex.Pattern;
 
 public class DirectionalTranslator implements Translator {
+        private static final Pattern OBJECT_PATTERN = Pattern.compile(".*:Ljava/lang/Object;:.*");
 
         private final TranslationDirection direction;
         private final Map<String, ClassMapping> classes;
@@ -98,7 +104,7 @@ public class DirectionalTranslator implements Translator {
         public FieldDefEntry getTranslatedFieldDef(FieldDefEntry entry) {
                 String translatedName = translateFieldName(entry);
                 if (translatedName == null) {
-                        return entry;
+                        translatedName = entry.getName();
                 }
                 ClassEntry translatedOwner = getTranslatedClass(entry.getOwnerClassEntry());
                 TypeDescriptor translatedDesc = getTranslatedTypeDesc(entry.getDesc());
@@ -110,7 +116,7 @@ public class DirectionalTranslator implements Translator {
         public FieldEntry getTranslatedField(FieldEntry entry) {
                 String translatedName = translateFieldName(entry);
                 if (translatedName == null) {
-                        return null;
+                        translatedName = entry.getName();
                 }
                 ClassEntry translatedOwner = getTranslatedClass(entry.getOwnerClassEntry());
                 TypeDescriptor translatedDesc = getTranslatedTypeDesc(entry.getDesc());
@@ -138,7 +144,7 @@ public class DirectionalTranslator implements Translator {
         public MethodDefEntry getTranslatedMethodDef(MethodDefEntry entry) {
                 String translatedName = translateMethodName(entry);
                 if (translatedName == null) {
-                        return entry;
+                        translatedName = entry.getName();
                 }
                 ClassEntry translatedOwner = getTranslatedClass(entry.getOwnerClassEntry());
                 MethodDescriptor translatedDesc = getTranslatedMethodDesc(entry.getDesc());
@@ -150,7 +156,7 @@ public class DirectionalTranslator implements Translator {
         public MethodEntry getTranslatedMethod(MethodEntry entry) {
                 String translatedName = translateMethodName(entry);
                 if (translatedName == null) {
-                        return null;
+                        translatedName = entry.getName();
                 }
                 ClassEntry translatedOwner = getTranslatedClass(entry.getOwnerClassEntry());
                 MethodDescriptor translatedDesc = getTranslatedMethodDesc(entry.getDesc());
@@ -181,11 +187,11 @@ public class DirectionalTranslator implements Translator {
                         translatedArgumentName = inheritLocalVariableName(entry);
                 }
                 if (translatedArgumentName == null) {
-                        return null;
+                        translatedArgumentName = entry.getName();
                 }
                 // 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 != null ? translatedOwner : entry.getOwnerEntry(), entry.getIndex(), translatedArgumentName);
+                return new LocalVariableEntry(translatedOwner, entry.getIndex(), translatedArgumentName);
         }
 
         @Override
@@ -200,6 +206,26 @@ public class DirectionalTranslator implements Translator {
                 return new LocalVariableDefEntry(translatedOwner, entry.getIndex(), translatedArgumentName != null ? translatedArgumentName : entry.getName(), translatedTypeDesc);
         }
 
+        @Override
+        public boolean hasClassMapping(ClassEntry entry) {
+                return classes.containsKey(entry.getName());
+        }
+
+        @Override
+        public boolean hasFieldMapping(FieldEntry entry) {
+                return translateFieldName(entry) != null;
+        }
+
+        @Override
+        public boolean hasMethodMapping(MethodEntry entry) {
+                return translateMethodName(entry) != null;
+        }
+
+        @Override
+        public boolean hasLocalVariableMapping(LocalVariableEntry entry) {
+                return translateLocalVariableName(entry) != null || inheritLocalVariableName(entry) != null;
+        }
+
         // TODO: support not identical behavior (specific to constructor)
         private String translateLocalVariableName(LocalVariableEntry entry) {
                 // look for identical behavior in superclasses
@@ -252,6 +278,26 @@ public class DirectionalTranslator implements Translator {
                 return new MethodDescriptor(translatedArguments, getTranslatedTypeDesc(descriptor.getReturnDesc()));
         }
 
+        @Override
+        public String getTranslatedSignature(String signature, boolean isType, int api) {
+                if (signature == null) {
+                        return null;
+                }
+                SignatureReader reader = new SignatureReader(signature);
+                SignatureWriter writer = new SignatureWriter();
+                SignatureVisitor visitor = new TranslationSignatureVisitor(this, api, writer);
+                if (isType) {
+                        reader.acceptType(visitor);
+                } else {
+                        reader.accept(visitor);
+                }
+                String translatedSignature = writer.toString();
+                if (OBJECT_PATTERN.matcher(signature).matches()) {
+                        translatedSignature = signature.replaceAll(":Ljava/lang/Object;:", "::");
+                }
+                return translatedSignature;
+        }
+
         private ClassMapping findClassMapping(ClassEntry entry) {
                 List<ClassMapping> mappingChain = getClassMappingChain(entry);
                 return mappingChain.get(mappingChain.size() - 1);
diff --git a/src/main/java/cuchaz/enigma/mapping/Translator.java b/src/main/java/cuchaz/enigma/mapping/Translator.java
index 6373da7..1bc2f37 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.Handle;
 import org.objectweb.asm.Type;
 
 public interface Translator {
@@ -31,20 +32,53 @@ public interface Translator {
 
         LocalVariableDefEntry getTranslatedVariableDef(LocalVariableDefEntry entry);
 
+        boolean hasClassMapping(ClassEntry entry);
+
+        boolean hasFieldMapping(FieldEntry entry);
+
+        boolean hasMethodMapping(MethodEntry entry);
+
+        boolean hasLocalVariableMapping(LocalVariableEntry entry);
+
         TypeDescriptor getTranslatedTypeDesc(TypeDescriptor desc);
 
         MethodDescriptor getTranslatedMethodDesc(MethodDescriptor descriptor);
 
+        String getTranslatedSignature(String signature, boolean isType, int api);
+
         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());
+                switch (type.getSort()) {
+                        case Type.OBJECT: {
+                                ClassEntry classEntry = new ClassEntry(type.getInternalName());
+                                return Type.getObjectType(getTranslatedClass(classEntry).getName());
+                        }
+                        case Type.ARRAY: {
+                                TypeDescriptor descriptor = new TypeDescriptor(descString);
+                                return Type.getType(getTranslatedTypeDesc(descriptor).toString());
+                        }
+                        case Type.METHOD: {
+                                MethodDescriptor descriptor = new MethodDescriptor(descString);
+                                return Type.getMethodType(getTranslatedMethodDesc(descriptor).toString());
+                        }
+                }
+                return type;
+        }
+
+        default Handle getTranslatedHandle(Handle handle) {
+                MethodEntry entry = new MethodEntry(new ClassEntry(handle.getOwner()), handle.getName(), new MethodDescriptor(handle.getDesc()));
+                MethodEntry translatedMethod = getTranslatedMethod(entry);
+                ClassEntry ownerClass = translatedMethod.getOwnerClassEntry();
+                return new Handle(handle.getTag(), ownerClass.getName(), translatedMethod.getName(), translatedMethod.getDesc().toString(), handle.isInterface());
+        }
+
+        default Object getTranslatedValue(Object value) {
+                if (value instanceof Type) {
+                        return this.getTranslatedType((Type) value);
+                } else if (value instanceof Handle) {
+                        return getTranslatedHandle((Handle) value);
                 }
+                return value;
         }
 
         @SuppressWarnings("unchecked")