Move utility to JetStubbedPsiUtil

Extract utility to obtain child of several possible element types

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetStubbedPsiUtil.java

@@ -18,7 +18,9 @@ package org.jetbrains.jet.lang.psi;
 
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.stubs.StubElement;
+import com.intellij.psi.tree.TokenSet;
 import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.ArrayFactory;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -52,6 +54,19 @@ public final class JetStubbedPsiUtil {
         return PsiTreeUtil.getParentOfType(element, declarationClass, strict);
     }
 
+    @Nullable
+    public static <T extends JetElement> T getStubOrPsiChild(
+            @NotNull JetElementImplStub<?> element,
+            @NotNull TokenSet types,
+            @NotNull ArrayFactory<T> factory
+    ) {
+        T[] typeElements = element.getStubOrPsiChildren(types, factory);
+        if (typeElements.length == 0) {
+            return null;
+        }
+        return typeElements[0];
+    }
+
     private JetStubbedPsiUtil() {
     }
 }

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetTypeElement.java

@@ -16,11 +16,22 @@
 
 package org.jetbrains.jet.lang.psi;
 
+import com.intellij.util.ArrayFactory;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.List;
 
 public interface JetTypeElement extends JetElement {
+    JetTypeElement[] EMPTY_ARRAY = new JetTypeElement[0];
+
+    ArrayFactory<JetTypeElement> ARRAY_FACTORY = new ArrayFactory<JetTypeElement>() {
+        @NotNull
+        @Override
+        public JetTypeElement[] create(int count) {
+            return count == 0 ? EMPTY_ARRAY : new JetTypeElement[count];
+        }
+    };
+
     @NotNull
     List<JetTypeReference> getTypeArgumentsAsTypes();
 }

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetTypeReference.java

@@ -17,11 +17,9 @@
 package org.jetbrains.jet.lang.psi;
 
 import com.intellij.lang.ASTNode;
-import com.intellij.psi.tree.TokenSet;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.psi.stubs.PsiJetPlaceHolderStub;
-import org.jetbrains.jet.lang.psi.stubs.elements.JetStubElementType;
 import org.jetbrains.jet.lang.psi.stubs.elements.JetStubElementTypes;
 
 import java.util.ArrayList;
@@ -36,9 +34,6 @@ import static org.jetbrains.jet.lang.psi.stubs.elements.JetStubElementTypes.ANNO
  */
 public class JetTypeReference extends JetElementImplStub<PsiJetPlaceHolderStub<JetTypeReference>> {
 
-    public static final TokenSet TYPE_ELEMENTS
-            = TokenSet.create(JetStubElementTypes.USER_TYPE, JetStubElementTypes.NULLABLE_TYPE, JetStubElementTypes.FUNCTION_TYPE);
-
     public JetTypeReference(@NotNull ASTNode node) {
         super(node);
     }
@@ -59,12 +54,7 @@ public class JetTypeReference extends JetElementImplStub<PsiJetPlaceHolderStub<J
 
     @Nullable
     public JetTypeElement getTypeElement() {
-        //TODO:
+        return JetStubbedPsiUtil.getStubOrPsiChild(this, JetStubElementTypes.TYPE_ELEMENT_TYPES, JetTypeElement.ARRAY_FACTORY);
-        JetTypeElement[] typeElements = getStubOrPsiChildren(TYPE_ELEMENTS, new JetTypeElement[] {});
-        if (typeElements.length == 0) {
-            return null;
-        }
-        return typeElements[0];
     }
 
     public List<JetAnnotationEntry> getAnnotations() {

compiler/frontend/src/org/jetbrains/jet/lang/psi/stubs/elements/JetStubElementTypes.java

@@ -100,4 +100,6 @@ public interface JetStubElementTypes {
               TokenSet.create(CLASS, OBJECT_DECLARATION, CLASS_OBJECT, FUNCTION, PROPERTY, ANONYMOUS_INITIALIZER, ENUM_ENTRY);
 
     TokenSet DELEGATION_SPECIFIER_TYPES = TokenSet.create(DELEGATOR_BY, DELEGATOR_SUPER_CALL, DELEGATOR_SUPER_CLASS, THIS_CALL);
+
+    TokenSet TYPE_ELEMENT_TYPES = TokenSet.create(USER_TYPE, NULLABLE_TYPE, FUNCTION_TYPE);
 }