diff --git a/idea/src/org/jetbrains/jet/lang/psi/JetFunction.java b/idea/src/org/jetbrains/jet/lang/psi/JetFunction.java index d332f49e3c7..73414e49175 100644 --- a/idea/src/org/jetbrains/jet/lang/psi/JetFunction.java +++ b/idea/src/org/jetbrains/jet/lang/psi/JetFunction.java @@ -1,7 +1,15 @@ package org.jetbrains.jet.lang.psi; import com.intellij.lang.ASTNode; +import com.intellij.psi.PsiElement; +import com.intellij.psi.tree.IElementType; import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.Nullable; +import org.jetbrains.jet.JetNodeTypes; +import org.jetbrains.jet.lexer.JetTokens; + +import java.util.Collections; +import java.util.List; /** * @author max @@ -15,4 +23,54 @@ public class JetFunction extends JetTypeParameterListOwner { public void accept(@NotNull JetVisitor visitor) { visitor.visitFunction(this); } + + @Nullable + public JetParameterList getParameterList() { + return (JetParameterList) findChildByType(JetNodeTypes.VALUE_PARAMETER_LIST); + } + + @NotNull + public List getParameters() { + JetParameterList list = getParameterList(); + return list != null ? list.getParameters() : Collections.emptyList(); + } + + @Nullable + public JetExpression getBodyExpression() { + return findChildByClass(JetExpression.class); + } + + + @Nullable + public JetTypeReference getReceiverTypeRef() { + PsiElement child = getFirstChild(); + while (child != null) { + IElementType tt = child.getNode().getElementType(); + if (tt == JetTokens.LPAR || tt == JetTokens.COLON) break; + if (child instanceof JetTypeReference) { + return (JetTypeReference) child; + } + child = child.getNextSibling(); + } + + return null; + } + + @Nullable + public JetTypeReference getReturnTypeRef() { + boolean colonPassed = false; + PsiElement child = getFirstChild(); + while (child != null) { + IElementType tt = child.getNode().getElementType(); + if (tt == JetTokens.COLON) { + colonPassed = true; + } + if (colonPassed && child instanceof JetTypeReference) { + return (JetTypeReference) child; + } + child = child.getNextSibling(); + } + + return null; + } }