8 files changed, 189 insertions, 26 deletions

diff --git a/src/main/java/cuchaz/enigma/mapping/ClassMapping.java b/src/main/java/cuchaz/enigma/mapping/ClassMapping.java
index 923c8ec..017a5b9 100644
--- a/src/main/java/cuchaz/enigma/mapping/ClassMapping.java
+++ b/src/main/java/cuchaz/enigma/mapping/ClassMapping.java
@@ -30,12 +30,20 @@ public class ClassMapping implements Comparable<ClassMapping> {
     private Map<String, MethodMapping> m_methodsByObf;
     private Map<String, MethodMapping> m_methodsByDeobf;
     private boolean isDirty;
+    private Mappings.EntryModifier modifier;
 
-    public ClassMapping(String obfFullName) {
-        this(obfFullName, null);
+    public ClassMapping(String obfFullName)
+    {
+        this(obfFullName, null, Mappings.EntryModifier.UNCHANGED);
     }
 
-    public ClassMapping(String obfFullName, String deobfName) {
+    public ClassMapping(String obfFullName, String deobfName)
+    {
+        this(obfFullName, deobfName, Mappings.EntryModifier.UNCHANGED);
+    }
+
+    public ClassMapping(String obfFullName, String deobfName, Mappings.EntryModifier modifier)
+    {
         m_obfFullName = obfFullName;
         ClassEntry classEntry = new ClassEntry(obfFullName);
         m_obfSimpleName = classEntry.isInnerClass() ? classEntry.getInnermostClassName() : classEntry.getSimpleName();
@@ -48,6 +56,7 @@ public class ClassMapping implements Comparable<ClassMapping> {
         m_methodsByObf = Maps.newHashMap();
         m_methodsByDeobf = Maps.newHashMap();
         isDirty = true;
+        this.modifier = modifier;
     }
 
     public String getObfFullName() {
@@ -186,15 +195,16 @@ public class ClassMapping implements Comparable<ClassMapping> {
         if (m_fieldsByObf.containsKey(obfKey)) {
             throw new Error("Already have mapping for " + m_obfFullName + "." + obfKey);
         }
-        String deobfKey = getFieldKey(fieldMapping.getDeobfName(), fieldMapping.getObfType());
-        if (m_fieldsByDeobf.containsKey(deobfKey)) {
-            throw new Error("Already have mapping for " + m_deobfName + "." + deobfKey);
+        if (fieldMapping.getDeobfName() != null) {
+            String deobfKey = getFieldKey(fieldMapping.getDeobfName(), fieldMapping.getObfType());
+            if (m_fieldsByDeobf.containsKey(deobfKey)) {
+                throw new Error("Already have mapping for " + m_deobfName + "." + deobfKey);
+            }
+            boolean deobfWasAdded = m_fieldsByDeobf.put(deobfKey, fieldMapping) == null;
+            assert (deobfWasAdded);
         }
         boolean obfWasAdded = m_fieldsByObf.put(obfKey, fieldMapping) == null;
         assert (obfWasAdded);
-        boolean deobfWasAdded = m_fieldsByDeobf.put(deobfKey, fieldMapping) == null;
-        assert (deobfWasAdded);
-        assert (m_fieldsByObf.size() == m_fieldsByDeobf.size());
         this.isDirty = true;
     }
 
@@ -242,12 +252,11 @@ public class ClassMapping implements Comparable<ClassMapping> {
         return name + ":" + type;
     }
 
-
     public void setFieldName(String obfName, Type obfType, String deobfName) {
         assert (deobfName != null);
         FieldMapping fieldMapping = m_fieldsByObf.get(getFieldKey(obfName, obfType));
         if (fieldMapping == null) {
-            fieldMapping = new FieldMapping(obfName, obfType, deobfName);
+            fieldMapping = new FieldMapping(obfName, obfType, deobfName, Mappings.EntryModifier.UNCHANGED);
             boolean obfWasAdded = m_fieldsByObf.put(getFieldKey(obfName, obfType), fieldMapping) == null;
             assert (obfWasAdded);
         } else {
@@ -492,4 +501,44 @@ public class ClassMapping implements Comparable<ClassMapping> {
     {
         this.isDirty = false;
     }
+
+    public void setModifier(Mappings.EntryModifier modifier)
+    {
+        if (this.modifier != modifier)
+            this.isDirty = true;
+        this.modifier = modifier;
+    }
+
+    public Mappings.EntryModifier getModifier()
+    {
+        return modifier;
+    }
+
+    public void setFieldModifier(String obfName, Type obfType, Mappings.EntryModifier modifier) {
+        FieldMapping fieldMapping = m_fieldsByObf.get(getFieldKey(obfName, obfType));
+        if (fieldMapping == null) {
+            fieldMapping = new FieldMapping(obfName, obfType, null, Mappings.EntryModifier.UNCHANGED);
+            m_fieldsByObf.put(getFieldKey(obfName, obfType), fieldMapping);
+        }
+
+        if (fieldMapping.getModifier() != modifier)
+        {
+            fieldMapping.setModifier(modifier);
+            this.isDirty = true;
+        }
+    }
+
+    public void setMethodModifier(String obfName, Signature sig, Mappings.EntryModifier modifier) {
+        MethodMapping methodMapping = m_methodsByObf.get(getMethodKey(obfName, sig));
+        if (methodMapping == null) {
+            methodMapping = new MethodMapping(obfName, sig, null, Mappings.EntryModifier.UNCHANGED);
+            m_methodsByObf.put(getMethodKey(obfName, sig), methodMapping);
+        }
+
+        if (methodMapping.getModifier() != modifier)
+        {
+            methodMapping.setModifier(modifier);
+            this.isDirty = true;
+        }
+    }
 }
diff --git a/src/main/java/cuchaz/enigma/mapping/FieldMapping.java b/src/main/java/cuchaz/enigma/mapping/FieldMapping.java
index 1b59660..e75485c 100644
--- a/src/main/java/cuchaz/enigma/mapping/FieldMapping.java
+++ b/src/main/java/cuchaz/enigma/mapping/FieldMapping.java
@@ -15,16 +15,19 @@ public class FieldMapping implements Comparable<FieldMapping>, MemberMapping<Fie
     private String obfName;
     private String deobfName;
     private Type obfType;
+    private Mappings.EntryModifier modifier;
 
-    public FieldMapping(String obfName, Type obfType, String deobfName) {
+    public FieldMapping(String obfName, Type obfType, String deobfName, Mappings.EntryModifier modifier) {
         this.obfName = obfName;
         this.deobfName = NameValidator.validateFieldName(deobfName);
         this.obfType = obfType;
+        this.modifier = modifier;
     }
 
     public FieldMapping(FieldMapping other, ClassNameReplacer obfClassNameReplacer) {
         this.obfName = other.obfName;
         this.deobfName = other.deobfName;
+        this.modifier = other.modifier;
         this.obfType = new Type(other.obfType, obfClassNameReplacer);
     }
 
@@ -58,6 +61,16 @@ public class FieldMapping implements Comparable<FieldMapping>, MemberMapping<Fie
         this.obfType = val;
     }
 
+    public void setModifier(Mappings.EntryModifier modifier)
+    {
+        this.modifier = modifier;
+    }
+
+    public Mappings.EntryModifier getModifier()
+    {
+        return modifier;
+    }
+
     @Override
     public int compareTo(FieldMapping other) {
         return (this.obfName + this.obfType).compareTo(other.obfName + other.obfType);
diff --git a/src/main/java/cuchaz/enigma/mapping/Mappings.java b/src/main/java/cuchaz/enigma/mapping/Mappings.java
index 2166bb9..d493dcf 100644
--- a/src/main/java/cuchaz/enigma/mapping/Mappings.java
+++ b/src/main/java/cuchaz/enigma/mapping/Mappings.java
@@ -240,4 +240,14 @@ public class Mappings {
     {
         ENIGMA_FILE, ENIGMA_DIRECTORY, SRG_FILE
     }
+
+    public enum EntryModifier
+    {
+        UNCHANGED, PUBLIC, PROTECTED, PRIVATE;
+
+        public String getFormattedName()
+        {
+            return " ACC:" + super.toString();
+        }
+    }
 }
diff --git a/src/main/java/cuchaz/enigma/mapping/MappingsEnigmaReader.java b/src/main/java/cuchaz/enigma/mapping/MappingsEnigmaReader.java
index 93276a6..cdfed72 100644
--- a/src/main/java/cuchaz/enigma/mapping/MappingsEnigmaReader.java
+++ b/src/main/java/cuchaz/enigma/mapping/MappingsEnigmaReader.java
@@ -135,21 +135,51 @@ public class MappingsEnigmaReader
     private ClassMapping readClass(String[] parts, boolean makeSimple) {
         if (parts.length == 2) {
             return new ClassMapping(parts[1]);
-        } else {
-            return new ClassMapping(parts[1], parts[2]);
-        }
+        } else if (parts.length == 3) {
+            boolean access = parts[2].startsWith("ACC:");
+            ClassMapping mapping;
+            if (access)
+                mapping = new ClassMapping(parts[1], null, Mappings.EntryModifier.valueOf(parts[2].substring(4)));
+            else
+                mapping = new ClassMapping(parts[1], parts[2]);
+
+            return mapping;
+        } else if (parts.length == 4)
+            return new ClassMapping(parts[1], parts[2], Mappings.EntryModifier.valueOf(parts[3].substring(4)));
+        return null;
     }
 
     /* TEMP */
     protected FieldMapping readField(String[] parts) {
-        return new FieldMapping(parts[1], new Type(parts[3]), parts[2]);
+        FieldMapping mapping = null;
+        if (parts.length == 4)
+        {
+            boolean access = parts[3].startsWith("ACC:");
+            if (access)
+                mapping = new FieldMapping(parts[1], new Type(parts[2]), null,
+                        Mappings.EntryModifier.valueOf(parts[3].substring(4)));
+            else
+                mapping = new FieldMapping(parts[1], new Type(parts[3]), parts[2], Mappings.EntryModifier.UNCHANGED);
+        }
+        else if (parts.length == 5)
+            mapping = new FieldMapping(parts[1], new Type(parts[3]), parts[2], Mappings.EntryModifier.valueOf(parts[4].substring(4)));
+        return mapping;
     }
 
     private MethodMapping readMethod(String[] parts) {
-        if (parts.length == 3) {
-            return new MethodMapping(parts[1], new Signature(parts[2]));
-        } else {
-            return new MethodMapping(parts[1], new Signature(parts[3]), parts[2]);
+        MethodMapping mapping = null;
+        if (parts.length == 3)
+            mapping = new MethodMapping(parts[1], new Signature(parts[2]));
+        else if (parts.length == 4){
+            boolean access = parts[3].startsWith("ACC:");
+            if (access)
+                mapping = new MethodMapping(parts[1], new Signature(parts[2]), null, Mappings.EntryModifier.valueOf(parts[3].substring(4)));
+            else
+                mapping = new MethodMapping(parts[1], new Signature(parts[3]), parts[2]);
         }
+        else if (parts.length == 5)
+            mapping = new MethodMapping(parts[1], new Signature(parts[3]), parts[2],
+                    Mappings.EntryModifier.valueOf(parts[4].substring(4)));
+        return mapping;
     }
 }
diff --git a/src/main/java/cuchaz/enigma/mapping/MappingsEnigmaWriter.java b/src/main/java/cuchaz/enigma/mapping/MappingsEnigmaWriter.java
index e1763d0..9589689 100644
--- a/src/main/java/cuchaz/enigma/mapping/MappingsEnigmaWriter.java
+++ b/src/main/java/cuchaz/enigma/mapping/MappingsEnigmaWriter.java
@@ -110,9 +110,9 @@ public class MappingsEnigmaWriter {
         
         private void write(PrintWriter out, ClassMapping classMapping, int depth) throws IOException {
                 if (classMapping.getDeobfName() == null) {
-                        out.format("%sCLASS %s\n", getIndent(depth), classMapping.getObfFullName());
+                        out.format("%sCLASS %s%s\n", getIndent(depth), classMapping.getObfFullName(), classMapping.getModifier() == Mappings.EntryModifier.UNCHANGED ? "" : classMapping.getModifier().getFormattedName());
                 } else {
-                        out.format("%sCLASS %s %s\n", getIndent(depth), classMapping.getObfFullName(), classMapping.getDeobfName());
+                        out.format("%sCLASS %s %s%s\n", getIndent(depth), classMapping.getObfFullName(), classMapping.getDeobfName(), classMapping.getModifier() == Mappings.EntryModifier.UNCHANGED ? "" :  classMapping.getModifier().getFormattedName());
                 }
                 
                 for (ClassMapping innerClassMapping : sorted(classMapping.innerClasses())) {
@@ -129,14 +129,17 @@ public class MappingsEnigmaWriter {
         }
         
         private void write(PrintWriter out, FieldMapping fieldMapping, int depth) throws IOException {
-                out.format("%sFIELD %s %s %s\n", getIndent(depth), fieldMapping.getObfName(), fieldMapping.getDeobfName(), fieldMapping.getObfType().toString());
+                if (fieldMapping.getDeobfName() == null)
+                        out.format("%sFIELD %s %s%s\n", getIndent(depth), fieldMapping.getObfName(), fieldMapping.getObfType().toString(), fieldMapping.getModifier() == Mappings.EntryModifier.UNCHANGED ? "" : fieldMapping.getModifier().getFormattedName());
+                else
+                        out.format("%sFIELD %s %s %s%s\n", getIndent(depth), fieldMapping.getObfName(), fieldMapping.getDeobfName(), fieldMapping.getObfType().toString(), fieldMapping.getModifier() == Mappings.EntryModifier.UNCHANGED ? "" : fieldMapping.getModifier().getFormattedName());
         }
         
         private void write(PrintWriter out, MethodMapping methodMapping, int depth) throws IOException {
                 if (methodMapping.getDeobfName() == null) {
-                        out.format("%sMETHOD %s %s\n", getIndent(depth), methodMapping.getObfName(), methodMapping.getObfSignature());
+                        out.format("%sMETHOD %s %s%s\n", getIndent(depth), methodMapping.getObfName(), methodMapping.getObfSignature(), methodMapping.getModifier() == Mappings.EntryModifier.UNCHANGED ? "" :methodMapping.getModifier().getFormattedName());
                 } else {
-                        out.format("%sMETHOD %s %s %s\n", getIndent(depth), methodMapping.getObfName(), methodMapping.getDeobfName(), methodMapping.getObfSignature());
+                        out.format("%sMETHOD %s %s %s%s\n", getIndent(depth), methodMapping.getObfName(), methodMapping.getDeobfName(), methodMapping.getObfSignature(), methodMapping.getModifier() == Mappings.EntryModifier.UNCHANGED ? "" : methodMapping.getModifier().getFormattedName());
                 }
                 
                 for (ArgumentMapping argumentMapping : sorted(methodMapping.arguments())) {
diff --git a/src/main/java/cuchaz/enigma/mapping/MappingsRenamer.java b/src/main/java/cuchaz/enigma/mapping/MappingsRenamer.java
index 4e812b1..bac6250 100644
--- a/src/main/java/cuchaz/enigma/mapping/MappingsRenamer.java
+++ b/src/main/java/cuchaz/enigma/mapping/MappingsRenamer.java
@@ -321,4 +321,22 @@ public class MappingsRenamer {
         }
         return mappingChain;
     }
+
+    public void setClassModifier(ClassEntry obEntry, Mappings.EntryModifier modifier)
+    {
+        ClassMapping classMapping = getOrCreateClassMapping(obEntry);
+        classMapping.setModifier(modifier);
+    }
+
+    public void setFieldModifier(FieldEntry obEntry, Mappings.EntryModifier modifier)
+    {
+        ClassMapping classMapping = getOrCreateClassMapping(obEntry.getClassEntry());
+        classMapping.setFieldModifier(obEntry.getName(), obEntry.getType(), modifier);
+    }
+
+    public void setMethodModifier(BehaviorEntry obEntry, Mappings.EntryModifier modifier)
+    {
+        ClassMapping classMapping = getOrCreateClassMapping(obEntry.getClassEntry());
+        classMapping.setMethodModifier(obEntry.getName(), obEntry.getSignature(), modifier);
+    }
 }
diff --git a/src/main/java/cuchaz/enigma/mapping/MethodMapping.java b/src/main/java/cuchaz/enigma/mapping/MethodMapping.java
index 99b9c88..455ff6b 100644
--- a/src/main/java/cuchaz/enigma/mapping/MethodMapping.java
+++ b/src/main/java/cuchaz/enigma/mapping/MethodMapping.java
@@ -22,12 +22,17 @@ public class MethodMapping implements Comparable<MethodMapping>, MemberMapping<B
     private String deobfName;
     private Signature obfSignature;
     private Map<Integer, ArgumentMapping> arguments;
+    private Mappings.EntryModifier modifier;
 
     public MethodMapping(String obfName, Signature obfSignature) {
-        this(obfName, obfSignature, null);
+        this(obfName, obfSignature, null,Mappings.EntryModifier.UNCHANGED);
     }
 
     public MethodMapping(String obfName, Signature obfSignature, String deobfName) {
+        this(obfName, obfSignature, deobfName, Mappings.EntryModifier.UNCHANGED);
+    }
+
+    public MethodMapping(String obfName, Signature obfSignature, String deobfName, Mappings.EntryModifier modifier) {
         if (obfName == null) {
             throw new IllegalArgumentException("obf name cannot be null!");
         }
@@ -38,11 +43,13 @@ public class MethodMapping implements Comparable<MethodMapping>, MemberMapping<B
         this.deobfName = NameValidator.validateMethodName(deobfName);
         this.obfSignature = obfSignature;
         this.arguments = Maps.newTreeMap();
+        this.modifier = modifier;
     }
     
     public MethodMapping(MethodMapping other, ClassNameReplacer obfClassNameReplacer) {
         this.obfName = other.obfName;
         this.deobfName = other.deobfName;
+        this.modifier = other.modifier;
         this.obfSignature = new Signature(other.obfSignature, obfClassNameReplacer);
         this.arguments = Maps.newTreeMap();
         for (Map.Entry<Integer,ArgumentMapping> entry : other.arguments.entrySet()) {
@@ -187,4 +194,14 @@ public class MethodMapping implements Comparable<MethodMapping>, MemberMapping<B
             return new MethodEntry(classEntry, this.obfName, this.obfSignature);
         }
     }
+
+    public Mappings.EntryModifier getModifier()
+    {
+        return modifier;
+    }
+
+    public void setModifier(Mappings.EntryModifier modifier)
+    {
+        this.modifier = modifier;
+    }
 }
diff --git a/src/main/java/cuchaz/enigma/mapping/Translator.java b/src/main/java/cuchaz/enigma/mapping/Translator.java
index 6b636f2..e94009e 100644
--- a/src/main/java/cuchaz/enigma/mapping/Translator.java
+++ b/src/main/java/cuchaz/enigma/mapping/Translator.java
@@ -329,4 +329,27 @@ public class Translator {
         assert (mappingsChain.size() == parts.length);
         return mappingsChain;
     }
+
+    public Mappings.EntryModifier getModifier(Entry entry)
+    {
+        ClassMapping classMapping = findClassMapping(entry.getClassEntry());
+        if (classMapping != null && !entry.getName().equals("<clinit>"))
+        {
+            if (entry instanceof ClassEntry)
+                return classMapping.getModifier();
+            else if (entry instanceof FieldEntry)
+            {
+                FieldMapping fieldMapping = classMapping.getFieldByObf(entry.getName(), ((FieldEntry) entry).getType());
+                return fieldMapping != null ? fieldMapping.getModifier() : Mappings.EntryModifier.UNCHANGED;
+            }
+            else if (entry instanceof BehaviorEntry)
+            {
+                MethodMapping methodMapping = classMapping.getMethodByObf(entry.getName(), ((BehaviorEntry) entry).getSignature());
+                return methodMapping != null ? methodMapping.getModifier() : Mappings.EntryModifier.UNCHANGED;
+            }
+            else
+                throw new Error("Unknown entry type: " + entry.getClass().getName());
+        }
+        return Mappings.EntryModifier.UNCHANGED;
+    }
 }