Added JetCallableDeclaration.getColon()

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();
 }

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();

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() {

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

@@ -186,6 +186,12 @@ public class JetNamedFunction extends JetTypeParameterListOwnerStub<PsiJetFuncti
         return TypeRefHelpersPackage.setTypeReference(this, getValueParameterList(), typeRef);
     }
 
+    @Nullable
+    @Override
+    public PsiElement getColon() {
+        return findChildByType(JetTokens.COLON);
+    }
+
     @Override
     public boolean isLocal() {
         PsiElement parent = getParent();

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

@@ -19,6 +19,7 @@ package org.jetbrains.jet.lang.psi;
 import com.intellij.lang.ASTNode;
 import com.intellij.navigation.ItemPresentation;
 import com.intellij.navigation.ItemPresentationProviders;
+import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.psi.stubs.PsiJetParameterStub;
@@ -56,6 +57,12 @@ public class JetParameter extends JetNamedDeclarationStub<PsiJetParameterStub> 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) {

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

@@ -144,6 +144,12 @@ public class JetProperty extends JetTypeParameterListOwnerStub<PsiJetPropertyStu
         return TypeRefHelpersPackage.setTypeReference(this, getNameIdentifier(), typeRef);
     }
 
+    @Nullable
+    @Override
+    public PsiElement getColon() {
+        return findChildByType(JetTokens.COLON);
+    }
+
     @NotNull
     public List<JetPropertyAccessor> getAccessors() {
         return getStubOrPsiChildrenAsList(JetStubElementTypes.PROPERTY_ACCESSOR);

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<JetTypeReference>())
             .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)
         }

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())