fix: Renaming works again

src/main/java/com/falsepattern/zigbrains/lsp/contributors/rename/LSPInplaceRenamer.java

@@ -49,7 +49,7 @@ public class LSPInplaceRenamer extends MemberInplaceRenamer {
         EditorEventManager eventManager = EditorEventManagerBase.forEditor(editor);
         if (eventManager != null) {
             Pair<List<PsiElement>, List<VirtualFile>> results = eventManager
-                    .references(editor.getCaretModel().getCurrentCaret().getOffset(), true, false);
+                    .referencesForRename(editor.getCaretModel().getCurrentCaret().getOffset(), true, false);
             List<PsiElement> references = results.getFirst();
             List<VirtualFile> toClose = results.getSecond();
             LSPRenameProcessor.addEditors(toClose);

src/main/java/com/falsepattern/zigbrains/lsp/contributors/rename/LSPRenameHandler.java

@@ -44,6 +44,7 @@ import com.intellij.refactoring.rename.inplace.MemberInplaceRenamer;
 import org.jetbrains.annotations.NotNull;
 import com.falsepattern.zigbrains.lsp.contributors.psi.LSPPsiElement;
 
+import java.lang.ref.WeakReference;
 import java.util.List;
 
 import static com.intellij.openapi.command.impl.StartMarkAction.START_MARK_ACTION_KEY;
@@ -161,9 +162,11 @@ public class LSPRenameHandler implements RenameHandler {
                 .filter(e -> e.getTextRange().getStartOffset() <= offset && offset <= e.getTextRange().getEndOffset())
                 .findAny().orElse(null);
         if (curElement != null) {
-            return new LSPPsiElement(curElement.getText(), curElement.getProject(),
+            var newElement = new LSPPsiElement(curElement.getText(), curElement.getProject(),
-                    curElement.getTextRange().getStartOffset(), curElement.getTextRange().getEndOffset(),
+                                               curElement.getTextRange().getStartOffset(), curElement.getTextRange().getEndOffset(),
-                    curElement.getContainingFile());
+                                               curElement.getContainingFile());
+            eventManager.renameCache = new Pair<>(new WeakReference<>(newElement), refResponse);
+            return newElement;
         }
         return null;
     }

src/main/java/com/falsepattern/zigbrains/lsp/contributors/rename/LSPRenameProcessor.java

@@ -73,7 +73,7 @@ public class LSPRenameProcessor extends RenamePsiElementProcessor {
             EditorEventManager
                     manager = EditorEventManagerBase.forEditor(FileUtils.editorFromPsiFile(element.getContainingFile()));
             if (manager != null) {
-                Pair<List<PsiElement>, List<VirtualFile>> refs = manager.references(element.getTextOffset(), true, false);
+                Pair<List<PsiElement>, List<VirtualFile>> refs = manager.referencesForRename(element, true, false);
                 if (refs.getFirst() != null && refs.getSecond() != null) {
                     addEditors(refs.getSecond());
                     return refs.getFirst().stream().map(PsiElement::getReference).filter(Objects::nonNull).collect(Collectors.toList());

src/main/java/com/falsepattern/zigbrains/lsp/editor/EditorEventManager.java

@@ -76,6 +76,7 @@ import com.falsepattern.zigbrains.lsp.utils.GUIUtils;
 
 import javax.swing.*;
 import java.awt.*;
+import java.lang.ref.WeakReference;
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.util.List;
@@ -364,6 +365,34 @@ public class EditorEventManager {
         return references(offset, false, false);
     }
 
+    public Pair<WeakReference<PsiElement>, Pair<List<PsiElement>, List<VirtualFile>>> renameCache = null;
+
+    public Pair<List<PsiElement>, List<VirtualFile>> referencesForRename(PsiElement originalElement, boolean getOriginalElement, boolean close) {
+        if (renameCache != null) {
+            var cached = renameCache.getFirst().get();
+            if (cached == originalElement) {
+                return renameCache.getSecond();
+            }
+        }
+        var newValue = references(originalElement.getTextOffset(), getOriginalElement, close);
+        renameCache = new Pair<>(new WeakReference<>(originalElement), newValue);
+        return newValue;
+    }
+
+    public Pair<List<PsiElement>, List<VirtualFile>> referencesForRename(int offset, boolean getOriginalElement, boolean close) {
+        if (renameCache != null) {
+            var cached = renameCache.getFirst().get();
+            if (cached != null) {
+                var textRange = cached.getTextRange();
+                if (textRange.contains(offset) || textRange.contains(offset-1)) {
+                    return renameCache.getSecond();
+                }
+            }
+        }
+        renameCache = null;
+        return references(offset, getOriginalElement, close);
+    }
+
     /**
      * Returns the references given the position of the word to search for
      * Must be called from main thread
@@ -372,6 +401,7 @@ public class EditorEventManager {
      * @return An array of PsiElement
      */
     public Pair<List<PsiElement>, List<VirtualFile>> references(int offset, boolean getOriginalElement, boolean close) {
+        renameCache = null;
         Position lspPos = DocumentUtils.offsetToLSPPos(editor, offset);
         TextDocumentIdentifier textDocumentIdentifier = new TextDocumentIdentifier(FileUtils.editorToURIString(editor));
         ReferenceParams params = new ReferenceParams(textDocumentIdentifier, lspPos, new ReferenceContext(getOriginalElement));