1 files changed, 21 insertions, 15 deletions

diff --git a/src/cuchaz/enigma/mapping/ClassMapping.java b/src/cuchaz/enigma/mapping/ClassMapping.java
index 095cb38..1219e7c 100644
--- a/src/cuchaz/enigma/mapping/ClassMapping.java
+++ b/src/cuchaz/enigma/mapping/ClassMapping.java
@@ -301,31 +301,37 @@ public class ClassMapping implements Serializable, Comparable<ClassMapping>
                 return m_obfName.compareTo( other.m_obfName );
         }
         
-        public void renameObfClasses( Map<String,String> nameMap )
+        public boolean renameObfClass( String oldObfClassName, String newObfClassName )
         {
-                // rename self
+                // rename inner classes
+                for( ClassMapping innerClassMapping : new ArrayList<ClassMapping>( m_innerClassesByObf.values() ) )
                 {
-                        String newName = nameMap.get( m_obfName );
-                        if( newName != null )
+                        if( innerClassMapping.renameObfClass( oldObfClassName, newObfClassName ) )
                         {
-                                m_obfName = newName;
+                                m_innerClassesByObf.remove( oldObfClassName );
+                                m_innerClassesByObf.put( newObfClassName, innerClassMapping );
+                                assert( m_innerClassesByObf.size() == m_innerClassesByDeobf.size() );
                         }
                 }
                 
-                // rename inner classes
-                for( ClassMapping classMapping : new ArrayList<ClassMapping>( m_innerClassesByObf.values() ) )
+                // rename method signatures
+                for( MethodMapping methodMapping : new ArrayList<MethodMapping>( m_methodsByObf.values() ) )
                 {
-                        m_innerClassesByObf.remove( classMapping.getObfName() );
-                        classMapping.renameObfClasses( nameMap );
-                        m_innerClassesByObf.put( classMapping.getObfName(), classMapping );
+                        String oldMethodKey = getMethodKey( methodMapping.getObfName(), methodMapping.getObfSignature() );
+                        if( methodMapping.renameObfClass( oldObfClassName, newObfClassName ) )
+                        {
+                                m_methodsByObf.remove( oldMethodKey );
+                                m_methodsByObf.put( getMethodKey( methodMapping.getObfName(), methodMapping.getObfSignature() ), methodMapping );
+                                assert( m_methodsByObf.size() == m_methodsByDeobf.size() );
+                        }
                 }
                 
-                // rename method signatures
-                for( MethodMapping methodMapping : new ArrayList<MethodMapping>( m_methodsByObf.values() ) )
+                if( m_obfName.equals( oldObfClassName ) )
                 {
-                        m_methodsByObf.remove( getMethodKey( methodMapping.getObfName(), methodMapping.getObfSignature() ) );
-                        methodMapping.renameObfClasses( nameMap );
-                        m_methodsByObf.put( getMethodKey( methodMapping.getObfName(), methodMapping.getObfSignature() ), methodMapping );
+                        // rename this class
+                        m_obfName = newObfClassName;
+                        return true;
                 }
+                return false;
         }
 }