diff --git a/idea/src/org/jetbrains/jet/lang/psi/JetDoWhileExpression.java b/idea/src/org/jetbrains/jet/lang/psi/JetDoWhileExpression.java index 45d24509c22..638ff95ef33 100644 --- a/idea/src/org/jetbrains/jet/lang/psi/JetDoWhileExpression.java +++ b/idea/src/org/jetbrains/jet/lang/psi/JetDoWhileExpression.java @@ -24,6 +24,6 @@ public class JetDoWhileExpression extends JetExpression { } public JetExpression getBody() { - return findExpressionUnder(JetNodeTypes.THEN); + return findExpressionUnder(JetNodeTypes.BODY); } } diff --git a/idea/src/org/jetbrains/jet/lang/psi/JetForExpression.java b/idea/src/org/jetbrains/jet/lang/psi/JetForExpression.java index 35bf5b7c77d..826f60c7a50 100644 --- a/idea/src/org/jetbrains/jet/lang/psi/JetForExpression.java +++ b/idea/src/org/jetbrains/jet/lang/psi/JetForExpression.java @@ -26,6 +26,6 @@ public class JetForExpression extends JetExpression { } public JetExpression getBody() { - return findExpressionUnder(JetNodeTypes.THEN); + return findExpressionUnder(JetNodeTypes.BODY); } } diff --git a/idea/src/org/jetbrains/jet/lang/psi/JetFunctionLiteralExpression.java b/idea/src/org/jetbrains/jet/lang/psi/JetFunctionLiteralExpression.java index e459a6ac365..5d255070594 100644 --- a/idea/src/org/jetbrains/jet/lang/psi/JetFunctionLiteralExpression.java +++ b/idea/src/org/jetbrains/jet/lang/psi/JetFunctionLiteralExpression.java @@ -1,7 +1,16 @@ package org.jetbrains.jet.lang.psi; import com.intellij.lang.ASTNode; +import com.intellij.psi.PsiElement; +import com.intellij.psi.tree.IElementType; +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.lexer.JetTokens; + +import java.util.Collections; +import java.util.List; /** * @author max @@ -15,4 +24,61 @@ public class JetFunctionLiteralExpression extends JetExpression { public void accept(JetVisitor visitor) { visitor.visitFunctionLiteralExpression(this); } + + //TODO: Due to the lack of multiple implementation inheritance these getters are copied from JetFunction + @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(); + } + + @NotNull + public ASTNode getBodyNode() { + ASTNode answer = getNode().findChildByType(JetNodeTypes.BODY); + assert answer != null; + return answer; + } + + @NotNull + public List getBody() { + return PsiTreeUtil.getChildrenOfTypeAsList(getBodyNode().getPsi(), 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; + } } diff --git a/idea/src/org/jetbrains/jet/lang/psi/JetWhileExpression.java b/idea/src/org/jetbrains/jet/lang/psi/JetWhileExpression.java index a4fcf901d3b..89c1b43b591 100644 --- a/idea/src/org/jetbrains/jet/lang/psi/JetWhileExpression.java +++ b/idea/src/org/jetbrains/jet/lang/psi/JetWhileExpression.java @@ -24,6 +24,6 @@ public class JetWhileExpression extends JetExpression { } public JetExpression getBody() { - return findExpressionUnder(JetNodeTypes.THEN); + return findExpressionUnder(JetNodeTypes.BODY); } }