From 78c1717eed852449459c3e973abcd44b0f76ef86 Mon Sep 17 00:00:00 2001
From: "Pavel V. Talanov" <talanov.pavel@gmail.com>
Date: Thu, 20 Mar 2014 13:37:32 +0400
Subject: [PATCH] JetFile#getImportList() and JetImportList#getImports() by
 stubs

---
 .../src/org/jetbrains/jet/lang/psi/JetFile.java | 17 +++++++++++++++--
 .../jet/lang/psi/JetImportDirective.java        | 11 ++++++++++-
 .../jetbrains/jet/lang/psi/JetImportList.java   |  5 ++---
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetFile.java b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetFile.java
index 5c27804f96a..1dcec7c7d80 100644
--- a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetFile.java
+++ b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetFile.java
@@ -24,11 +24,13 @@ import com.intellij.psi.FileViewProvider;
 import com.intellij.psi.PsiClass;
 import com.intellij.psi.PsiClassOwner;
 import com.intellij.psi.PsiElementVisitor;
+import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.util.PsiTreeUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.JetNodeTypes;
 import org.jetbrains.jet.lang.psi.stubs.PsiJetFileStub;
+import org.jetbrains.jet.lang.psi.stubs.elements.JetStubElementTypes;
 import org.jetbrains.jet.lang.resolve.name.FqName;
 import org.jetbrains.jet.plugin.JetFileType;
 import org.jetbrains.jet.plugin.JetLanguage;
@@ -73,6 +75,11 @@ public class JetFile extends PsiFileBase implements JetDeclarationContainer, Jet
 
     @Nullable
     public JetImportList getImportList() {
+        PsiJetFileStub stub = getStub();
+        if (stub != null) {
+            StubElement<JetImportList> importListStub = stub.findChildStubByType(JetStubElementTypes.IMPORT_LIST);
+            return importListStub != null ? importListStub.getPsi() : null;
+        }
         return findChildByClass(JetImportList.class);
     }
 
@@ -108,7 +115,7 @@ public class JetFile extends PsiFileBase implements JetDeclarationContainer, Jet
 
     @NotNull
     public FqName getPackageFqName() {
-        PsiJetFileStub stub = (PsiJetFileStub) getStub();
+        PsiJetFileStub stub = getStub();
         if (stub != null) {
             return stub.getPackageFqName();
         }
@@ -120,6 +127,12 @@ public class JetFile extends PsiFileBase implements JetDeclarationContainer, Jet
         return packageDirective.getFqName();
     }
 
+    @Override
+    @Nullable
+    public PsiJetFileStub getStub() {
+        return (PsiJetFileStub) super.getStub();
+    }
+
     @NotNull
     @Override
     public PsiClass[] getClasses() {
@@ -136,7 +149,7 @@ public class JetFile extends PsiFileBase implements JetDeclarationContainer, Jet
     }
 
     public boolean isScript() {
-        PsiJetFileStub stub = (PsiJetFileStub) getStub();
+        PsiJetFileStub stub = getStub();
         if (stub != null) {
             return stub.isScript();
         }
diff --git a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetImportDirective.java b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetImportDirective.java
index 5910915ed3e..999a8add1b2 100644
--- a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetImportDirective.java
+++ b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetImportDirective.java
@@ -17,8 +17,8 @@
 package org.jetbrains.jet.lang.psi;
 
 import com.intellij.lang.ASTNode;
-import com.intellij.psi.stubs.IStubElementType;
 import com.intellij.psi.tree.IElementType;
+import com.intellij.util.ArrayFactory;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.psi.stubs.PsiJetImportDirectiveStub;
@@ -26,6 +26,15 @@ import org.jetbrains.jet.lang.psi.stubs.elements.JetStubElementTypes;
 import org.jetbrains.jet.lexer.JetTokens;
 
 public class JetImportDirective extends JetElementImplStub<PsiJetImportDirectiveStub> {
+    public static final JetImportDirective[] EMPTY_ARRAY = new JetImportDirective[0];
+
+    public static final ArrayFactory<JetImportDirective> ARRAY_FACTORY = new ArrayFactory<JetImportDirective>() {
+        @Override
+        public JetImportDirective[] create(int count) {
+            return count == 0 ? EMPTY_ARRAY : new JetImportDirective[count];
+        }
+    };
+
     public JetImportDirective(@NotNull ASTNode node) {
         super(node);
     }
diff --git a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetImportList.java b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetImportList.java
index 62e6db41d49..963406a7716 100644
--- a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetImportList.java
+++ b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetImportList.java
@@ -17,12 +17,11 @@
 package org.jetbrains.jet.lang.psi;
 
 import com.intellij.lang.ASTNode;
-import com.intellij.psi.stubs.StubElement;
-import com.intellij.psi.util.PsiTreeUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.lang.psi.stubs.PsiJetImportListStub;
 import org.jetbrains.jet.lang.psi.stubs.elements.JetStubElementTypes;
 
+import java.util.Arrays;
 import java.util.List;
 
 public class JetImportList extends JetElementImplStub<PsiJetImportListStub> {
@@ -42,6 +41,6 @@ public class JetImportList extends JetElementImplStub<PsiJetImportListStub> {
 
     @NotNull
     public List<JetImportDirective> getImports() {
-        return PsiTreeUtil.getChildrenOfTypeAsList(this, JetImportDirective.class);
+        return Arrays.asList(getStubOrPsiChildren(JetStubElementTypes.IMPORT_DIRECTIVE, JetImportDirective.ARRAY_FACTORY));
     }
 }