diff --git a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetCallableDeclaration.java b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetCallableDeclaration.java index ba448d2d4b7..87a725b9e7c 100644 --- a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetCallableDeclaration.java +++ b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetCallableDeclaration.java @@ -16,6 +16,7 @@ package org.jetbrains.jet.lang.psi; +import com.intellij.psi.PsiElement; import org.jetbrains.annotations.Nullable; public interface JetCallableDeclaration extends JetNamedDeclaration, JetTypeParameterListOwner { @@ -30,4 +31,7 @@ public interface JetCallableDeclaration extends JetNamedDeclaration, JetTypePara @Nullable JetTypeReference setTypeReference(@Nullable JetTypeReference typeRef); + + @Nullable + PsiElement getColon(); } diff --git a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetFunctionNotStubbed.java b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetFunctionNotStubbed.java index d0f99ef79cd..1356e73de4c 100644 --- a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetFunctionNotStubbed.java +++ b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetFunctionNotStubbed.java @@ -87,6 +87,12 @@ abstract public class JetFunctionNotStubbed extends JetTypeParameterListOwnerNot return TypeRefHelpersPackage.setTypeReference(this, getValueParameterList(), typeRef); } + @Nullable + @Override + public PsiElement getColon() { + return findChildByType(JetTokens.COLON); + } + @Override public boolean isLocal() { PsiElement parent = getParent(); diff --git a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetMultiDeclarationEntry.java b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetMultiDeclarationEntry.java index d9c5a17b311..4c39f093696 100644 --- a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetMultiDeclarationEntry.java +++ b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetMultiDeclarationEntry.java @@ -17,6 +17,7 @@ package org.jetbrains.jet.lang.psi; import com.intellij.lang.ASTNode; +import com.intellij.psi.PsiElement; import com.intellij.psi.tree.TokenSet; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -47,6 +48,12 @@ public class JetMultiDeclarationEntry extends JetNamedDeclarationNotStubbed impl return TypeRefHelpersPackage.setTypeReference(this, getNameIdentifier(), typeRef); } + @Nullable + @Override + public PsiElement getColon() { + return findChildByType(JetTokens.COLON); + } + @Nullable @Override public JetParameterList getValueParameterList() { diff --git a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetNamedFunction.java b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetNamedFunction.java index e43435ce3de..6bee7de0eb2 100644 --- a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetNamedFunction.java +++ b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetNamedFunction.java @@ -186,6 +186,12 @@ public class JetNamedFunction extends JetTypeParameterListOwnerStub i return TypeRefHelpersPackage.setTypeReference(this, getNameIdentifier(), typeRef); } + @Nullable + @Override + public PsiElement getColon() { + return findChildByType(JetTokens.COLON); + } + public boolean hasDefaultValue() { PsiJetParameterStub stub = getStub(); if (stub != null) { diff --git a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetProperty.java b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetProperty.java index b5541335f47..ce660038948 100644 --- a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetProperty.java +++ b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetProperty.java @@ -144,6 +144,12 @@ public class JetProperty extends JetTypeParameterListOwnerStub getAccessors() { return getStubOrPsiChildrenAsList(JetStubElementTypes.PROPERTY_ACCESSOR); diff --git a/compiler/frontend/src/org/jetbrains/jet/lang/psi/TypeRefHelpers.kt b/compiler/frontend/src/org/jetbrains/jet/lang/psi/TypeRefHelpers.kt index 9b4c790e2f0..fc6e1698415 100644 --- a/compiler/frontend/src/org/jetbrains/jet/lang/psi/TypeRefHelpers.kt +++ b/compiler/frontend/src/org/jetbrains/jet/lang/psi/TypeRefHelpers.kt @@ -18,22 +18,21 @@ package org.jetbrains.jet.lang.psi.typeRefHelpers import org.jetbrains.jet.lang.psi.JetTypeReference import org.jetbrains.jet.lexer.JetTokens -import org.jetbrains.jet.lang.psi.JetDeclaration import org.jetbrains.jet.lang.psi.psiUtil.siblings import com.intellij.psi.PsiElement import org.jetbrains.jet.lang.psi.JetPsiFactory import com.intellij.psi.PsiWhiteSpace -import org.jetbrains.jet.lang.psi.JetNamedDeclaration import com.intellij.psi.PsiErrorElement +import org.jetbrains.jet.lang.psi.JetCallableDeclaration -fun getTypeReference(declaration: JetDeclaration): JetTypeReference? { +fun getTypeReference(declaration: JetCallableDeclaration): JetTypeReference? { return declaration.getFirstChild()!!.siblings(forward = true) .dropWhile { it.getNode()!!.getElementType() != JetTokens.COLON } .filterIsInstance(javaClass()) .firstOrNull() } -fun setTypeReference(declaration: JetNamedDeclaration, addAfter: PsiElement?, typeRef: JetTypeReference?): JetTypeReference? { +fun setTypeReference(declaration: JetCallableDeclaration, addAfter: PsiElement?, typeRef: JetTypeReference?): JetTypeReference? { val oldTypeRef = getTypeReference(declaration) if (typeRef != null) { if (oldTypeRef != null) { @@ -48,7 +47,7 @@ fun setTypeReference(declaration: JetNamedDeclaration, addAfter: PsiElement?, ty } else { if (oldTypeRef != null) { - val colon = declaration.getNode()!!.findChildByType(JetTokens.COLON)!!.getPsi()!! + val colon = declaration.getColon()!! val removeFrom = colon.getPrevSibling() as? PsiWhiteSpace ?: colon declaration.deleteChildRange(removeFrom, oldTypeRef) } diff --git a/idea/src/org/jetbrains/jet/plugin/intentions/ConvertToExpressionBodyAction.kt b/idea/src/org/jetbrains/jet/plugin/intentions/ConvertToExpressionBodyAction.kt index 7a6bbc5aa79..1dd4f4e2c9e 100644 --- a/idea/src/org/jetbrains/jet/plugin/intentions/ConvertToExpressionBodyAction.kt +++ b/idea/src/org/jetbrains/jet/plugin/intentions/ConvertToExpressionBodyAction.kt @@ -55,7 +55,7 @@ public class ConvertToExpressionBodyAction : PsiElementBaseIntentionAction() { if (declaration.hasDeclaredReturnType() && declaration is JetCallableDeclaration && canOmitType(declaration)) { val typeRef = declaration.getTypeReference()!! - val colon = typeRef.siblings(forward = false, withItself = false).first { it.getNode().getElementType() == JetTokens.COLON } + val colon = declaration.getColon()!! val range = TextRange(colon.getTextRange().getStartOffset(), typeRef.getTextRange().getEndOffset()) editor.getSelectionModel().setSelection(range.getStartOffset(), range.getEndOffset()) editor.getCaretModel().moveToOffset(range.getEndOffset())