From ed51dfd4abf7f1955d5a8fdcd43dcc9e0b40a075 Mon Sep 17 00:00:00 2001
From: asie
Date: Sun, 9 Dec 2018 09:18:12 +0100
Subject: fix SourceIndex remap bugs

---
 .../java/cuchaz/enigma/analysis/SourceIndex.java    | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

(limited to 'src/main/java/cuchaz/enigma/analysis/SourceIndex.java')

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));
 		}
 	}
 
-- 
cgit v1.2.3