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 ad71774db019..9ae2f3ebda16 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
@@ -277,52 +277,22 @@ public abstract class PsiFileImpl extends ElementBase implements PsiFileEx, PsiF
       return Collections.emptyList();
     }
 
-    final Iterator<StubElement<?>> stubs = stubTree.getPlainList().iterator();
-    stubs.next(); // Skip file stub;
-    final List<Pair<StubElement, AstPath>> result = ContainerUtil.newArrayList();
-    final IStubFileElementType elementType = getElementTypeForStubBuilder();
-    assert elementType != null;
-    final StubBuilder builder = elementType.getBuilder();
-
-    root.acceptTree(new RecursiveTreeElementWalkingVisitor() {
-      @Override
-      protected void visitNode(TreeElement node) {
-        CompositeElement parent = node.getTreeParent();
-        if (parent != null && builder.skipChildProcessingWhenBuildingStubs(parent, node)) {
-          return;
-        }
-
-
-        IElementType type = node.getElementType();
-        if (type instanceof IStubElementType && ((IStubElementType)type).shouldCreateStub(node)) {
-          if (!stubs.hasNext()) {
-            reportStubAstMismatch("Stub list is less than AST, last AST element: " + node.getElementType() + " " + node, stubTree);
-          }
-
-          final StubElement stub = stubs.next();
-          if (stub.getStubType() != node.getElementType()) {
-            reportStubAstMismatch("Stub and PSI element type mismatch in " + getName() + ": stub " + stub + ", AST " +
-                                  node.getElementType() + "; " + node, stubTree);
-          }
-
-          AstPath path = AstPath.getNodePath((CompositeElement)node);
-          assert path != null;
-          result.add(Pair.create(stub, path));
-        }
-
-        super.visitNode(node);
+    try {
+      List<Pair<StubBase, TreeElement>> result = TreeUtil.calcStubAstBindings(stubTree, root);
+      synchronized (myPsiLock) {
+        return ContainerUtil.map(result, pair -> {
+          StubElement stub = pair.first;
+          PsiElement psi = stub.getPsi();
+          assert psi != null : "Stub " + stub + " (" + stub.getClass() + ") has returned null PSI";
+          AstPath path = AstPath.getNodePath((CompositeElement)pair.second);
+          assert path != null : "Null path";
+          return Pair.create((StubBasedPsiElementBase)psi, path);
+        });
       }
-    });
-    if (stubs.hasNext()) {
-      reportStubAstMismatch("Stub list in " + getName() + " has more elements than PSI", stubTree);
     }
-    synchronized (myPsiLock) {
-      return ContainerUtil.map(result, pair -> {
-        StubElement stub = pair.first;
-        PsiElement psi = stub.getPsi();
-        assert psi != null : "Stub " + stub + " (" + stub.getClass() + ") has returned null PSI";
-        return Pair.create((StubBasedPsiElementBase)psi, pair.second);
-      });
+    catch (TreeUtil.StubBindingException e) {
+      reportStubAstMismatch(e.getMessage(), stubTree);
+      return Collections.emptyList();
     }
   }
 
@@ -741,7 +711,7 @@ public abstract class PsiFileImpl extends ElementBase implements PsiFileEx, PsiF
         // Set references from these stubs to AST, because:
         // Stub index might call getStubTree on main PSI file, but then use getPlainListFromAllRoots and return stubs from another file.
         // Even if that file already has AST, stub.getPsi() should be the same as in AST
-        TreeUtil.bindStubsToTree(eachPsiRoot, stubTree, fileElement);
+        TreeUtil.bindStubsToTree(stubTree, fileElement);
       } else {
         eachPsiRoot.bindStubsToCachedPsi(stubTree);
         bindings.put(eachPsiRoot, stubTree);
@@ -1103,7 +1073,7 @@ public abstract class PsiFileImpl extends ElementBase implements PsiFileEx, PsiF
         tree = new StubTree((PsiFileStub)currentStubTree);
         tree.setDebugInfo("created in calcStubTree");
         try {
-          TreeUtil.bindStubsToTree(this, tree, fileElement);
+          TreeUtil.bindStubsToTree(tree, fileElement);
         }
         catch (TreeUtil.StubBindingException e) {
           rebuildStub();
diff --git a/platform/core-impl/src/com/intellij/psi/impl/source/tree/TreeUtil.java b/platform/core-impl/src/com/intellij/psi/impl/source/tree/TreeUtil.java
index 4b92db07228a..4df4a6884595 100644
--- a/platform/core-impl/src/com/intellij/psi/impl/source/tree/TreeUtil.java
+++ b/platform/core-impl/src/com/intellij/psi/impl/source/tree/TreeUtil.java
@@ -23,6 +23,7 @@ import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.progress.ProgressIndicator;
 import com.intellij.openapi.util.Couple;
 import com.intellij.openapi.util.Key;
+import com.intellij.openapi.util.Pair;
 import com.intellij.psi.PsiComment;
 import com.intellij.psi.PsiWhiteSpace;
 import com.intellij.psi.StubBuilder;
@@ -40,10 +41,7 @@ import com.intellij.psi.tree.TokenSet;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-import java.util.HashSet;
-import java.util.LinkedList;
-import java.util.ListIterator;
-import java.util.Set;
+import java.util.*;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 public class TreeUtil {
@@ -452,10 +450,26 @@ public class TreeUtil {
     }
   }
 
-  public static void bindStubsToTree(@NotNull PsiFileImpl file, @NotNull StubTree stubTree, @NotNull FileElement tree) throws StubBindingException {
+  public static void bindStubsToTree(@NotNull StubTree stubTree, @NotNull FileElement tree) throws StubBindingException {
+    List<Pair<StubBase, TreeElement>> bindings = calcStubAstBindings(stubTree, tree);
+
+    for (int i = 0; i < bindings.size(); i++) {
+      Pair<StubBase, TreeElement> pair = bindings.get(i);
+      StubBasedPsiElementBase psi = (StubBasedPsiElementBase)pair.second.getPsi();
+      //noinspection unchecked
+      pair.first.setPsi(psi);
+      psi.setStubIndex(i + 1);
+    }
+  }
+
+  @NotNull
+  public static List<Pair<StubBase, TreeElement>> calcStubAstBindings(@NotNull StubTree stubTree, @NotNull FileElement tree) throws StubBindingException {
+    List<Pair<StubBase, TreeElement>> bindings = new ArrayList<>();
+
     final ListIterator<StubElement<?>> stubs = stubTree.getPlainList().listIterator();
     stubs.next();  // skip file root stub
 
+    PsiFileImpl file = (PsiFileImpl)tree.getPsi();
     final IStubFileElementType type = file.getElementTypeForStubBuilder();
     assert type != null;
     final StubBuilder builder = type.getBuilder();
@@ -474,15 +488,16 @@ public class TreeUtil {
             throw new StubBindingException("stub:" + stub + ", AST:" + type);
           }
 
-          StubBasedPsiElementBase psi = (StubBasedPsiElementBase)node.getPsi();
-          //noinspection unchecked
-          ((StubBase)stub).setPsi(psi);
-          psi.setStubIndex(stubs.previousIndex());
+          bindings.add(Pair.create((StubBase)stub, node));
         }
 
         super.visitNode(node);
       }
     });
+    if (stubs.hasNext()) {
+      throw new StubBindingException("Stub list in " + file.getName() + " has more elements than PSI");
+    }
+    return bindings;
   }
 
   @Nullable