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

5 files changed, 146 insertions, 39 deletions

diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/panels/EditorPanel.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/panels/EditorPanel.java
index 9bf5fa72..3409fc14 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/panels/EditorPanel.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/panels/EditorPanel.java
@@ -69,6 +69,7 @@ public class EditorPanel {
         private final Gui gui;
 
         private EntryReference<Entry<?>, Entry<?>> cursorReference;
+        private EntryReference<Entry<?>, Entry<?>> nextReference;
         private boolean mouseIsPressed = false;
         private boolean shouldNavigateOnClick;
 
@@ -332,13 +333,14 @@ public class EditorPanel {
                         } else {
                                 this.displayError(res.unwrapErr());
                         }
+                        this.nextReference = null;
                 });
         }
 
         public void displayError(ClassHandleError t) {
                 this.setDisplayMode(DisplayMode.ERRORED);
                 String str;
-                switch(t.type) {
+                switch (t.type) {
                         case DECOMPILE:
                                 str = "editor.decompile_error";
                                 break;
@@ -404,21 +406,13 @@ public class EditorPanel {
         }
 
         public void onCaretMove(int pos, boolean fromClick) {
+                if (this.settingSource) return;
                 if (this.controller.project == null) return;
 
                 EntryRemapper mapper = this.controller.project.getMapper();
                 Token token = getToken(pos);
 
-                if (this.settingSource) {
-                        EntryReference<Entry<?>, Entry<?>> ref = getCursorReference();
-                        EntryReference<Entry<?>, Entry<?>> refAtCursor = getReference(token);
-                        if (this.editor.getDocument().getLength() != 0 && ref != null && !ref.equals(refAtCursor)) {
-                                showReference0(ref);
-                        }
-                        return;
-                } else {
-                        setCursorReference(getReference(token));
-                }
+                setCursorReference(getReference(token));
 
                 Entry<?> referenceEntry = this.cursorReference != null ? this.cursorReference.entry : null;
 
@@ -484,17 +478,49 @@ public class EditorPanel {
                 if (source == null) return;
                 try {
                         this.settingSource = true;
+
+                        int newCaretPos = 0;
+                        if (this.source != null && this.source.getEntry().equals(source.getEntry())) {
+                                int caretPos = this.editor.getCaretPosition();
+
+                                if (this.source.getTokenStore().isCompatible(source.getTokenStore())) {
+                                        newCaretPos = this.source.getTokenStore().mapPosition(source.getTokenStore(), caretPos);
+                                } else {
+                                        // if the class is the same but the token stores aren't
+                                        // compatible, then the user probably switched decompilers
+
+                                        // check if there's a selected reference we can navigate to,
+                                        // but only if there's none already queued up for being selected
+                                        if (this.getCursorReference() != null && this.nextReference == null) {
+                                                this.nextReference = this.getCursorReference();
+                                        }
+
+                                        // otherwise fall back to just using the same average
+                                        // position in the file
+                                        float scale = (float) source.toString().length() / this.source.toString().length();
+                                        newCaretPos = (int) (caretPos * scale);
+                                }
+                        }
+
                         this.source = source;
                         this.editor.getHighlighter().removeAllHighlights();
                         this.editor.setText(source.toString());
+                        if (this.source != null) {
+                                this.editor.setCaretPosition(newCaretPos);
+                        }
                         setHighlightedTokens(source.getHighlightedTokens());
+                        setCursorReference(getReference(getToken(this.editor.getCaretPosition())));
                 } finally {
                         this.settingSource = false;
                 }
-                showReference0(getCursorReference());
+
+                if (this.nextReference != null) {
+                        this.showReference0(this.nextReference);
+                        this.nextReference = null;
+                }
         }
 
-        public void setHighlightedTokens(Map<RenamableTokenType, Collection<Token>> tokens) {
+        public void setHighlightedTokens(Map<RenamableTokenType, ? extends Collection<Token>> tokens) {
                 // remove any old highlighters
                 this.editor.getHighlighter().removeAllHighlights();
 
@@ -526,8 +552,11 @@ public class EditorPanel {
         }
 
         public void showReference(EntryReference<Entry<?>, Entry<?>> reference) {
-                setCursorReference(reference);
-                showReference0(reference);
+                if (this.mode == DisplayMode.SUCCESS) {
+                        showReference0(reference);
+                } else if (this.mode != DisplayMode.ERRORED) {
+                        this.nextReference = reference;
+                }
         }
 
         /**
diff --git a/enigma/src/main/java/cuchaz/enigma/classhandle/ClassHandleProvider.java b/enigma/src/main/java/cuchaz/enigma/classhandle/ClassHandleProvider.java
index bc38e61b..22722075 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 e4a71150..9ac611f5 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 fad733fa..7d729abb 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 00000000..f32d918c
--- /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;
+        }
+
+}