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 +++++++++++----------
 .../enigma/analysis/SourceIndexClassVisitor.java    |  4 ----
 src/main/java/cuchaz/enigma/gui/GuiController.java  |  8 ++++----
 3 files changed, 15 insertions(+), 18 deletions(-)

(limited to 'src/main/java/cuchaz')

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);
 	}
 }
diff --git a/src/main/java/cuchaz/enigma/gui/GuiController.java b/src/main/java/cuchaz/enigma/gui/GuiController.java
index e2f332d..acb0ebb 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 {
-- 
cgit v1.2.3