IDEA-144655 Smart PSI pointers should retain old offsets after inserting text after them

java/java-tests/testSrc/com/intellij/psi/impl/smartPointers/SmartPsiElementPointersTest.java

@@ -24,6 +24,7 @@ import com.intellij.lang.FileASTNode;
 import com.intellij.openapi.command.WriteCommandAction;
 import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.editor.EditorFactory;
+import com.intellij.openapi.editor.EditorModificationUtil;
 import com.intellij.openapi.editor.event.DocumentEvent;
 import com.intellij.openapi.editor.event.DocumentListener;
 import com.intellij.openapi.editor.event.EditorEventMulticaster;
@@ -754,6 +755,24 @@ public class SmartPsiElementPointersTest extends CodeInsightTestCase {
     PsiDocumentManager.getInstance(myProject).commitAllDocuments();
     assertEquals(TextRange.create(2, 3), pointer.getRange());
   }
+  
+  public void testUpdateAfterInsertingIdenticalText() {
+    PsiJavaFile file = (PsiJavaFile)configureByText(StdFileTypes.JAVA, "class Foo {\n" +
+                                                                       "    void m() {\n" +
+                                                                       "    }\n" +
+                                                                       "<caret>}\n");
+    PsiMethod method = file.getClasses()[0].getMethods()[0];
+    TextRange originalRange = method.getTextRange();
+    SmartPsiElementPointer pointer = SmartPointerManager.getInstance(myProject).createSmartPsiElementPointer(method);
+
+    EditorModificationUtil.insertStringAtCaret(myEditor, "    void m() {\n" +
+                                                         "    }\n");
+    PsiDocumentManager.getInstance(myProject).commitDocument(myEditor.getDocument());
+    PsiElement element = pointer.getElement();
+    assertNotNull(element);
+    TextRange newRange = element.getTextRange();
+    assertEquals(originalRange, newRange);
+  }
 
   public void testAnchorInfoSurvivesPsiChange() {
     PsiJavaFile file = (PsiJavaFile)configureByText(JavaFileType.INSTANCE, "class C1{}\nclass C2 {}");

platform/core-impl/src/com/intellij/psi/impl/smartPointers/MarkerCache.java

@@ -23,6 +23,7 @@ import com.intellij.openapi.editor.impl.event.RetargetRangeMarkers;
 import com.intellij.openapi.util.ProperTextRange;
 import com.intellij.openapi.util.Trinity;
 import com.intellij.openapi.vfs.VirtualFile;
+import com.intellij.psi.PsiElement;
 import com.intellij.util.NullableFunction;
 import com.intellij.util.containers.ContainerUtil;
 import gnu.trove.TLongObjectHashMap;
@@ -113,17 +114,40 @@ class MarkerCache {
   synchronized void updateMarkers(@NotNull FrozenDocument frozen, @NotNull List<DocumentEvent> events) {
     TLongObjectHashMap<ManualRangeMarker> updated = getUpdatedMarkers(frozen, events);
 
-    for (SelfElementInfo info : getInfos()) {
-      long key = info.markerCacheKey();
-      if (key != 0) {
-        ManualRangeMarker newRange = updated.get(key);
-        info.setRange(newRange == null ? null : newRange.getRange());
+    for (SmartPsiElementPointerImpl pointer : myPointers.getAlivePointers()) {
+      SmartPointerElementInfo info = pointer.getElementInfo();
+      if (info instanceof SelfElementInfo) {
+        long key = ((SelfElementInfo)info).markerCacheKey();
+        if (key != 0) {
+          ManualRangeMarker newRangeMarker = updated.get(key);
+          ProperTextRange newRange = newRangeMarker == null ? null : newRangeMarker.getRange();
+          ((SelfElementInfo)info).setRange(newRange);
+
+          if (newRange != null && !(pointer instanceof SmartPsiFileRangePointerImpl)) {
+            updatePointerTarget(pointer, newRange);
+          }
+        }
+
       }
     }
 
     myUpdatedRanges = null;
   }
 
+  // after reparse and its complex tree diff, the element might have "moved" to other range
+  // but if an element of the same type can still be found at the old range, let's point there
+  private static <E extends PsiElement> void updatePointerTarget(@NotNull SmartPsiElementPointerImpl<E> pointer, @NotNull ProperTextRange newRange) {
+    E cachedElement = pointer.getCachedElement();
+    if (cachedElement == null || cachedElement.isValid() && newRange.equals(cachedElement.getTextRange())) {
+      return;
+    }
+
+    E newTarget = pointer.doRestoreElement();
+    if (newTarget != null) {
+      pointer.cacheElement(newTarget);
+    }
+  }
+
   @NotNull
   private List<SelfElementInfo> getInfos() {
     return ContainerUtil.findAll(ContainerUtil.map(myPointers.getAlivePointers(), new NullableFunction<SmartPsiElementPointerImpl, SmartPointerElementInfo>() {

platform/core-impl/src/com/intellij/psi/impl/smartPointers/SmartPointerManagerImpl.java

@@ -134,8 +134,8 @@ public class SmartPointerManagerImpl extends SmartPointerManager {
     Reference<SmartPointerEx> data = element.getUserData(CACHED_SMART_POINTER_KEY);
     SmartPointerEx cachedPointer = SoftReference.dereference(data);
     if (cachedPointer != null) {
-      PsiElement cachedElement = cachedPointer.getCachedElement();
-      if (cachedElement != null && cachedElement != element) {
+      PsiElement cachedElement = cachedPointer.getElement();
+      if (cachedElement == null || cachedElement != element) {
         return null;
       }
     }

platform/core-impl/src/com/intellij/psi/impl/smartPointers/SmartPsiElementPointerImpl.java

@@ -80,23 +80,24 @@ class SmartPsiElementPointerImpl<E extends PsiElement> implements SmartPointerEx
   @Nullable
   public E getElement() {
     E element = getCachedElement();
-    if (element != null && !element.isValid()) {
-      element = null;
-    }
-    if (element == null) {
-      //noinspection unchecked
-      element = (E)myElementInfo.restoreElement();
-      if (element != null && (!element.getClass().equals(myElementClass) || !element.isValid())) {
-        element = null;
-      }
-
+    if (element == null || !element.isValid()) {
+      element = doRestoreElement();
       cacheElement(element);
     }
-
     return element;
   }
 
-  private void cacheElement(E element) {
+  @Nullable
+  E doRestoreElement() {
+    //noinspection unchecked
+    E element = (E)myElementInfo.restoreElement();
+    if (element != null && (!element.getClass().equals(myElementClass) || !element.isValid())) {
+      return null;
+    }
+    return element;
+  }
+
+  void cacheElement(@Nullable E element) {
     myElement = element == null ? null : 
                 ((PsiManagerEx)PsiManager.getInstance(getProject())).isBatchFilesProcessingMode() ? new WeakReference<E>(element) : 
                 new SoftReference<E>(element);