1 files changed, 26 insertions, 27 deletions

diff --git a/src/main/java/cuchaz/enigma/analysis/SourceIndex.java b/src/main/java/cuchaz/enigma/analysis/SourceIndex.java
index 3e0d66b..abdec92 100644
--- a/src/main/java/cuchaz/enigma/analysis/SourceIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/SourceIndex.java
@@ -11,17 +11,15 @@
 
 package cuchaz.enigma.analysis;
 
-import com.google.common.collect.HashMultimap;
-import com.google.common.collect.Lists;
-import com.google.common.collect.Maps;
-import com.google.common.collect.Multimap;
+import com.google.common.collect.*;
 import com.strobel.decompiler.languages.Region;
 import com.strobel.decompiler.languages.java.ast.AstNode;
 import com.strobel.decompiler.languages.java.ast.ConstructorDeclaration;
 import com.strobel.decompiler.languages.java.ast.Identifier;
 import com.strobel.decompiler.languages.java.ast.TypeDeclaration;
-import cuchaz.enigma.mapping.entry.Entry;
+import cuchaz.enigma.translation.representation.entry.Entry;
 
+import javax.annotation.Nullable;
 import java.util.*;
 import java.util.regex.Pattern;
 
@@ -29,9 +27,9 @@ public class SourceIndex {
         private static Pattern ANONYMOUS_INNER = Pattern.compile("\\$\\d+$");
 
         private String source;
-        private TreeMap<Token, EntryReference<Entry, Entry>> tokenToReference;
-        private Multimap<EntryReference<Entry, Entry>, Token> referenceToTokens;
-        private Map<Entry, Token> declarationToToken;
+        private TreeMap<Token, EntryReference<Entry<?>, Entry<?>>> tokenToReference;
+        private Multimap<EntryReference<Entry<?>, Entry<?>>, Token> referenceToTokens;
+        private Map<Entry<?>, Token> declarationToToken;
         private List<Integer> lineOffsets;
         private boolean ignoreBadTokens;
 
@@ -42,7 +40,7 @@ public class SourceIndex {
         public SourceIndex(String source, boolean ignoreBadTokens) {
                 this.source = source;
                 this.ignoreBadTokens = ignoreBadTokens;
-                this.tokenToReference = Maps.newTreeMap();
+                this.tokenToReference = new TreeMap<>();
                 this.referenceToTokens = HashMultimap.create();
                 this.declarationToToken = Maps.newHashMap();
                 calculateLineOffsets();
@@ -63,12 +61,12 @@ public class SourceIndex {
                 this.source = source;
                 calculateLineOffsets();
 
-                for (Entry entry : Lists.newArrayList(declarationToToken.keySet())) {
+                for (Entry<?> entry : Lists.newArrayList(declarationToToken.keySet())) {
                         Token token = declarationToToken.get(entry);
                         declarationToToken.put(entry, tokenMap.getOrDefault(token, token));
                 }
 
-                for (EntryReference<Entry, Entry> ref : referenceToTokens.keySet()) {
+                for (EntryReference<Entry<?>, Entry<?>> ref : referenceToTokens.keySet()) {
                         Collection<Token> oldTokens = referenceToTokens.get(ref);
                         List<Token> newTokens = new ArrayList<>(oldTokens.size());
 
@@ -79,7 +77,8 @@ public class SourceIndex {
                         referenceToTokens.replaceValues(ref, newTokens);
                 }
 
-                Map<Token, EntryReference<Entry, Entry>> tokenToReferenceCopy = Maps.newHashMap(tokenToReference);
+                TreeMap<Token, EntryReference<Entry<?>, Entry<?>>> tokenToReferenceCopy = new TreeMap<>(tokenToReference);
+
                 tokenToReference.clear();
                 for (Token token : tokenToReferenceCopy.keySet()) {
                         tokenToReference.put(tokenMap.getOrDefault(token, token), tokenToReferenceCopy.get(token));
@@ -112,9 +111,9 @@ public class SourceIndex {
                         return null;
                 }
 
-                if (node instanceof Identifier && name.indexOf('$') >=0 && node.getParent() instanceof ConstructorDeclaration && name.lastIndexOf('$') >= 0 && !ANONYMOUS_INNER.matcher(name).matches()){
+                if (node instanceof Identifier && name.indexOf('$') >= 0 && node.getParent() instanceof ConstructorDeclaration && name.lastIndexOf('$') >= 0 && !ANONYMOUS_INNER.matcher(name).matches()) {
                         TypeDeclaration type = node.getParent().getParent() instanceof TypeDeclaration ? (TypeDeclaration) node.getParent().getParent() : null;
-                        if (type != null){
+                        if (type != null) {
                                 name = type.getName();
                                 token.end = token.start + name.length();
                         }
@@ -133,19 +132,19 @@ public class SourceIndex {
                 return token;
         }
 
-        public void addReference(AstNode node, Entry deobfEntry, Entry deobfContext) {
+        public void addReference(AstNode node, Entry<?> deobfEntry, Entry<?> deobfContext) {
                 Token token = getToken(node);
                 if (token != null) {
-                        EntryReference<Entry, Entry> deobfReference = new EntryReference<>(deobfEntry, token.text, deobfContext);
+                        EntryReference<Entry<?>, Entry<?>> deobfReference = new EntryReference<>(deobfEntry, token.text, deobfContext);
                         this.tokenToReference.put(token, deobfReference);
                         this.referenceToTokens.put(deobfReference, token);
                 }
         }
 
-        public void addDeclaration(AstNode node, Entry deobfEntry) {
+        public void addDeclaration(AstNode node, Entry<?> deobfEntry) {
                 Token token = getToken(node);
                 if (token != null) {
-                        EntryReference<Entry, Entry> reference = new EntryReference<>(deobfEntry, token.text);
+                        EntryReference<Entry<?>, Entry<?>> reference = new EntryReference<>(deobfEntry, token.text);
                         this.tokenToReference.put(token, reference);
                         this.referenceToTokens.put(reference, token);
                         this.declarationToToken.put(deobfEntry, token);
@@ -160,22 +159,22 @@ public class SourceIndex {
                 return null;
         }
 
-        public Collection<Token> getReferenceTokens(EntryReference<Entry, Entry> deobfReference) {
+        public Collection<Token> getReferenceTokens(EntryReference<Entry<?>, Entry<?>> deobfReference) {
                 return this.referenceToTokens.get(deobfReference);
         }
 
-        public EntryReference<Entry, Entry> getDeobfReference(Token token) {
+        @Nullable
+        public EntryReference<Entry<?>, Entry<?>> getDeobfReference(Token token) {
                 if (token == null) {
                         return null;
                 }
                 return this.tokenToReference.get(token);
         }
 
-        public void replaceDeobfReference(Token token, EntryReference<Entry, Entry> newDeobfReference) {
-                EntryReference<Entry, Entry> oldDeobfReference = this.tokenToReference.get(token);
-                this.tokenToReference.put(token, newDeobfReference);
-                Collection<Token> tokens = this.referenceToTokens.get(oldDeobfReference);
-                this.referenceToTokens.removeAll(oldDeobfReference);
+        public void replaceDeobfReference(Token token, EntryReference<Entry<?>, Entry<?>> newDeobfReference) {
+                EntryReference<Entry<?>, Entry<?>> oldDeobfReferences = this.tokenToReference.replace(token, newDeobfReference);
+
+                Collection<Token> tokens = this.referenceToTokens.removeAll(oldDeobfReferences);
                 this.referenceToTokens.putAll(newDeobfReference, tokens);
         }
 
@@ -187,11 +186,11 @@ public class SourceIndex {
                 return this.declarationToToken.values();
         }
 
-        public Iterable<Entry> declarations() {
+        public Iterable<Entry<?>> declarations() {
                 return this.declarationToToken.keySet();
         }
 
-        public Token getDeclarationToken(Entry deobfEntry) {
+        public Token getDeclarationToken(Entry<?> deobfEntry) {
                 return this.declarationToToken.get(deobfEntry);
         }