implemented mark-as-deobfuscated and reset-to-obfuscated

4 files changed, 164 insertions, 12 deletions

diff --git a/src/cuchaz/enigma/mapping/ClassMapping.java b/src/cuchaz/enigma/mapping/ClassMapping.java
index b551d71..88106df 100644
--- a/src/cuchaz/enigma/mapping/ClassMapping.java
+++ b/src/cuchaz/enigma/mapping/ClassMapping.java
@@ -79,6 +79,17 @@ public class ClassMapping implements Serializable, Comparable<ClassMapping>
                 }
         }
         
+        public void removeInnerClassMapping( ClassMapping classMapping )
+        {
+                boolean obfWasRemoved = m_innerClassesByObf.remove( classMapping.getObfName() ) != null;
+                assert( obfWasRemoved );
+                if( classMapping.getDeobfName() != null )
+                {
+                        boolean deobfWasRemoved = m_innerClassesByDeobf.remove( classMapping.getDeobfName() ) != null;
+                        assert( deobfWasRemoved );
+                }
+        }
+        
         public ClassMapping getOrCreateInnerClass( String obfName )
         {
                 ClassMapping classMapping = m_innerClassesByObf.get( obfName );
@@ -124,11 +135,17 @@ public class ClassMapping implements Serializable, Comparable<ClassMapping>
         public void setInnerClassName( String obfName, String deobfName )
         {
                 ClassMapping classMapping = getOrCreateInnerClass( obfName );
-                boolean wasRemoved = m_innerClassesByDeobf.remove( classMapping.getDeobfName() ) != null;
-                assert( wasRemoved );
+                if( classMapping.getDeobfName() != null )
+                {
+                        boolean wasRemoved = m_innerClassesByDeobf.remove( classMapping.getDeobfName() ) != null;
+                        assert( wasRemoved );
+                }
                 classMapping.setDeobfName( deobfName );
-                boolean wasAdded = m_innerClassesByDeobf.put( deobfName, classMapping ) == null;
-                assert( wasAdded );
+                if( deobfName != null )
+                {
+                        boolean wasAdded = m_innerClassesByDeobf.put( deobfName, classMapping ) == null;
+                        assert( wasAdded );
+                }
         }
         
         //// FIELDS ////////
@@ -176,6 +193,16 @@ public class ClassMapping implements Serializable, Comparable<ClassMapping>
                         assert( deobfWasRemoved );
                 }
         }
+        
+        public FieldMapping getFieldByObf( String obfName )
+        {
+                return m_fieldsByObf.get( obfName );
+        }
+        
+        public FieldMapping getFieldByDeobf( String deobfName )
+        {
+                return m_fieldsByDeobf.get( deobfName );
+        }
 
         public String getObfFieldName( String deobfName )
         {
@@ -212,8 +239,11 @@ public class ClassMapping implements Serializable, Comparable<ClassMapping>
                         assert( wasRemoved );
                 }
                 fieldMapping.setDeobfName( deobfName );
-                boolean wasAdded = m_fieldsByDeobf.put( deobfName, fieldMapping ) == null;
-                assert( wasAdded );
+                if( deobfName != null )
+                {
+                        boolean wasAdded = m_fieldsByDeobf.put( deobfName, fieldMapping ) == null;
+                        assert( wasAdded );
+                }
         }
         
         //// METHODS ////////
@@ -303,20 +333,28 @@ public class ClassMapping implements Serializable, Comparable<ClassMapping>
                         assert( wasRemoved );
                 }
                 methodMapping.setDeobfName( deobfName );
-                boolean wasAdded = m_methodsByDeobf.put( getMethodKey( deobfName, obfSignature ), methodMapping ) == null;
-                assert( wasAdded );
+                if( deobfName != null )
+                {
+                        boolean wasAdded = m_methodsByDeobf.put( getMethodKey( deobfName, obfSignature ), methodMapping ) == null;
+                        assert( wasAdded );
+                }
         }
         
         //// ARGUMENTS ////////
         
         public void setArgumentName( String obfMethodName, String obfMethodSignature, int argumentIndex, String argumentName )
         {
-                MethodMapping methodIndex = m_methodsByObf.get( getMethodKey( obfMethodName, obfMethodSignature ) );
-                if( methodIndex == null )
+                MethodMapping methodMapping = m_methodsByObf.get( getMethodKey( obfMethodName, obfMethodSignature ) );
+                if( methodMapping == null )
                 {
-                        methodIndex = createMethodMapping( obfMethodName, obfMethodSignature );
+                        methodMapping = createMethodMapping( obfMethodName, obfMethodSignature );
                 }
-                methodIndex.setArgumentName( argumentIndex, argumentName );
+                methodMapping.setArgumentName( argumentIndex, argumentName );
+        }
+        
+        public void removeArgumentName( String obfMethodName, String obfMethodSignature, int argumentIndex )
+        {
+                m_methodsByObf.get( getMethodKey( obfMethodName, obfMethodSignature ) ).removeArgumentName( argumentIndex );
         }
 
         private MethodMapping createMethodMapping( String obfName, String obfSignature )
diff --git a/src/cuchaz/enigma/mapping/Mappings.java b/src/cuchaz/enigma/mapping/Mappings.java
index 0b4e7f3..45b41bc 100644
--- a/src/cuchaz/enigma/mapping/Mappings.java
+++ b/src/cuchaz/enigma/mapping/Mappings.java
@@ -94,6 +94,17 @@ public class Mappings implements Serializable
                 }
         }
         
+        public void removeClassMapping( ClassMapping classMapping )
+        {
+                boolean obfWasRemoved = m_classesByObf.remove( classMapping.getObfName() ) != null;
+                assert( obfWasRemoved );
+                if( classMapping.getDeobfName() != null )
+                {
+                        boolean deobfWasRemoved = m_classesByDeobf.remove( classMapping.getDeobfName() ) != null;
+                        assert( deobfWasRemoved );
+                }
+        }
+        
         public ClassMapping getClassByObf( ClassEntry entry )
         {
                 return getClassByObf( entry.getName() );
diff --git a/src/cuchaz/enigma/mapping/MappingsRenamer.java b/src/cuchaz/enigma/mapping/MappingsRenamer.java
index 182bbd1..957b6d6 100644
--- a/src/cuchaz/enigma/mapping/MappingsRenamer.java
+++ b/src/cuchaz/enigma/mapping/MappingsRenamer.java
@@ -57,6 +57,36 @@ public class MappingsRenamer
                 }
         }
         
+        public void removeClassMapping( ClassEntry obf )
+        {
+                ClassMapping classMapping = getClassMapping( obf );
+                if( obf.isInnerClass() )
+                {
+                        classMapping.setInnerClassName( obf.getName(), null );
+                }
+                else
+                {
+                        boolean wasRemoved = m_mappings.m_classesByDeobf.remove( classMapping.getDeobfName() ) != null;
+                        assert( wasRemoved );
+                        classMapping.setDeobfName( null );
+                }
+        }
+        
+        public void markClassAsDeobfuscated( ClassEntry obf )
+        {
+                ClassMapping classMapping = getOrCreateClassMapping( obf );
+                if( obf.isInnerClass() )
+                {
+                        classMapping.setInnerClassName( obf.getName(), obf.getName() );
+                }
+                else
+                {
+                        classMapping.setDeobfName( obf.getName() );
+                        boolean wasAdded = m_mappings.m_classesByDeobf.put( obf.getName(), classMapping ) == null;
+                        assert( wasAdded );
+                }
+        }
+        
         public void setFieldName( FieldEntry obf, String deobfName )
         {
                 deobfName = NameValidator.validateFieldName( deobfName );
@@ -70,6 +100,18 @@ public class MappingsRenamer
                 classMapping.setFieldName( obf.getName(), deobfName );
         }
         
+        public void removeFieldMapping( FieldEntry obf )
+        {
+                ClassMapping classMapping = getClassMappingOrInnerClassMapping( obf.getClassEntry() );
+                classMapping.setFieldName( obf.getName(), null );
+        }
+        
+        public void markFieldAsDeobfuscated( FieldEntry obf )
+        {
+                ClassMapping classMapping = getOrCreateClassMappingOrInnerClassMapping( obf.getClassEntry() );
+                classMapping.setFieldName( obf.getName(), obf.getName() );
+        }
+        
         public void setMethodTreeName( MethodEntry obf, String deobfName )
         {
                 Set<MethodEntry> implementations = m_index.getRelatedMethodImplementations( obf );
@@ -106,6 +148,34 @@ public class MappingsRenamer
                 classMapping.setMethodName( obf.getName(), obf.getSignature(), deobfName );
         }
         
+        public void removeMethodTreeMapping( MethodEntry obf )
+        {
+                for( MethodEntry implementation : m_index.getRelatedMethodImplementations( obf ) )
+                {
+                        removeMethodMapping( implementation );
+                }
+        }
+        
+        public void removeMethodMapping( MethodEntry obf )
+        {
+                ClassMapping classMapping = getOrCreateClassMappingOrInnerClassMapping( obf.getClassEntry() );
+                classMapping.setMethodName( obf.getName(), obf.getSignature(), null );
+        }
+        
+        public void markMethodTreeAsDeobfuscated( MethodEntry obf )
+        {
+                for( MethodEntry implementation : m_index.getRelatedMethodImplementations( obf ) )
+                {
+                        markMethodAsDeobfuscated( implementation );
+                }
+        }
+        
+        public void markMethodAsDeobfuscated( MethodEntry obf )
+        {
+                ClassMapping classMapping = getOrCreateClassMappingOrInnerClassMapping( obf.getClassEntry() );
+                classMapping.setMethodName( obf.getName(), obf.getSignature(), obf.getName() );
+        }
+        
         public void setArgumentName( ArgumentEntry obf, String deobfName )
         {
                 deobfName = NameValidator.validateArgumentName( deobfName );
@@ -119,6 +189,18 @@ public class MappingsRenamer
                 classMapping.setArgumentName( obf.getMethodName(), obf.getMethodSignature(), obf.getIndex(), deobfName );
         }
         
+        public void removeArgumentMapping( ArgumentEntry obf )
+        {
+                ClassMapping classMapping = getClassMappingOrInnerClassMapping( obf.getClassEntry() );
+                classMapping.removeArgumentName( obf.getMethodName(), obf.getMethodSignature(), obf.getIndex() );
+        }
+        
+        public void markArgumentAsDeobfuscated( ArgumentEntry obf )
+        {
+                ClassMapping classMapping = getOrCreateClassMappingOrInnerClassMapping( obf.getClassEntry() );
+                classMapping.setArgumentName( obf.getMethodName(), obf.getMethodSignature(), obf.getIndex(), obf.getName() );
+        }
+        
         public boolean moveFieldToObfClass( ClassMapping classMapping, FieldMapping fieldMapping, ClassEntry obfClass )
         {
                 classMapping.removeFieldMapping( fieldMapping );
@@ -167,6 +249,11 @@ public class MappingsRenamer
                 gzipout.finish();
         }
         
+        private ClassMapping getClassMapping( ClassEntry obfClassEntry )
+        {
+                return m_mappings.m_classesByObf.get( obfClassEntry.getOuterClassName() );
+        }
+        
         private ClassMapping getOrCreateClassMapping( ClassEntry obfClassEntry )
         {
                 String obfClassName = obfClassEntry.getOuterClassName();
@@ -180,6 +267,16 @@ public class MappingsRenamer
                 return classMapping;
         }
         
+        private ClassMapping getClassMappingOrInnerClassMapping( ClassEntry obfClassEntry )
+        {
+                ClassMapping classMapping = getClassMapping( obfClassEntry );
+                if( obfClassEntry.isInDefaultPackage() )
+                {
+                        classMapping = classMapping.getInnerClassByObf( obfClassEntry.getInnerClassName() );
+                }
+                return classMapping;
+        }
+        
         private ClassMapping getOrCreateClassMappingOrInnerClassMapping( ClassEntry obfClassEntry )
         {
                 ClassMapping classMapping = getOrCreateClassMapping( obfClassEntry );
diff --git a/src/cuchaz/enigma/mapping/MethodMapping.java b/src/cuchaz/enigma/mapping/MethodMapping.java
index b076fa3..c51b011 100644
--- a/src/cuchaz/enigma/mapping/MethodMapping.java
+++ b/src/cuchaz/enigma/mapping/MethodMapping.java
@@ -118,6 +118,12 @@ public class MethodMapping implements Serializable, Comparable<MethodMapping>
                 }
         }
         
+        public void removeArgumentName( int index )
+        {
+                boolean wasRemoved = m_arguments.remove( index ) != null;
+                assert( wasRemoved );
+        }
+        
         @Override
         public String toString( )
         {