From 8a146b1014b619f348cd0972f37315db4df0d2e9 Mon Sep 17 00:00:00 2001
From: peter <peter@jetbrains.com>
Date: Fri, 20 Jan 2017 09:20:11 +0100
Subject: [PATCH] allow psi.getStub to return something after AST is GC-ed

---
 .../com/intellij/psi/StubAstSwitchTest.groovy | 13 +++++++++
 .../extapi/psi/StubBasedPsiElementBase.java   |  2 +-
 .../psi/impl/source/SubstrateRef.java         |  6 ++---
 .../psi/impl/source/tree/AstPath.java         | 27 +++++++++++++++++++
 4 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/java/java-tests/testSrc/com/intellij/psi/StubAstSwitchTest.groovy b/java/java-tests/testSrc/com/intellij/psi/StubAstSwitchTest.groovy
index f499391ce805..1c100cda9815 100644
--- a/java/java-tests/testSrc/com/intellij/psi/StubAstSwitchTest.groovy
+++ b/java/java-tests/testSrc/com/intellij/psi/StubAstSwitchTest.groovy
@@ -19,7 +19,9 @@ import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.command.WriteCommandAction
 import com.intellij.openapi.fileEditor.FileDocumentManager
 import com.intellij.openapi.vfs.VfsUtil
+import com.intellij.psi.impl.source.PsiClassImpl
 import com.intellij.psi.impl.source.PsiFileImpl
+import com.intellij.psi.impl.source.PsiJavaFileImpl
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.search.searches.DirectClassInheritorsSearch
 import com.intellij.psi.search.searches.OverridingMethodsSearch
@@ -235,4 +237,15 @@ class B {
     assert !file.node.parsed
   }
 
+  void "test load stub from non-file PSI after AST is unloaded"() {
+    PsiJavaFileImpl file = (PsiJavaFileImpl)myFixture.addFileToProject("a.java", "class A<T>{}")
+    def cls = file.classes[0]
+    assert cls.nameIdentifier
+
+    GCUtil.tryGcSoftlyReachableObjects()
+    assert !file.treeElement
+
+    assert ((PsiClassImpl) cls).stub
+  }
+
 }
diff --git a/platform/core-impl/src/com/intellij/extapi/psi/StubBasedPsiElementBase.java b/platform/core-impl/src/com/intellij/extapi/psi/StubBasedPsiElementBase.java
index fd25aa2ba4e6..2e672c04023f 100644
--- a/platform/core-impl/src/com/intellij/extapi/psi/StubBasedPsiElementBase.java
+++ b/platform/core-impl/src/com/intellij/extapi/psi/StubBasedPsiElementBase.java
@@ -371,7 +371,7 @@ public class StubBasedPsiElementBase<T extends StubElement> extends ASTDelegateP
   public T getStub() {
     ProgressIndicatorProvider.checkCanceled(); // Hope, this is called often
     //noinspection unchecked
-    return (T)mySubstrateRef.getStub();
+    return (T)mySubstrateRef.getStub(myStubIndex);
   }
 
   /**
diff --git a/platform/core-impl/src/com/intellij/psi/impl/source/SubstrateRef.java b/platform/core-impl/src/com/intellij/psi/impl/source/SubstrateRef.java
index 154c474288f4..45491ba69508 100644
--- a/platform/core-impl/src/com/intellij/psi/impl/source/SubstrateRef.java
+++ b/platform/core-impl/src/com/intellij/psi/impl/source/SubstrateRef.java
@@ -42,13 +42,13 @@ public abstract class SubstrateRef {
   public abstract ASTNode getNode();
 
   @Nullable
-  public Stub getStub() {
+  public Stub getStub(int stubIndex) {
     return null;
   }
 
   @Nullable
   public Stub getGreenStub(int index) {
-    return getStub();
+    return getStub(index);
   }
 
   public abstract boolean isValid();
@@ -119,7 +119,7 @@ public abstract class SubstrateRef {
 
     @NotNull
     @Override
-    public Stub getStub() {
+    public Stub getStub(int stubIndex) {
       return myStub;
     }
 
diff --git a/platform/core-impl/src/com/intellij/psi/impl/source/tree/AstPath.java b/platform/core-impl/src/com/intellij/psi/impl/source/tree/AstPath.java
index 86bdee26bbc2..96ddd6ae10fa 100644
--- a/platform/core-impl/src/com/intellij/psi/impl/source/tree/AstPath.java
+++ b/platform/core-impl/src/com/intellij/psi/impl/source/tree/AstPath.java
@@ -57,6 +57,18 @@ public abstract class AstPath extends SubstrateRef {
 
   protected abstract int getDepth();
 
+  @Nullable
+  @Override
+  public Stub getStub(int stubIndex) {
+    if (stubIndex < 0) return null;
+
+    StubTree stubTree = getFileStubTree();
+    return stubTree == null ? null : stubTree.getPlainList().get(stubIndex);
+  }
+
+  @Nullable
+  protected abstract StubTree getFileStubTree();
+
   @Nullable
   @Override
   public Stub getGreenStub(int stubIndex) {
@@ -165,6 +177,11 @@ public abstract class AstPath extends SubstrateRef {
       return 1 + myParent.getDepth();
     }
 
+    @Override
+    protected StubTree getFileStubTree() {
+      return myParent.getFileStubTree();
+    }
+
     @Override
     public boolean equals(Object o) {
       if (this == o) return true;
@@ -204,6 +221,11 @@ public abstract class AstPath extends SubstrateRef {
       return node;
     }
 
+    @Override
+    protected StubTree getFileStubTree() {
+      return SoftReference.dereference(myNode) == null ? myFile.getStubTree() : null;
+    }
+
     @NotNull
     @Override
     public PsiFileImpl getContainingFile() {
@@ -249,6 +271,11 @@ public abstract class AstPath extends SubstrateRef {
     protected int getDepth() {
       return 0;
     }
+
+    @Override
+    protected StubTree getFileStubTree() {
+      return myFile.getStubTree();
+    }
   }
 
 }