introduce StubbedSpine to relieve clients from getStubTree/calcStubTree dualism

and save memory on not instantiating stubs when AST is enough
2026-09-27 10:03:11 +07:00 · 2017-09-26 14:10:18 +02:00
parent 45650b36f7
commit cf10982ace
14 changed files with 244 additions and 140 deletions
@@ -23,14 +23,13 @@ import com.intellij.openapi.module.Module;
 import com.intellij.openapi.roots.ProjectFileIndex;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.vfs.VirtualFile;
+import com.intellij.psi.PsiAnonymousClass;
 import com.intellij.psi.PsiClass;
 import com.intellij.psi.PsiFunctionalExpression;
-import com.intellij.psi.impl.java.stubs.PsiClassStub;
-import com.intellij.psi.impl.source.PsiFileImpl;
+import com.intellij.psi.impl.source.Constants;
 import com.intellij.psi.impl.source.PsiFileWithStubSupport;
+import com.intellij.psi.impl.source.StubbedSpine;
 import com.intellij.psi.impl.source.tree.JavaElementType;
-import com.intellij.psi.stubs.StubElement;
-import com.intellij.psi.stubs.StubTree;
 import com.intellij.psi.tree.TokenSet;
 import com.intellij.util.SmartList;
 import com.intellij.util.containers.ContainerUtil;
@@ -48,18 +47,14 @@ public class JavaCompilerElementRetriever {
  @NotNull
  static PsiFunctionalExpression[] retrieveFunExpressionsByIndices(@NotNull TIntHashSet indices,
                                                                   @NotNull PsiFileWithStubSupport psiFile) {
-    StubTree tree = psiFile.getStubTree();
-    boolean foreign = tree == null;
-    if (foreign) {
-      tree = ((PsiFileImpl)psiFile).calcStubTree();
-    }
+    StubbedSpine spine = psiFile.getStubbedSpine();

    PsiFunctionalExpression[] result = new PsiFunctionalExpression[indices.size()];
    int resIdx = 0;
    int funExprIdx = 0;
-    for (StubElement<?> element : tree.getPlainList()) {
-      if (FUN_EXPR.contains(element.getStubType()) && indices.contains(funExprIdx++)) {
-        result[resIdx++] = (PsiFunctionalExpression)element.getPsi();
+    for (int i = 0; i < spine.getStubCount(); i++) {
+      if (FUN_EXPR.contains(spine.getStubType(i)) && indices.contains(funExprIdx++)) {
+        result[resIdx++] = (PsiFunctionalExpression)spine.getStubPsi(i);
      }
    }

@@ -88,7 +83,7 @@ public class JavaCompilerElementRetriever {
  }

  private interface InternalNameMatcher {
-    boolean matches(PsiClassStub stub);
+    boolean matches(PsiClass psiClass);

    class ByName implements InternalNameMatcher {
      private final String myName;
@@ -96,8 +91,8 @@ public class JavaCompilerElementRetriever {
      public ByName(String name) {myName = name;}

      @Override
-      public boolean matches(PsiClassStub stub) {
-        return myName.equals(stub.getName());
+      public boolean matches(PsiClass psiClass) {
+        return myName.equals(psiClass.getName());
      }
    }

@@ -107,8 +102,8 @@ public class JavaCompilerElementRetriever {
      public ByQualifiedName(String name) {myQName = name;}

      @Override
-      public boolean matches(PsiClassStub stub) {
-        return myQName.equals(stub.getQualifiedName());
+      public boolean matches(PsiClass psiClass) {
+        return myQName.equals(psiClass.getQualifiedName());
      }
    }
  }
@@ -127,35 +122,32 @@ public class JavaCompilerElementRetriever {


    private PsiClass[] retrieveClasses(PsiFileWithStubSupport file) {
-      StubTree tree = file.getStubTree();
-      boolean foreign = tree == null;
-      if (foreign) {
-        tree = ((PsiFileImpl)file).calcStubTree();
-      }
+      StubbedSpine spine = file.getStubbedSpine();

      List<PsiClass> result = new ArrayList<>(myClassNameMatchers.size() + (myAnonymousIndices == null ? 0 : myAnonymousIndices.size()));
      int anonymousId = 0;
-      for (StubElement<?> element : tree.getPlainList()) {
-        if (element instanceof PsiClassStub) {
-          if (((PsiClassStub)element).isAnonymous()) {
+      for (int i = 0; i < spine.getStubCount(); i++) {
+        if (Constants.CLASS_BIT_SET.contains(spine.getStubType(i))) {
+          PsiClass element = (PsiClass)spine.getStubPsi(i);
+          if (element instanceof PsiAnonymousClass) {
            if (myAnonymousIndices != null && !myAnonymousIndices.isEmpty()) {
              if (myAnonymousIndices.contains(anonymousId)) {
-                result.add((PsiClass)element.getPsi());
+                result.add(element);
              }
              anonymousId++;
            }
          }
-          else if (match((PsiClassStub)element, myClassNameMatchers)) {
-            result.add((PsiClass)element.getPsi());
+          else if (match(element, myClassNameMatchers)) {
+            result.add(element);
          }
        }
      }
      return result.toArray(PsiClass.EMPTY_ARRAY);
    }

-    private static boolean match(PsiClassStub stub, Collection<InternalNameMatcher> matchers) {
+    private static boolean match(PsiClass psiClass, Collection<InternalNameMatcher> matchers) {
      for (InternalNameMatcher matcher : matchers) {
-        if (matcher.matches(stub)) {
+        if (matcher.matches(psiClass)) {
          //qualified name is unique among file's classes
          if (matcher instanceof InternalNameMatcher.ByQualifiedName) {
            matchers.remove(matcher);
@@ -22,10 +22,7 @@ import com.intellij.psi.PsiClass;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiManager;
-import com.intellij.psi.impl.source.PsiFileImpl;
 import com.intellij.psi.impl.source.PsiFileWithStubSupport;
-import com.intellij.psi.stubs.StubElement;
-import com.intellij.psi.stubs.StubTree;
 import com.intellij.util.ObjectUtils;
 import com.intellij.util.containers.IntIntHashMap;
 import gnu.trove.TByteArrayList;
@@ -33,8 +30,6 @@ import gnu.trove.TIntArrayList;
 import gnu.trove.TIntIntHashMap;
 import org.jetbrains.annotations.NotNull;

-import java.util.List;
-
 /**
 * @author peter
 */
@@ -97,30 +92,10 @@ class AnchorRepository {
  static PsiClass retrieveClass(@NotNull Project project, int fileId, int stubId) {
    PsiFile psiFile = PsiManager.getInstance(project).findFile(retrieveFile(fileId));
    assert psiFile != null : anchorToString(stubId, fileId);
-    PsiElement element = restoreFromStubIndex((PsiFileWithStubSupport)psiFile, stubId);
+    PsiElement element = ((PsiFileWithStubSupport)psiFile).getStubbedSpine().getStubPsi(stubId);
    if (!(element instanceof PsiClass)) {
      throw new AssertionError(anchorToString(stubId, fileId) + "; " + psiFile);
    }
    return (PsiClass)element;
  }
-
-  private static PsiElement restoreFromStubIndex(PsiFileWithStubSupport fileImpl, int index) {
-    StubTree tree = fileImpl.getStubTree();
-    if (tree == null) {
-      if (fileImpl instanceof PsiFileImpl) {
-        tree = ((PsiFileImpl)fileImpl).calcStubTree();
-      }
-      else {
-        return null;
-      }
-    }
-
-    List<StubElement<?>> list = tree.getPlainList();
-    if (index >= list.size()) {
-      return null;
-    }
-
-    return ((StubElement)list.get(index)).getPsi();
-  }
-
 }
@@ -24,15 +24,10 @@ import com.intellij.openapi.progress.ProgressManager;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.vfs.VirtualFile;
-import com.intellij.psi.CommonClassNames;
-import com.intellij.psi.PsiClass;
-import com.intellij.psi.PsiManager;
-import com.intellij.psi.PsiTypeParameter;
-import com.intellij.psi.impl.source.PsiFileImpl;
+import com.intellij.psi.*;
 import com.intellij.psi.impl.source.PsiFileWithStubSupport;
+import com.intellij.psi.impl.source.StubbedSpine;
 import com.intellij.psi.search.GlobalSearchScope;
-import com.intellij.psi.stubs.StubElement;
-import com.intellij.psi.stubs.StubTree;
 import com.intellij.psi.stubsHierarchy.ClassHierarchy;
 import com.intellij.psi.stubsHierarchy.HierarchyService;
 import com.intellij.psi.stubsHierarchy.SmartClassAnchor;
@@ -51,7 +46,7 @@ public class TestStubHierarchyAction extends InheritanceAction {
  public void actionPerformed(AnActionEvent e) {
    final Project project = e.getData(CommonDataKeys.PROJECT);
    if (project != null) {
-      ProgressManager.getInstance().runProcessWithProgressSynchronously((Runnable)() -> ReadAction.run(new TestHierarchy(project)::run),
+      ProgressManager.getInstance().runProcessWithProgressSynchronously(() -> ReadAction.run(new TestHierarchy(project)::run),
                                                                        "Testing Hierarchy", true, project);
    }
  }
@@ -89,8 +84,9 @@ public class TestStubHierarchyAction extends InheritanceAction {
    }

    private void checkFile(MultiMap<SmartClassAnchor, SmartClassAnchor> supers, ClassHierarchy hierarchy, VirtualFile vFile) {
-      for (StubElement<?> element : getStubTree(vFile).getPlainListFromAllRoots()) {
-        Object psi = element.getPsi();
+      StubbedSpine spine = getStubTree(vFile);
+      for (int i = 0; i < spine.getStubCount(); i++) {
+        PsiElement psi = spine.getStubPsi(i);
        if (psi instanceof PsiClass && !(psi instanceof PsiTypeParameter)) {
          SmartClassAnchor anchor = hierarchy.findAnchor((PsiClass)psi);
          if (anchor == null) {
@@ -116,11 +112,10 @@ public class TestStubHierarchyAction extends InheritanceAction {
    }

    @NotNull
-    private StubTree getStubTree(VirtualFile vFile) {
+    private StubbedSpine getStubTree(VirtualFile vFile) {
      PsiFileWithStubSupport psiFile = (PsiFileWithStubSupport)PsiManager.getInstance(myProject).findFile(vFile);
      assert psiFile != null : "No PSI for " + vFile;
-      StubTree stubTree = psiFile.getStubTree();
-      return stubTree != null ? stubTree : ((PsiFileImpl)psiFile).calcStubTree();
+      return psiFile.getStubbedSpine();
    }

    private void compareSupers(final SmartClassAnchor anchor, final Collection<SmartClassAnchor> superAnchors) {
@@ -52,6 +52,7 @@ import com.intellij.psi.impl.java.stubs.impl.PsiJavaFileStubImpl;
 import com.intellij.psi.impl.source.PsiFileImpl;
 import com.intellij.psi.impl.source.PsiFileWithStubSupport;
 import com.intellij.psi.impl.source.SourceTreeToPsiMap;
+import com.intellij.psi.impl.source.StubbedSpine;
 import com.intellij.psi.impl.source.resolve.FileContextUtil;
 import com.intellij.psi.impl.source.tree.JavaElementType;
 import com.intellij.psi.impl.source.tree.TreeElement;
@@ -561,6 +562,12 @@ public class ClsFileImpl extends ClsRepositoryPsiElement<PsiClassHolderFileStub>
    return null;
  }

+  @NotNull
+  @Override
+  public StubbedSpine getStubbedSpine() {
+    return getStubTree().getSpine();
+  }
+
  @Override
  public boolean isContentsLoaded() {
    return myStub != null;
@@ -15,31 +15,27 @@
 */
 package com.intellij.psi.stub;

-import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiMethod;
 import com.intellij.psi.impl.source.PsiFileImpl;
 import com.intellij.psi.impl.source.PsiMethodImpl;
+import com.intellij.psi.impl.source.StubbedSpine;
 import com.intellij.psi.impl.source.tree.JavaElementType;
-import com.intellij.psi.stubs.StubElement;
-import com.intellij.psi.stubs.StubTree;
-import com.intellij.util.ArrayUtil;

 public class JavaStubImplUtil {
  public static int getMethodStubIndex(PsiMethod method) {
    if (!(method instanceof PsiMethodImpl)) return -1;
    PsiFileImpl file = (PsiFileImpl)method.getContainingFile();
-    if (file.getElementTypeForStubBuilder() == null) return -1;
-    StubTree stubTree = file.getStubTree();
-    if (stubTree == null) {
-      stubTree = file.calcStubTree();
-    }
+    StubbedSpine spine = file.getStubbedSpine();

-    PsiElement[] stubs = stubTree
-      .getPlainList()
-      .stream()
-      .filter(e -> e.getStubType() == JavaElementType.METHOD)
-      .map(StubElement::getPsi)
-      .toArray(PsiElement[]::new);
-    return ArrayUtil.indexOf(stubs, method);
+    int result = 0;
+    for (int i = 0; i < spine.getStubCount(); i++) {
+      if (spine.getStubType(i) == JavaElementType.METHOD) {
+        if (spine.getStubPsi(i) == method) {
+          return result;
+        }
+        result++;
+      }
+    }
+    return -1;
  }
 }
@@ -33,19 +33,16 @@ import com.intellij.psi.impl.smartPointers.SelfElementInfo;
 import com.intellij.psi.impl.smartPointers.SmartPointerAnchorProvider;
 import com.intellij.psi.impl.source.PsiFileImpl;
 import com.intellij.psi.impl.source.PsiFileWithStubSupport;
-import com.intellij.psi.impl.source.tree.CompositeElement;
-import com.intellij.psi.impl.source.tree.FileElement;
+import com.intellij.psi.impl.source.StubbedSpine;
 import com.intellij.psi.stubs.IStubElementType;
 import com.intellij.psi.stubs.StubBase;
 import com.intellij.psi.stubs.StubElement;
-import com.intellij.psi.stubs.StubTree;
 import com.intellij.psi.tree.IStubFileElementType;
 import com.intellij.psi.util.PsiUtilCore;
 import org.jetbrains.annotations.NonNls;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;

-import java.util.List;
 import java.util.Set;

 /**
@@ -167,8 +164,7 @@ public abstract class PsiAnchor {
      return ((StubBase)liveStub).id;
    }

-    PsiFileImpl file = (PsiFileImpl)psi.getContainingFile();
-    return file.calcTreeElement().getStubbedDescendants().indexOf((CompositeElement)psi.getNode());
+    return ((PsiFileImpl)psi.getContainingFile()).calcTreeElement().getStubbedSpine().getStubIndex(psi);
  }

  private static class TreeRangeReference extends PsiAnchor {
@@ -400,51 +396,20 @@ public abstract class PsiAnchor {
      if (throwIfNull) throw new AssertionError("Null file");
      return null;
    }
-    StubTree tree = fileImpl.getStubTree();
-    if (tree != null) {
-      return findByStubs(index, elementType, throwIfNull, tree);
-    }
-    
-    if (fileImpl instanceof PsiFileImpl) {
-      return findByAst(index, elementType, throwIfNull, ((PsiFileImpl)fileImpl).calcTreeElement());
-    }
-    else {
-      if (throwIfNull) throw new AssertionError("Not PsiFileImpl: " + fileImpl.getClass());
-      return null;
-    }
-  }

-  @Nullable
-  private static PsiElement findByStubs(int index, @NotNull IStubElementType elementType, boolean throwIfNull, StubTree tree) {
-    List<StubElement<?>> list = tree.getPlainList();
-    if (index >= list.size()) {
-      if (throwIfNull) throw new AssertionError("Too large index: " + index + ">=" + list.size());
-      return null;
-    }
-    StubElement stub = list.get(index);
-
-    if (stub.getStubType() != elementType) {
-      if (throwIfNull) throw new AssertionError("Element type mismatch: " + stub.getStubType() + "!=" + elementType);
+    StubbedSpine spine = fileImpl.getStubbedSpine();
+    StubBasedPsiElement psi = spine.getStubPsi(index);
+    if (psi == null) {
+      if (throwIfNull) throw new AssertionError("Too large index: " + index + ">=" + spine.getStubCount());
      return null;
    }

-    return stub.getPsi();
-  }
-  @Nullable
-  private static PsiElement findByAst(int index, @NotNull IStubElementType elementType, boolean throwIfNull, FileElement tree) {
-    List<CompositeElement> list = tree.getStubbedDescendants();
-    if (index >= list.size()) {
-      if (throwIfNull) throw new AssertionError("Too large index: " + index + ">=" + list.size());
-      return null;
-    }
-    CompositeElement node = list.get(index);
-
-    if (node.getElementType() != elementType) {
-      if (throwIfNull) throw new AssertionError("Element type mismatch: " + node.getElementType() + "!=" + elementType);
+    if (psi.getElementType() != elementType) {
+      if (throwIfNull) throw new AssertionError("Element type mismatch: " + psi.getElementType() + "!=" + elementType);
      return null;
    }

-    return node.getPsi();
+    return psi;
  }

  public static class StubIndexReference extends PsiAnchor {
@@ -238,6 +238,13 @@ public abstract class PsiFileImpl extends ElementBase implements PsiFileEx, PsiF
    return findTreeForStub(ast, stubs, stub);
  }

+  @NotNull
+  @Override
+  public StubbedSpine getStubbedSpine() {
+    StubTree tree = getGreenStubTree();
+    return tree != null ? tree.getSpine() : calcTreeElement().getStubbedSpine();
+  }
+
  @Nullable
  private static ASTNode findTreeForStub(ASTNode tree, final Iterator<StubElement<?>> stubs, final StubElement stub) {
    final IElementType type = tree.getElementType();
@@ -24,6 +24,7 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.stubs.StubTree;
+import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;

 /**
@@ -44,4 +45,10 @@ public interface PsiFileWithStubSupport extends PsiFile {
   */
  @Nullable
  ASTNode findTreeForStub(StubTree tree, StubElement<?> stub);
+
+  /**
+   * @return StubbedSpine for accessing stubbed PSI, which can be backed up by stubs or AST
+   */
+  @NotNull
+  StubbedSpine getStubbedSpine();
 }
@@ -0,0 +1,40 @@
+/*
+ * Copyright 2000-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.intellij.psi.impl.source;
+
+import com.intellij.psi.StubBasedPsiElement;
+import com.intellij.psi.tree.IElementType;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * Represents the list of stubbed PSI elements in the file, backed by either AST or stubs.
+ */
+public interface StubbedSpine {
+  
+  /** @return the number of stubbed elements */
+  int getStubCount();
+
+  /** @return the stubbed PSI element at the given index or null if index is greater than the spine size */
+  @Nullable
+  StubBasedPsiElement getStubPsi(int index);
+
+  /** 
+   * @return the stubbed PSI element's type at the given index or null if index is greater than the spine size.
+   * This method avoids PSI element allocation if possible. 
+   */
+  @Nullable
+  IElementType getStubType(int index);
+}
@@ -0,0 +1,62 @@
+/*
+ * Copyright 2000-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.intellij.psi.impl.source.tree;
+
+import com.intellij.psi.StubBasedPsiElement;
+import com.intellij.psi.impl.source.StubbedSpine;
+import com.intellij.psi.tree.IElementType;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * @author peter
+ */
+public class AstSpine implements StubbedSpine {
+  static final AstSpine EMPTY_SPINE = new AstSpine(Collections.emptyList()); 
+  private final List<CompositeElement> myNodes;
+
+  AstSpine(@NotNull List<CompositeElement> nodes) {
+    myNodes = nodes;
+  }
+
+  @Override
+  public int getStubCount() {
+    return myNodes.size();
+  }
+
+  @Nullable
+  @Override
+  public StubBasedPsiElement getStubPsi(int index) {
+    return index >= myNodes.size() ? null : (StubBasedPsiElement)myNodes.get(index).getPsi();
+  }
+
+  public int getStubIndex(@NotNull StubBasedPsiElement psi) {
+    return myNodes.indexOf((CompositeElement)psi.getNode());
+  }
+
+  @Nullable
+  @Override
+  public IElementType getStubType(int index) {
+    return index >= myNodes.size() ? null : myNodes.get(index).getElementType();
+  }
+
+  List<CompositeElement> getNodes() {
+    return myNodes;
+  }
+}
@@ -34,14 +34,13 @@ import org.jetbrains.annotations.NotNull;
 import java.lang.ref.Reference;
 import java.lang.ref.WeakReference;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;

 public class FileElement extends LazyParseableElement implements FileASTNode, Getter<FileElement> {
  public static final FileElement[] EMPTY_ARRAY = new FileElement[0];
  private volatile CharTable myCharTable = new CharTableImpl();
  private volatile boolean myDetached;
-  private volatile Reference<List<CompositeElement>> myStubbedDescendants;
+  private volatile Reference<AstSpine> myStubbedSpine;

  @Override
  protected PsiElement createPsiNoLock() {
@@ -104,18 +103,18 @@ public class FileElement extends LazyParseableElement implements FileASTNode, Ge
  @Override
  public void clearCaches() {
    super.clearCaches();
-    myStubbedDescendants = null;
+    myStubbedSpine = null;
  }

  @NotNull
-  public final List<CompositeElement> getStubbedDescendants() {
-    List<CompositeElement> result = SoftReference.dereference(myStubbedDescendants);
+  public final AstSpine getStubbedSpine() {
+    AstSpine result = SoftReference.dereference(myStubbedSpine);
    if (result == null) {
      IStubFileElementType type = ((PsiFileImpl)getPsi()).getElementTypeForStubBuilder();
-      if (type == null) return Collections.emptyList();
+      if (type == null) return AstSpine.EMPTY_SPINE;

-      result = calcStubbedDescendants(type.getBuilder());
-      myStubbedDescendants = getManager().isBatchFilesProcessingMode() ? new WeakReference<>(result) : new SoftReference<>(result);
+      result = new AstSpine(calcStubbedDescendants(type.getBuilder()));
+      myStubbedSpine = getManager().isBatchFilesProcessingMode() ? new WeakReference<>(result) : new SoftReference<>(result);
    }
    return result;
  }
@@ -462,7 +462,7 @@ public class TreeUtil {
  @NotNull
  public static List<Pair<StubBase, TreeElement>> calcStubAstBindings(@NotNull StubTree stubTree, @NotNull FileElement tree) throws StubBindingException {
    PsiFile file = (PsiFile)tree.getPsi();
-    List<CompositeElement> nodes = tree.getStubbedDescendants();
+    List<CompositeElement> nodes = tree.getStubbedSpine().getNodes();
    List<StubElement<?>> stubs = stubTree.getPlainList();
    if (stubs.size() != nodes.size()) {
      throw new StubBindingException("Stub list in " + file.getName() + " length differs from PSI");
@@ -0,0 +1,54 @@
+/*
+ * Copyright 2000-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.intellij.psi.stubs;
+
+import com.intellij.psi.StubBasedPsiElement;
+import com.intellij.psi.impl.source.StubbedSpine;
+import com.intellij.psi.tree.IElementType;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.List;
+
+/**
+ * @author peter
+ */
+class StubSpine implements StubbedSpine {
+  private final StubTree myTree;
+
+  StubSpine(@NotNull StubTree tree) {
+    myTree = tree;
+  }
+
+  @Override
+  public int getStubCount() {
+    return myTree.getPlainList().size();
+  }
+
+  @Nullable
+  @Override
+  public StubBasedPsiElement getStubPsi(int index) {
+    List<StubElement<?>> stubs = myTree.getPlainList();
+    return index >= stubs.size() ? null : (StubBasedPsiElement)stubs.get(index).getPsi();
+  }
+
+  @Nullable
+  @Override
+  public IElementType getStubType(int index) {
+    List<StubElement<?>> stubs = myTree.getPlainList();
+    return index >= stubs.size() ? null : stubs.get(index).getStubType();
+  }
+}
@@ -19,14 +19,14 @@
 */
 package com.intellij.psi.stubs;

-import com.intellij.util.Function;
+import com.intellij.psi.impl.source.StubbedSpine;
 import com.intellij.util.containers.ContainerUtil;
 import org.jetbrains.annotations.NotNull;

-import java.util.Collection;
 import java.util.List;

 public class StubTree extends ObjectStubTree<StubElement<?>> {
+  private final StubSpine mySpine = new StubSpine(this);

  public StubTree(@NotNull final PsiFileStub root) {
    this(root, true);
@@ -54,4 +54,9 @@ public class StubTree extends ObjectStubTree<StubElement<?>> {
  public PsiFileStub getRoot() {
    return (PsiFileStub)myRoot;
  }
+
+  @NotNull
+  public StubbedSpine getSpine() {
+    return mySpine;
+  }
 }