From 0ea7826d3c22902a3cfd43c7e2a4079a208bb5c0 Mon Sep 17 00:00:00 2001
From: peter <peter@jetbrains.com>
Date: Fri, 30 Mar 2012 18:52:12 +0200
Subject: [PATCH] let stub reference stub tree by default to prevent it being
 collected (e.g. ClsFileImpl)

---
 .../src/com/intellij/psi/impl/source/PsiFileImpl.java    | 1 -
 .../core-impl/src/com/intellij/psi/stubs/StubTree.java   | 9 +++++++++
 .../src/com/intellij/psi/stubs/StubUpdatingIndex.java    | 2 +-
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/platform/core-impl/src/com/intellij/psi/impl/source/PsiFileImpl.java b/platform/core-impl/src/com/intellij/psi/impl/source/PsiFileImpl.java
index 347786657d07..ef0e78c11bf1 100644
--- a/platform/core-impl/src/com/intellij/psi/impl/source/PsiFileImpl.java
+++ b/platform/core-impl/src/com/intellij/psi/impl/source/PsiFileImpl.java
@@ -684,7 +684,6 @@ public abstract class PsiFileImpl extends ElementBase implements PsiFileEx, PsiF
       myStub = new SoftReference<StubTree>(stubHolder);
       StubBase<PsiFile> base = (StubBase)stubHolder.getRoot();
       base.setPsi(this);
-      base.putUserData(STUB_TREE_IN_PARSED_TREE, stubHolder); // This will prevent soft reference myStub to be collected before all of the stubs are collected.
       return stubHolder;
     }
   }
diff --git a/platform/core-impl/src/com/intellij/psi/stubs/StubTree.java b/platform/core-impl/src/com/intellij/psi/stubs/StubTree.java
index 1b5ea6483a35..d1021e2be84c 100644
--- a/platform/core-impl/src/com/intellij/psi/stubs/StubTree.java
+++ b/platform/core-impl/src/com/intellij/psi/stubs/StubTree.java
@@ -19,6 +19,7 @@
  */
 package com.intellij.psi.stubs;
 
+import com.intellij.openapi.util.Key;
 import gnu.trove.THashMap;
 import gnu.trove.TIntArrayList;
 import org.jetbrains.annotations.NotNull;
@@ -29,12 +30,20 @@ import java.util.List;
 import java.util.Map;
 
 public class StubTree {
+  private static final Key<StubTree> STUB_TO_TREE_REFERENCE = Key.create("stub to tree reference");
   private final PsiFileStub myRoot;
   private final List<StubElement<?>> myPlainList = new ArrayList<StubElement<?>>();
 
   public StubTree(@NotNull final PsiFileStub root) {
+    this(root, true);
+  }
+
+  public StubTree(@NotNull final PsiFileStub root, final boolean withBackReference) {
     myRoot = root;
     enumerateStubs(root, myPlainList);
+    if (withBackReference) {
+      myRoot.putUserData(STUB_TO_TREE_REFERENCE, this); // This will prevent soft references to stub tree to be collected before all of the stubs are collected.
+    }
   }
 
   private static void enumerateStubs(final StubElement<?> root, final List<StubElement<?>> result) {
diff --git a/platform/lang-impl/src/com/intellij/psi/stubs/StubUpdatingIndex.java b/platform/lang-impl/src/com/intellij/psi/stubs/StubUpdatingIndex.java
index b2a2c5334c80..6812876cb35c 100644
--- a/platform/lang-impl/src/com/intellij/psi/stubs/StubUpdatingIndex.java
+++ b/platform/lang-impl/src/com/intellij/psi/stubs/StubUpdatingIndex.java
@@ -306,7 +306,7 @@ public class StubUpdatingIndex extends CustomImplementationFileBasedIndexExtensi
       final Map<StubIndexKey, Map<Object, TIntArrayList>> stubTree;
       if (!data.isEmpty()) {
         final SerializedStubTree stub = data.values().iterator().next();
-        stubTree = new StubTree((PsiFileStub)stub.getStub()).indexStubTree();
+        stubTree = new StubTree((PsiFileStub)stub.getStub(), false).indexStubTree();
       }
       else {
         stubTree = Collections.emptyMap();