better api for lighter ast + caching

platform/core-api/src/com/intellij/lang/FCTSBackedLighterAST.java

@@ -23,7 +23,7 @@ import org.jetbrains.annotations.NotNull;
 
 import java.util.List;
 
-class FCTSBackedLighterAST extends LighterAST {
+public class FCTSBackedLighterAST extends LighterAST {
   private final FlyweightCapableTreeStructure<LighterASTNode> myTreeStructure;
 
   public FCTSBackedLighterAST(final CharTable charTable, final FlyweightCapableTreeStructure<LighterASTNode> treeStructure) {

platform/core-api/src/com/intellij/lang/FileASTNode.java

@@ -17,6 +17,7 @@ package com.intellij.lang;
 
 import com.intellij.util.CharTable;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 /**
  * Top-level node in AST, corresponds to PsiFile.
@@ -28,4 +29,7 @@ public interface FileASTNode extends ASTNode {
   CharTable getCharTable();
 
   boolean isParsed();
+
+  @Nullable
+  LighterAST getLighterAST();
 }

platform/core-api/src/com/intellij/lang/LighterAST.java

@@ -15,8 +15,6 @@
  */
 package com.intellij.lang;
 
-import com.intellij.psi.tree.IFileElementType;
-import com.intellij.psi.tree.ILightStubFileElementType;
 import com.intellij.util.CharTable;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -46,19 +44,4 @@ public abstract class LighterAST {
 
   @NotNull
   public abstract List<LighterASTNode> getChildren(@NotNull final LighterASTNode parent);
-
-  public static @NotNull LighterAST getLighterASTFromFileAST(@NotNull FileASTNode node, @NotNull Language language) {
-    final IFileElementType contentType = LanguageParserDefinitions.INSTANCE.forLanguage(language).getFileNodeType();
-    assert contentType instanceof ILightStubFileElementType;
-
-    final LighterAST tree;
-    if (!node.isParsed()) {
-      final ILightStubFileElementType<?> type = (ILightStubFileElementType)contentType;
-      tree = new FCTSBackedLighterAST(node.getCharTable(), type.parseContentsLight(node));
-    }
-    else {
-      tree = new TreeBackedLighterAST(node);
-    }
-    return tree;
-  }
 }

platform/core-api/src/com/intellij/psi/stubs/LightStubBuilder.java

@@ -47,7 +47,7 @@ public class LightStubBuilder implements StubBuilder {
       LOG.error("File is not of ILightStubFileElementType: " + contentType + ", " + file);
       return null;
     }
-    final LighterAST tree = LighterAST.getLighterASTFromFileAST(file.getNode(), language);
+    final LighterAST tree = file.getNode().getLighterAST();
     if (tree == null) return null;
 
     final StubElement rootStub = createStubForFile(file, tree);

platform/core-impl/src/com/intellij/psi/impl/source/tree/FileElement.java

@@ -16,16 +16,19 @@
 
 package com.intellij.psi.impl.source.tree;
 
-import com.intellij.lang.ASTNode;
-import com.intellij.lang.FileASTNode;
+import com.intellij.lang.*;
 import com.intellij.openapi.util.Getter;
+import com.intellij.openapi.util.Key;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.impl.PsiManagerEx;
 import com.intellij.psi.impl.source.CharTableImpl;
 import com.intellij.psi.impl.source.PsiFileImpl;
 import com.intellij.psi.tree.IElementType;
+import com.intellij.psi.tree.IFileElementType;
+import com.intellij.psi.tree.ILightStubFileElementType;
 import com.intellij.util.CharTable;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 public class FileElement extends LazyParseableElement implements FileASTNode, Getter<FileElement> {
   private volatile CharTable myCharTable = new CharTableImpl();
@@ -47,6 +50,35 @@ public class FileElement extends LazyParseableElement implements FileASTNode, Ge
     return myCharTable;
   }
 
+  private static final Key<LighterAST> ourTreeKey = Key.create("lighter.key");
+
+  @Nullable
+  @Override
+  public LighterAST getLighterAST() {
+    final IFileElementType contentType = (IFileElementType)getElementType();
+    assert contentType instanceof ILightStubFileElementType;
+
+    LighterAST tree;
+    if (!isParsed()) {
+      tree = getUserData(ourTreeKey);
+      if (tree == null) {
+        final ILightStubFileElementType<?> type = (ILightStubFileElementType)contentType;
+        tree = new FCTSBackedLighterAST(getCharTable(), type.parseContentsLight(this));
+        putUserData(ourTreeKey, tree);
+      }
+    }
+    else {
+      tree = new TreeBackedLighterAST(this);
+    }
+    return tree;
+  }
+
+  @Override
+  protected void resetCachesOnParsedStateUpdate() {
+    putUserData(ourTreeKey, null);
+    super.resetCachesOnParsedStateUpdate();
+  }
+
   public FileElement(IElementType type, CharSequence text) {
     super(type, text);
   }

platform/core-impl/src/com/intellij/psi/impl/source/tree/LazyParseableElement.java

@@ -181,8 +181,8 @@ public class LazyParseableElement extends CompositeElement {
         if (rawFirstChild() != null) {
           LOG.error("Reentrant parsing?");
         }
-  
-        myText = null;
+
+        resetCachesOnParsedStateUpdate();
   
         if (parsedNode == null) return;
         super.rawAddChildrenWithoutNotifications((TreeElement)parsedNode);
@@ -199,6 +199,10 @@ public class LazyParseableElement extends CompositeElement {
     }
   }
 
+  protected void resetCachesOnParsedStateUpdate() {
+    myText = null;
+  }
+
   @Override
   public void rawAddChildrenWithoutNotifications(@NotNull TreeElement first) {
     if (myText() != null) {