Fix losing current cursor position when renaming entries (#297)

* Fix losing current cursor position when renaming entries

* Set nextReference to null after applying it

* Extract token map into wrapper type

4 files changed, 102 insertions, 24 deletions

diff --git a/enigma/src/main/java/cuchaz/enigma/classhandle/ClassHandleProvider.java b/enigma/src/main/java/cuchaz/enigma/classhandle/ClassHandleProvider.java
index bc38e61..2272207 100644
--- a/enigma/src/main/java/cuchaz/enigma/classhandle/ClassHandleProvider.java
+++ b/enigma/src/main/java/cuchaz/enigma/classhandle/ClassHandleProvider.java
@@ -11,10 +11,9 @@ import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 import javax.annotation.Nullable;
 
+import cuchaz.enigma.EnigmaProject;
 import cuchaz.enigma.classprovider.CachingClassProvider;
 import cuchaz.enigma.classprovider.ObfuscationFixClassProvider;
-
-import cuchaz.enigma.EnigmaProject;
 import cuchaz.enigma.events.ClassHandleListener;
 import cuchaz.enigma.events.ClassHandleListener.InvalidationType;
 import cuchaz.enigma.source.*;
@@ -277,8 +276,8 @@ public final class ClassHandleProvider {
                                 if (res == null || mappedVersion.get() != v) return;
                                 res = res.andThen(source -> {
                                         try {
-                                                source.remapSource(p.project, p.project.getMapper().getDeobfuscator());
-                                                return Result.ok(source);
+                                                DecompiledClassSource remappedSource = source.remapSource(p.project, p.project.getMapper().getDeobfuscator());
+                                                return Result.ok(remappedSource);
                                         } catch (Throwable e) {
                                                 return Result.err(ClassHandleError.remap(e));
                                         }
diff --git a/enigma/src/main/java/cuchaz/enigma/source/DecompiledClassSource.java b/enigma/src/main/java/cuchaz/enigma/source/DecompiledClassSource.java
index e4a7115..9ac611f 100644
--- a/enigma/src/main/java/cuchaz/enigma/source/DecompiledClassSource.java
+++ b/enigma/src/main/java/cuchaz/enigma/source/DecompiledClassSource.java
@@ -22,30 +22,35 @@ public class DecompiledClassSource {
         private final ClassEntry classEntry;
 
         private final SourceIndex obfuscatedIndex;
-        private SourceIndex remappedIndex;
+        private final SourceIndex remappedIndex;
 
-        private final Map<RenamableTokenType, Collection<Token>> highlightedTokens = new EnumMap<>(RenamableTokenType.class);
+        private final TokenStore highlightedTokens;
 
-        public DecompiledClassSource(ClassEntry classEntry, SourceIndex index) {
+        private DecompiledClassSource(ClassEntry classEntry, SourceIndex obfuscatedIndex, SourceIndex remappedIndex, TokenStore highlightedTokens) {
                 this.classEntry = classEntry;
-                this.obfuscatedIndex = index;
-                this.remappedIndex = index;
+                this.obfuscatedIndex = obfuscatedIndex;
+                this.remappedIndex = remappedIndex;
+                this.highlightedTokens = highlightedTokens;
+        }
+
+        public DecompiledClassSource(ClassEntry classEntry, SourceIndex index) {
+                this(classEntry, index, index, TokenStore.empty());
         }
 
         public static DecompiledClassSource text(ClassEntry classEntry, String text) {
                 return new DecompiledClassSource(classEntry, new SourceIndex(text));
         }
 
-        public void remapSource(EnigmaProject project, Translator translator) {
-                highlightedTokens.clear();
-
+        public DecompiledClassSource remapSource(EnigmaProject project, Translator translator) {
                 SourceRemapper remapper = new SourceRemapper(obfuscatedIndex.getSource(), obfuscatedIndex.referenceTokens());
 
-                SourceRemapper.Result remapResult = remapper.remap((token, movedToken) -> remapToken(project, token, movedToken, translator));
-                remappedIndex = obfuscatedIndex.remapTo(remapResult);
+                TokenStore tokenStore = TokenStore.create(this.obfuscatedIndex);
+                SourceRemapper.Result remapResult = remapper.remap((token, movedToken) -> remapToken(tokenStore, project, token, movedToken, translator));
+                SourceIndex remappedIndex = obfuscatedIndex.remapTo(remapResult);
+                return new DecompiledClassSource(this.classEntry, this.obfuscatedIndex, remappedIndex, tokenStore);
         }
 
-        private String remapToken(EnigmaProject project, Token token, Token movedToken, Translator translator) {
+        private String remapToken(TokenStore target, EnigmaProject project, Token token, Token movedToken, Translator translator) {
                 EntryReference<Entry<?>, Entry<?>> reference = obfuscatedIndex.getReference(token);
 
                 Entry<?> entry = reference.getNameableEntry();
@@ -53,16 +58,16 @@ public class DecompiledClassSource {
 
                 if (project.isRenamable(reference)) {
                         if (!translatedEntry.isObfuscated()) {
-                                highlightToken(movedToken, translatedEntry.getType());
+                                target.add(translatedEntry.getType(), movedToken);
                                 return translatedEntry.getValue().getSourceRemapName();
                         } else {
                                 Optional<String> proposedName = proposeName(project, entry);
                                 if (proposedName.isPresent()) {
-                                        highlightToken(movedToken, RenamableTokenType.PROPOSED);
+                                        target.add(RenamableTokenType.PROPOSED, movedToken);
                                         return proposedName.get();
                                 }
 
-                                highlightToken(movedToken, RenamableTokenType.OBFUSCATED);
+                                target.add(RenamableTokenType.OBFUSCATED, movedToken);
                         }
                 }
 
@@ -113,12 +118,12 @@ public class DecompiledClassSource {
                 return remappedIndex;
         }
 
-        public Map<RenamableTokenType, Collection<Token>> getHighlightedTokens() {
-                return highlightedTokens;
+        public TokenStore getTokenStore() {
+                return this.highlightedTokens;
         }
 
-        private void highlightToken(Token token, RenamableTokenType highlightType) {
-                highlightedTokens.computeIfAbsent(highlightType, t -> new ArrayList<>()).add(token);
+        public Map<RenamableTokenType, ? extends Collection<Token>> getHighlightedTokens() {
+                return this.highlightedTokens.getByType();
         }
 
         public int getObfuscatedOffset(int deobfOffset) {
diff --git a/enigma/src/main/java/cuchaz/enigma/source/Token.java b/enigma/src/main/java/cuchaz/enigma/source/Token.java
index fad733f..7d729ab 100644
--- a/enigma/src/main/java/cuchaz/enigma/source/Token.java
+++ b/enigma/src/main/java/cuchaz/enigma/source/Token.java
@@ -23,9 +23,12 @@ public class Token implements Comparable<Token> {
                 this.text = text;
         }
 
+        public int length() {
+                return this.end - this.start;
+        }
+
         public int getRenameOffset(String to) {
-                int length = this.end - this.start;
-                return to.length() - length;
+                return to.length() - this.length();
         }
 
         public void rename(StringBuffer source, String to) {
diff --git a/enigma/src/main/java/cuchaz/enigma/source/TokenStore.java b/enigma/src/main/java/cuchaz/enigma/source/TokenStore.java
new file mode 100644
index 0000000..f32d918
--- /dev/null
+++ b/enigma/src/main/java/cuchaz/enigma/source/TokenStore.java
@@ -0,0 +1,71 @@
+package cuchaz.enigma.source;
+
+import java.util.*;
+
+public final class TokenStore {
+
+        private static final TokenStore EMPTY = new TokenStore(Collections.emptyNavigableSet(), Collections.emptyMap(), null);
+
+        private final NavigableSet<Token> tokens;
+        private final Map<RenamableTokenType, NavigableSet<Token>> byType;
+        private final String obfSource;
+
+        private TokenStore(NavigableSet<Token> tokens, Map<RenamableTokenType, NavigableSet<Token>> byType, String obfSource) {
+                this.tokens = tokens;
+                this.byType = byType;
+                this.obfSource = obfSource;
+        }
+
+        public static TokenStore create(SourceIndex obfuscatedIndex) {
+                EnumMap<RenamableTokenType, NavigableSet<Token>> map = new EnumMap<>(RenamableTokenType.class);
+                for (RenamableTokenType value : RenamableTokenType.values()) {
+                        map.put(value, new TreeSet<>(Comparator.comparing(t -> t.start)));
+                }
+                return new TokenStore(new TreeSet<>(Comparator.comparing(t -> t.start)), Collections.unmodifiableMap(map), obfuscatedIndex.getSource());
+        }
+
+        public static TokenStore empty() {
+                return TokenStore.EMPTY;
+        }
+
+        public void add(RenamableTokenType type, Token token) {
+                this.tokens.add(token);
+                this.byType.get(type).add(token);
+        }
+
+        public boolean isCompatible(TokenStore other) {
+                return this.obfSource != null && other.obfSource != null &&
+                                this.obfSource.equals(other.obfSource) &&
+                                this.tokens.size() == other.tokens.size();
+        }
+
+        public int mapPosition(TokenStore to, int position) {
+                if (!this.isCompatible(to)) return 0;
+
+                int newPos = position;
+                Iterator<Token> thisIter = this.tokens.iterator();
+                Iterator<Token> toIter = to.tokens.iterator();
+                while (thisIter.hasNext()) {
+                        Token token = thisIter.next();
+                        Token newToken = toIter.next();
+
+                        if (position < token.start) break;
+
+                        // if we're inside the token and the text changed,
+                        // snap the cursor to the beginning
+                        if (!token.text.equals(newToken.text) && position < token.end) {
+                                newPos = newToken.start;
+                                break;
+                        }
+
+                        newPos += newToken.length() - token.length();
+                }
+
+                return newPos;
+        }
+
+        public Map<RenamableTokenType, NavigableSet<Token>> getByType() {
+                return byType;
+        }
+
+}