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

7 files changed, 296 insertions, 29 deletions

diff --git a/src/cuchaz/enigma/Deobfuscator.java b/src/cuchaz/enigma/Deobfuscator.java
index 9a339176..44845ba2 100644
--- a/src/cuchaz/enigma/Deobfuscator.java
+++ b/src/cuchaz/enigma/Deobfuscator.java
@@ -508,23 +508,19 @@ public class Deobfuscator
                 Translator translator = getTranslator( TranslationDirection.Deobfuscating );
                 if( obfEntry instanceof ClassEntry )
                 {
-                        String deobfName = translator.translate( (ClassEntry)obfEntry );
-                        return deobfName != null && !deobfName.equals( obfEntry.getName() );
+                        return translator.translate( (ClassEntry)obfEntry ) != null;
                 }
                 else if( obfEntry instanceof FieldEntry )
                 {
-                        String deobfName = translator.translate( (FieldEntry)obfEntry );
-                        return deobfName != null && !deobfName.equals( obfEntry.getName() );
+                        return translator.translate( (FieldEntry)obfEntry ) != null;
                 }
                 else if( obfEntry instanceof MethodEntry )
                 {
-                        String deobfName = translator.translate( (MethodEntry)obfEntry );
-                        return deobfName != null && !deobfName.equals( obfEntry.getName() );
+                        return translator.translate( (MethodEntry)obfEntry ) != null;
                 }
                 else if( obfEntry instanceof ConstructorEntry )
                 {
-                        String deobfName = translator.translate( obfEntry.getClassEntry() );
-                        return deobfName != null && !deobfName.equals( obfEntry.getClassName() );
+                        return translator.translate( obfEntry.getClassEntry() ) != null;
                 }
                 else if( obfEntry instanceof ArgumentEntry )
                 {
@@ -540,4 +536,60 @@ public class Deobfuscator
         {
                 return m_jarIndex.containsObfEntry( obfEntry );
         }
+        
+        public void removeMapping( Entry obfEntry )
+        {
+                if( obfEntry instanceof ClassEntry )
+                {
+                        m_renamer.removeClassMapping( (ClassEntry)obfEntry );
+                }
+                else if( obfEntry instanceof FieldEntry )
+                {
+                        m_renamer.removeFieldMapping( (FieldEntry)obfEntry );
+                }
+                else if( obfEntry instanceof MethodEntry )
+                {
+                        m_renamer.removeMethodTreeMapping( (MethodEntry)obfEntry );
+                }
+                else if( obfEntry instanceof ConstructorEntry )
+                {
+                        m_renamer.removeClassMapping( obfEntry.getClassEntry() );
+                }
+                else if( obfEntry instanceof ArgumentEntry )
+                {
+                        m_renamer.removeArgumentMapping( (ArgumentEntry)obfEntry );
+                }
+                else
+                {
+                        throw new Error( "Unknown entry type: " + obfEntry );
+                }
+        }
+        
+        public void markAsDeobfuscated( Entry obfEntry )
+        {
+                if( obfEntry instanceof ClassEntry )
+                {
+                        m_renamer.markClassAsDeobfuscated( (ClassEntry)obfEntry );
+                }
+                else if( obfEntry instanceof FieldEntry )
+                {
+                        m_renamer.markFieldAsDeobfuscated( (FieldEntry)obfEntry );
+                }
+                else if( obfEntry instanceof MethodEntry )
+                {
+                        m_renamer.markMethodTreeAsDeobfuscated( (MethodEntry)obfEntry );
+                }
+                else if( obfEntry instanceof ConstructorEntry )
+                {
+                        m_renamer.markClassAsDeobfuscated( obfEntry.getClassEntry() );
+                }
+                else if( obfEntry instanceof ArgumentEntry )
+                {
+                        m_renamer.markArgumentAsDeobfuscated( (ArgumentEntry)obfEntry );
+                }
+                else
+                {
+                        throw new Error( "Unknown entry type: " + obfEntry );
+                }
+        }
 }
diff --git a/src/cuchaz/enigma/gui/Gui.java b/src/cuchaz/enigma/gui/Gui.java
index 1f04aa35..8bf6ce95 100644
--- a/src/cuchaz/enigma/gui/Gui.java
+++ b/src/cuchaz/enigma/gui/Gui.java
@@ -122,6 +122,7 @@ public class Gui
         private JMenuItem m_openPreviousMenu;
         private JMenuItem m_showCallsMenu;
         private JMenuItem m_showImplementationsMenu;
+        private JMenuItem m_toggleMappingMenu;
         
         // state
         private EntryReference<Entry,Entry> m_reference;
@@ -136,17 +137,20 @@ public class Gui
                 final Container pane = m_frame.getContentPane();
                 pane.setLayout( new BorderLayout() );
                 
-                // install a global exception handler to the event thread
-                CrashDialog.init( m_frame );
-                Thread.setDefaultUncaughtExceptionHandler( new UncaughtExceptionHandler( )
+                if( Boolean.parseBoolean( System.getProperty( "enigma.catchExceptions", "true" ) ) )
                 {
-                        @Override
-                        public void uncaughtException( Thread thread, Throwable ex )
+                        // install a global exception handler to the event thread
+                        CrashDialog.init( m_frame );
+                        Thread.setDefaultUncaughtExceptionHandler( new UncaughtExceptionHandler( )
                         {
-                                ex.printStackTrace( System.err );
-                                CrashDialog.show( ex );
-                        }
-                } );
+                                @Override
+                                public void uncaughtException( Thread thread, Throwable ex )
+                                {
+                                        ex.printStackTrace( System.err );
+                                        CrashDialog.show( ex );
+                                }
+                        } );
+                }
                 
                 m_controller = new GuiController( this );
                 
@@ -251,6 +255,10 @@ public class Gui
                                         case KeyEvent.VK_C:
                                                 m_showCallsMenu.doClick();
                                         break;
+                                        
+                                        case KeyEvent.VK_T:
+                                                m_toggleMappingMenu.doClick();
+                                        break;
                                 }
                         }
                 } );
@@ -352,6 +360,21 @@ public class Gui
                         popupMenu.add( menu );
                         m_openPreviousMenu = menu;
                 }
+                {
+                        JMenuItem menu = new JMenuItem( "Mark as deobfuscated" );
+                        menu.addActionListener( new ActionListener( )
+                        {
+                                @Override
+                                public void actionPerformed( ActionEvent event )
+                                {
+                                        toggleMapping();
+                                }
+                        } );
+                        menu.setAccelerator( KeyStroke.getKeyStroke( KeyEvent.VK_T, 0 ) );
+                        menu.setEnabled( false );
+                        popupMenu.add( menu );
+                        m_toggleMappingMenu = menu;
+                }
                 
                 // init inheritance panel
                 m_inheritanceTree = new JTree();
@@ -1031,6 +1054,16 @@ public class Gui
                 m_showCallsMenu.setEnabled( isClassEntry || isFieldEntry || isMethodEntry || isConstructorEntry );
                 m_openEntryMenu.setEnabled( isInJar && ( isClassEntry || isFieldEntry || isMethodEntry || isConstructorEntry ) );
                 m_openPreviousMenu.setEnabled( m_controller.hasPreviousLocation() );
+                m_toggleMappingMenu.setEnabled( isInJar && isToken );
+                
+                if( isToken && m_controller.entryHasMapping( m_reference.entry ) )
+                {
+                        m_toggleMappingMenu.setText( "Reset to obfuscated" );
+                }
+                else
+                {
+                        m_toggleMappingMenu.setText( "Mark as deobfuscated" );
+                }
         }
         
         private void navigateTo( Entry entry )
@@ -1233,6 +1266,18 @@ public class Gui
                 redraw();
         }
         
+        private void toggleMapping()
+        {
+                if( m_controller.entryHasMapping( m_reference.entry ) )
+                {
+                        m_controller.removeMapping( m_reference );
+                }
+                else
+                {
+                        m_controller.markAsDeobfuscated( m_reference );
+                }
+        }
+        
         private TreePath getPathToRoot( TreeNode node )
         {
                 List<TreeNode> nodes = Lists.newArrayList();
diff --git a/src/cuchaz/enigma/gui/GuiController.java b/src/cuchaz/enigma/gui/GuiController.java
index 098e065d..3adaf91d 100644
--- a/src/cuchaz/enigma/gui/GuiController.java
+++ b/src/cuchaz/enigma/gui/GuiController.java
@@ -249,6 +249,24 @@ public class GuiController
                 refreshCurrentClass( obfReference );
         }
         
+        public void removeMapping( EntryReference<Entry,Entry> deobfReference )
+        {
+                EntryReference<Entry,Entry> obfReference = m_deobfuscator.obfuscateReference( deobfReference );
+                m_deobfuscator.removeMapping( obfReference.entry );
+                m_isDirty = true;
+                refreshClasses();
+                refreshCurrentClass( obfReference );
+        }
+        
+        public void markAsDeobfuscated( EntryReference<Entry,Entry> deobfReference )
+        {
+                EntryReference<Entry,Entry> obfReference = m_deobfuscator.obfuscateReference( deobfReference );
+                m_deobfuscator.markAsDeobfuscated( obfReference.entry );
+                m_isDirty = true;
+                refreshClasses();
+                refreshCurrentClass( obfReference );
+        }
+        
         public void openDeclaration( Entry deobfEntry )
         {
                 if( deobfEntry == null )
diff --git a/src/cuchaz/enigma/mapping/ClassMapping.java b/src/cuchaz/enigma/mapping/ClassMapping.java
index b551d71c..88106dfa 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 0b4e7f3c..45b41bcd 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 182bbd19..957b6d68 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 b076fa33..c51b0110 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( )
         {