5 files changed, 50 insertions, 30 deletions

diff --git a/src/main/java/cuchaz/enigma/TranslatingTypeLoader.java b/src/main/java/cuchaz/enigma/TranslatingTypeLoader.java
index 6ddfa634..d9dc67b2 100644
--- a/src/main/java/cuchaz/enigma/TranslatingTypeLoader.java
+++ b/src/main/java/cuchaz/enigma/TranslatingTypeLoader.java
@@ -51,7 +51,8 @@ public class TranslatingTypeLoader implements ITypeLoader {
         this.cache = Maps.newHashMap();
         this.defaultTypeLoader = new ClasspathTypeLoader();
         this.docMapping = docMapping;
-        this.docMapping.cleanBehaviors();
+        if (docMapping != null)
+            this.docMapping.cleanBehaviors();
     }
 
     public void clearCache() {
diff --git a/src/main/java/cuchaz/enigma/gui/MemberMatchingGui.java b/src/main/java/cuchaz/enigma/gui/MemberMatchingGui.java
index 50550910..91067b50 100644
--- a/src/main/java/cuchaz/enigma/gui/MemberMatchingGui.java
+++ b/src/main/java/cuchaz/enigma/gui/MemberMatchingGui.java
@@ -382,13 +382,11 @@ public class MemberMatchingGui<T extends Entry> {
     }
 
     private void setDest(T obfEntry) {
-        if (obfEntry == null) {
-            m_obfDestEntry = null;
+        m_obfDestEntry = obfEntry;
+        if (obfEntry == null)
             m_destLabel.setText("");
-        } else {
-            m_obfDestEntry = obfEntry;
+        else
             m_destLabel.setText(getEntryLabel(obfEntry, m_destDeobfuscator));
-        }
     }
 
     private String getEntryLabel(T obfEntry, Deobfuscator deobfuscator) {
diff --git a/src/main/java/cuchaz/enigma/mapping/NameValidator.java b/src/main/java/cuchaz/enigma/mapping/NameValidator.java
index 6925b72a..15b0314c 100644
--- a/src/main/java/cuchaz/enigma/mapping/NameValidator.java
+++ b/src/main/java/cuchaz/enigma/mapping/NameValidator.java
@@ -30,6 +30,8 @@ public class NameValidator {
     );
 
     static {
+
+        // java allows all kinds of weird characters...
         String identifierRegex = "[A-Za-z_<][A-Za-z0-9_>]*";
         IdentifierPattern = Pattern.compile(identifierRegex);
         ClassPattern = Pattern.compile(String.format("^(%s(\\.|/))*(%s)$", identifierRegex, identifierRegex));
diff --git a/src/main/java/cuchaz/enigma/mapping/javadoc/JavaDocMapping.java b/src/main/java/cuchaz/enigma/mapping/javadoc/JavaDocMapping.java
index 98996fdb..8b68a12d 100644
--- a/src/main/java/cuchaz/enigma/mapping/javadoc/JavaDocMapping.java
+++ b/src/main/java/cuchaz/enigma/mapping/javadoc/JavaDocMapping.java
@@ -35,7 +35,7 @@ public class JavaDocMapping
 
         // TODO: File format
         addField("a", "none/akw", "Lnone/kp;", "Hello from Enigma");
-        addMethod("<init>", "none/akw", "(Lnone/ayo;)V", "You know what? I love constructors!", "The material of the block!");
+        addMethod("<init>", "none/akw", "(Lnone/ayo;)V", "You know what? I love constructors!", null, new String[]{ "The material of the block!" });
         addClass("none/akw", "HEY I'M A BLOCK YOU KNOW THAT?!");
     }
 
@@ -49,7 +49,7 @@ public class JavaDocMapping
         this.javaDocFieldByID.put(new FieldEntry(new ClassEntry(className), fieldName, new Type(type)), new JavaDocField(new JavaDocClass(className, "LOL"), fieldName, type, comment));
     }
 
-    public void addMethod(String methodName, String className, String signature, String comment, String... argComment)
+    public void addMethod(String methodName, String className, String signature, String comment, String returnComment, String[] argComment)
     {
         BehaviorEntry behaviorEntry;
         if (methodName.equals("<init>"))
@@ -58,7 +58,7 @@ public class JavaDocMapping
             behaviorEntry = new MethodEntry(new ClassEntry(className), methodName, new Signature(signature));
 
         this.javaDocMethodByID.put(behaviorEntry, new JavaDocMethod(new JavaDocClass(className, "LOL"), methodName, signature, comment,
-                Arrays.asList(argComment)));
+                returnComment, argComment));
     }
 
     public void cleanBehaviors()
@@ -121,10 +121,17 @@ public class JavaDocMapping
             // Prepare the annotation
             AnnotationsAttribute attr = new AnnotationsAttribute(constpool, AnnotationsAttribute.visibleTag);
             Annotation annot = new Annotation("enigma.remapper.MethodDoc", constpool);
-            annot.addMemberValue("comment", new StringMemberValue(javaDocMethod.getComment(), constpool));
+            String comment = javaDocMethod.getComment();
+            if (comment != null)
+                annot.addMemberValue("comment", new StringMemberValue(comment, constpool));
+
             behaviorByID.add(entry);
             annot.addMemberValue("behavior", new IntegerMemberValue(constpool, behaviorByID.size() - 1));
             annot.addMemberValue("args", translateStringArray(constpool, javaDocMethod.getArgsComments()));
+
+            String returnComment = javaDocMethod.getReturnComment();
+            if (returnComment != null)
+                annot.addMemberValue("return", new StringMemberValue(returnComment, constpool));
             attr.addAnnotation(annot);
             ctBehavior.getMethodInfo().addAttribute(attr);
         }
@@ -154,7 +161,7 @@ public class JavaDocMapping
     {
         MemberValue[] result = new MemberValue[args.length];
         for (int i = 0; i < result.length; i++)
-            result[i] = new StringMemberValue(args[i], pool);
+            result[i] = new StringMemberValue(args[i] == null ? "null" : args[i], pool);
 
         return result;
     }
@@ -202,15 +209,18 @@ public class JavaDocMapping
                 if (child instanceof PrimitiveExpression)
                 {
                     PrimitiveExpression data = (PrimitiveExpression) child;
-                    if (id.equals("comment"))
+                    switch (id)
                     {
-                        builder.append(spaces);
-                        builder.append(" * ");
-                        builder.append(data.getValue());
-                        builder.append("\n");
+                        case "comment":
+                            addCommentLine(builder, (String) data.getValue(), spaces);
+                            break;
+                        case "behavior":
+                            behaviorID = (Integer) data.getValue();
+                            break;
+                        case "return":
+                            addCommentLine(builder, "@return " + data.getValue(), spaces);
+                            break;
                     }
-                    else if (id.equals("behavior"))
-                        behaviorID = (Integer) data.getValue();
                 }
                 else if (child instanceof ArrayInitializerExpression)
                 {
@@ -226,16 +236,14 @@ public class JavaDocMapping
                         {
                             if (expArg instanceof PrimitiveExpression)
                             {
-                                PrimitiveExpression primitiveExpression = (PrimitiveExpression) expArg;
+                                PrimitiveExpression argComment = (PrimitiveExpression) expArg;
+                                String commentStr = (String) argComment.getValue();
+                                if (commentStr.equals("null"))
+                                    continue;
                                 String argName = entry == null ? null : this.translator.translate(new ArgumentEntry(entry, i, ""));
                                 if (argName == null)
                                     argName = "a" + (i + 1);
-                                builder.append(spaces);
-                                builder.append(" * @param ");
-                                builder.append(argName);
-                                builder.append(" ");
-                                builder.append(primitiveExpression.getValue());
-                                builder.append("\n");
+                                addCommentLine(builder, "@param " + argName + " " + argComment.getValue(), spaces);
                             }
                             i++;
                         }
@@ -248,6 +256,14 @@ public class JavaDocMapping
         return builder.toString();
     }
 
+    public void addCommentLine(StringBuilder builder, String value, String spaces)
+    {
+        builder.append(spaces);
+        builder.append(" * ");
+        builder.append(value);
+        builder.append("\n");
+    }
+
     private String buildLineSpace(int spacesCount)
     {
         StringBuilder builder = new StringBuilder();
diff --git a/src/main/java/cuchaz/enigma/mapping/javadoc/JavaDocMethod.java b/src/main/java/cuchaz/enigma/mapping/javadoc/JavaDocMethod.java
index 1508aede..6ed7ff8c 100644
--- a/src/main/java/cuchaz/enigma/mapping/javadoc/JavaDocMethod.java
+++ b/src/main/java/cuchaz/enigma/mapping/javadoc/JavaDocMethod.java
@@ -17,13 +17,16 @@ public class JavaDocMethod extends BaseJavaDocEntry
 {
     private final JavaDocClass classDocEntry;
     private final String       name;
-    private final List<String> argsComments;
+    private final String returnComment;
+    private final String[] argsComments;
 
-    public JavaDocMethod(JavaDocClass classDocEntry, String name, String identifier, String comment, List<String> argsComments)
+    public JavaDocMethod(JavaDocClass classDocEntry, String name, String identifier, String comment,
+            String returnComment, String[] argsComments)
     {
         super(identifier, comment);
         this.name = name;
         this.classDocEntry = classDocEntry;
+        this.returnComment = returnComment;
         this.argsComments = argsComments;
     }
 
@@ -39,11 +42,11 @@ public class JavaDocMethod extends BaseJavaDocEntry
 
     public String[] getArgsComments()
     {
-        return argsComments.toArray(new String[argsComments.size()]);
+        return argsComments;
     }
 
-    public String getArgsComment(int id)
+    public String getReturnComment()
     {
-        return argsComments.size() > id ? argsComments.get(id) : null;
+        return returnComment;
     }
 }