fix SourceIndex remap bugs

3 files changed, 15 insertions, 18 deletions

diff --git a/src/main/java/cuchaz/enigma/analysis/SourceIndex.java b/src/main/java/cuchaz/enigma/analysis/SourceIndex.java
index 4c84e69d..3e0d66b4 100644
--- a/src/main/java/cuchaz/enigma/analysis/SourceIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/SourceIndex.java
@@ -68,20 +68,21 @@ public class SourceIndex {
                         declarationToToken.put(entry, tokenMap.getOrDefault(token, token));
                 }
 
-                for (Token token : Lists.newArrayList(tokenToReference.keySet())) {
-                        EntryReference<Entry, Entry> e = tokenToReference.remove(token);
-                        tokenToReference.put(tokenMap.getOrDefault(token, token), e);
-                }
-
-                for (EntryReference<Entry, Entry> ref : Lists.newArrayList(referenceToTokens.keySet())) {
-                        List<Token> newTokens = new ArrayList<>();
+                for (EntryReference<Entry, Entry> ref : referenceToTokens.keySet()) {
+                        Collection<Token> oldTokens = referenceToTokens.get(ref);
+                        List<Token> newTokens = new ArrayList<>(oldTokens.size());
 
-                        for (Token token : referenceToTokens.get(ref)) {
+                        for (Token token : oldTokens) {
                                 newTokens.add(tokenMap.getOrDefault(token, token));
                         }
 
-                        referenceToTokens.removeAll(ref);
-                        referenceToTokens.putAll(ref, newTokens);
+                        referenceToTokens.replaceValues(ref, newTokens);
+                }
+
+                Map<Token, EntryReference<Entry, Entry>> tokenToReferenceCopy = Maps.newHashMap(tokenToReference);
+                tokenToReference.clear();
+                for (Token token : tokenToReferenceCopy.keySet()) {
+                        tokenToReference.put(tokenMap.getOrDefault(token, token), tokenToReferenceCopy.get(token));
                 }
         }
 
diff --git a/src/main/java/cuchaz/enigma/analysis/SourceIndexClassVisitor.java b/src/main/java/cuchaz/enigma/analysis/SourceIndexClassVisitor.java
index dd5bcef0..cad08574 100644
--- a/src/main/java/cuchaz/enigma/analysis/SourceIndexClassVisitor.java
+++ b/src/main/java/cuchaz/enigma/analysis/SourceIndexClassVisitor.java
@@ -24,9 +24,7 @@ import cuchaz.enigma.mapping.entry.*;
 public class SourceIndexClassVisitor extends SourceIndexVisitor {
         private final ReferencedEntryPool entryPool;
         private final ProcyonEntryFactory entryFactory;
-
         private ClassDefEntry classEntry;
-        private boolean isEnum;
 
         public SourceIndexClassVisitor(ReferencedEntryPool entryPool, ClassDefEntry classEntry) {
                 super(entryPool);
@@ -88,7 +86,6 @@ public class SourceIndexClassVisitor extends SourceIndexVisitor {
                 assert (node.getVariables().size() == 1);
                 VariableInitializer variable = node.getVariables().firstOrNullObject();
                 index.addDeclaration(variable.getNameToken(), fieldEntry);
-
                 return recurse(node, index);
         }
 
@@ -98,7 +95,6 @@ public class SourceIndexClassVisitor extends SourceIndexVisitor {
                 FieldDefinition def = node.getUserData(Keys.FIELD_DEFINITION);
                 FieldDefEntry fieldEntry = entryFactory.getFieldDefEntry(def);
                 index.addDeclaration(node.getNameToken(), fieldEntry);
-                this.isEnum = true;
                 return recurse(node, index);
         }
 }
diff --git a/src/main/java/cuchaz/enigma/gui/GuiController.java b/src/main/java/cuchaz/enigma/gui/GuiController.java
index e2f332d0..acb0ebbb 100644
--- a/src/main/java/cuchaz/enigma/gui/GuiController.java
+++ b/src/main/java/cuchaz/enigma/gui/GuiController.java
@@ -345,7 +345,7 @@ public class GuiController {
                         List<Token> otherTokens = Lists.newArrayList();
 
                         int offset = 0;
-                        Map<Token, Token> tokenRemap = new IdentityHashMap<>();
+                        Map<Token, Token> tokenRemap = new HashMap<>();
                         boolean remapped = false;
 
                         for (Token inToken : index.referenceTokens()) {
@@ -353,7 +353,7 @@ public class GuiController {
                                 Token token = inToken.move(offset);
 
                                 if (referenceIsRenameable(reference)) {
-                                        boolean renamed = false;
+                                        boolean added = false;
 
                                         if (!entryHasDeobfuscatedName(reference.getNameableEntry())) {
                                                 Entry obfEntry = deobfuscator.obfuscateEntry(reference.getNameableEntry());
@@ -364,7 +364,7 @@ public class GuiController {
                                                                         proposedTokens.add(token);
                                                                         offset += token.getRenameOffset(proposal);
                                                                         sourceString = token.rename(sourceString, proposal);
-                                                                        renamed = true;
+                                                                        added = true;
                                                                         remapped = true;
                                                                         break;
                                                                 }
@@ -372,7 +372,7 @@ public class GuiController {
                                                 }
                                         }
 
-                                        if (!renamed) {
+                                        if (!added) {
                                                 if (entryHasDeobfuscatedName(reference.getNameableEntry())) {
                                                         deobfuscatedTokens.add(token);
                                                 } else {