fix SourceIndex remap bugs

2 files changed, 11 insertions, 14 deletions

diff --git a/src/main/java/cuchaz/enigma/analysis/SourceIndex.java b/src/main/java/cuchaz/enigma/analysis/SourceIndex.java
index 4c84e69..3e0d66b 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 dd5bcef..cad0857 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);
         }
 }