Added API for changing type ref for function/variable and refactored code in some quickfixes to use it

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

@@ -27,4 +27,7 @@ public interface JetCallableDeclaration extends JetNamedDeclaration, JetTypePara
 
     @Nullable
     JetTypeReference getReturnTypeRef();
+
+    @Nullable
+    JetTypeReference setReturnTypeRef(@Nullable JetTypeReference typeRef);
 }

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

@@ -22,6 +22,7 @@ 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.lang.psi.typeRefHelpers.TypeRefHelpersPackage;
 import org.jetbrains.jet.lexer.JetTokens;
 
 import java.util.Collections;
@@ -77,20 +78,13 @@ abstract public class JetFunctionNotStubbed extends JetTypeParameterListOwnerNot
     @Override
     @Nullable
     public JetTypeReference getReturnTypeRef() {
-        boolean colonPassed = false;
+        return TypeRefHelpersPackage.getTypeRef(this);
-        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;
+    @Nullable
+    @Override
+    public JetTypeReference setReturnTypeRef(@Nullable JetTypeReference typeRef) {
+        return TypeRefHelpersPackage.setTypeRef(this, getValueParameterList(), typeRef);
     }
 
     @Override

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

@@ -21,6 +21,7 @@ import com.intellij.psi.tree.TokenSet;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.JetNodeTypes;
+import org.jetbrains.jet.lang.psi.typeRefHelpers.TypeRefHelpersPackage;
 import org.jetbrains.jet.lang.resolve.name.FqName;
 import org.jetbrains.jet.lexer.JetTokens;
 
@@ -34,7 +35,12 @@ public class JetMultiDeclarationEntry extends JetNamedDeclarationNotStubbed impl
 
     @Override
     public JetTypeReference getTypeRef() {
-        return (JetTypeReference) findChildByType(JetNodeTypes.TYPE_REFERENCE);
+        return TypeRefHelpersPackage.getTypeRef(this);
+    }
+
+    @Nullable
+    public JetTypeReference setTypeRef(@Nullable JetTypeReference typeRef) {
+        return TypeRefHelpersPackage.setTypeRef(this, getNameIdentifier(), typeRef);
     }
 
     @Override

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

@@ -26,6 +26,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.psi.stubs.PsiJetFunctionStub;
 import org.jetbrains.jet.lang.psi.stubs.elements.JetStubElementTypes;
+import org.jetbrains.jet.lang.psi.typeRefHelpers.TypeRefHelpersPackage;
 import org.jetbrains.jet.lexer.JetTokens;
 
 import java.util.Collections;
@@ -174,25 +175,13 @@ public class JetNamedFunction extends JetTypeParameterListOwnerStub<PsiJetFuncti
             }
             return typeReferences.get(returnTypeIndex);
         }
-        return getReturnTypeRefByPsi();
+        return TypeRefHelpersPackage.getTypeRef(this);
     }
 
+    @Override
     @Nullable
-    private JetTypeReference getReturnTypeRefByPsi() {
+    public JetTypeReference setReturnTypeRef(@Nullable JetTypeReference typeRef) {
-        boolean colonPassed = false;
+        return TypeRefHelpersPackage.setTypeRef(this, getValueParameterList(), typeRef);
-        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;
     }
 
     @Override

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

@@ -23,6 +23,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.psi.stubs.PsiJetParameterStub;
 import org.jetbrains.jet.lang.psi.stubs.elements.JetStubElementTypes;
+import org.jetbrains.jet.lang.psi.typeRefHelpers.TypeRefHelpersPackage;
 import org.jetbrains.jet.lexer.JetTokens;
 
 public class JetParameter extends JetNamedDeclarationStub<PsiJetParameterStub> {
@@ -40,11 +41,17 @@ public class JetParameter extends JetNamedDeclarationStub<PsiJetParameterStub> {
         return visitor.visitParameter(this, data);
     }
 
+    //TODO: rename it to getTypeRef for consistency
     @Nullable
     public JetTypeReference getTypeReference() {
         return getStubOrPsiChild(JetStubElementTypes.TYPE_REFERENCE);
     }
 
+    @Nullable
+    public JetTypeReference setTypeRef(@Nullable JetTypeReference typeRef) {
+        return TypeRefHelpersPackage.setTypeRef(this, getNameIdentifier(), typeRef);
+    }
+
     public boolean hasDefaultValue() {
         PsiJetParameterStub stub = getStub();
         if (stub != null) {

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

@@ -29,6 +29,7 @@ import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.JetNodeTypes;
 import org.jetbrains.jet.lang.psi.stubs.PsiJetPropertyStub;
 import org.jetbrains.jet.lang.psi.stubs.elements.JetStubElementTypes;
+import org.jetbrains.jet.lang.psi.typeRefHelpers.TypeRefHelpersPackage;
 import org.jetbrains.jet.lexer.JetTokens;
 
 import java.util.List;
@@ -121,6 +122,12 @@ public class JetProperty extends JetTypeParameterListOwnerStub<PsiJetPropertyStu
         return getTypeRef();
     }
 
+    @Nullable
+    @Override
+    public JetTypeReference setReturnTypeRef(@Nullable JetTypeReference typeRef) {
+        return setTypeRef(typeRef);
+    }
+
     @Override
     @Nullable
     public JetTypeReference getTypeRef() {
@@ -139,25 +146,12 @@ public class JetProperty extends JetTypeParameterListOwnerStub<PsiJetPropertyStu
                 return typeReferences.get(returnTypeRefPositionInPsi);
             }
         }
-        return getTypeRefByTree();
+        return TypeRefHelpersPackage.getTypeRef(this);
     }
 
     @Nullable
-    private JetTypeReference getTypeRefByTree() {
+    public JetTypeReference setTypeRef(@Nullable JetTypeReference typeRef) {
-        ASTNode node = getNode().getFirstChildNode();
+        return TypeRefHelpersPackage.setTypeRef(this, getNameIdentifier(), typeRef);
-        boolean passedColon = false;
-        while (node != null) {
-            IElementType tt = node.getElementType();
-            if (tt == JetTokens.COLON) {
-                passedColon = true;
-            }
-            else if (tt == JetNodeTypes.TYPE_REFERENCE && passedColon) {
-                return (JetTypeReference) node.getPsi();
-            }
-            node = node.getTreeNext();
-        }
-
-        return null;
     }
 
     @NotNull

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

@@ -27,4 +27,7 @@ public interface JetVariableDeclaration extends JetDeclaration, JetWithExpressio
 
     @Nullable
     JetTypeReference getTypeRef();
+
+    @Nullable
+    JetTypeReference setTypeRef(@Nullable JetTypeReference typeRef);
 }

compiler/frontend/src/org/jetbrains/jet/lang/psi/TypeRefHelpers.kt

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2010-2014 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+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
+
+fun getTypeRef(declaration: JetDeclaration): JetTypeReference? {
+    return declaration.getFirstChild()!!.siblings(forward = true)
+            .dropWhile { it.getNode()!!.getElementType() != JetTokens.COLON }
+            .filterIsInstance(javaClass<JetTypeReference>())
+            .firstOrNull()
+}
+
+fun setTypeRef(declaration: JetNamedDeclaration, addAfter: PsiElement?, typeRef: JetTypeReference?): JetTypeReference? {
+    val oldTypeRef = getTypeRef(declaration)
+    if (typeRef != null) {
+        if (oldTypeRef != null) {
+            return oldTypeRef.replace(typeRef) as JetTypeReference
+        }
+        else {
+            var anchor = addAfter ?: declaration.getNameIdentifier()
+            val newTypeRef = declaration.addAfter(typeRef, anchor) as JetTypeReference
+            declaration.addAfter(JetPsiFactory(declaration.getProject()).createColon(), anchor)
+            return newTypeRef
+        }
+    }
+    else {
+        if (oldTypeRef != null) {
+            val colon = declaration.getNode()!!.findChildByType(JetTokens.COLON)!!.getPsi()!!
+            val removeFrom = colon.getPrevSibling() as? PsiWhiteSpace ?: colon
+            declaration.deleteChildRange(removeFrom, oldTypeRef)
+        }
+        return null
+    }
+}

idea/src/org/jetbrains/jet/plugin/intentions/SpecifyTypeExplicitlyAction.java

@@ -61,7 +61,7 @@ public class SpecifyTypeExplicitlyAction extends PsiElementBaseIntentionAction {
     }
 
     @Override
-    public void invoke(@NotNull Project project, Editor editor, @NotNull PsiElement element) {
+    public void invoke(@NotNull Project project, @NotNull Editor editor, @NotNull PsiElement element) {
         JetTypeReference typeRefParent = PsiTreeUtil.getTopmostParentOfType(element, JetTypeReference.class);
         if (typeRefParent != null) {
             element = typeRefParent;
@@ -74,7 +74,7 @@ public class SpecifyTypeExplicitlyAction extends PsiElementBaseIntentionAction {
                 addTypeAnnotation(project, editor, property, type);
             }
             else {
-                removeTypeAnnotation(property);
+                property.setTypeRef(null);
             }
         }
         else if (parent instanceof JetParameter) {
@@ -83,7 +83,7 @@ public class SpecifyTypeExplicitlyAction extends PsiElementBaseIntentionAction {
                 addTypeAnnotation(project, editor, parameter, type);
             }
             else {
-                removeTypeAnnotation(parameter);
+                parameter.setTypeRef(null);
             }
         }
         else if (parent instanceof JetNamedFunction) {
@@ -97,7 +97,7 @@ public class SpecifyTypeExplicitlyAction extends PsiElementBaseIntentionAction {
     }
 
     @Override
-    public boolean isAvailable(@NotNull Project project, Editor editor, @NotNull PsiElement element) {
+    public boolean isAvailable(@NotNull Project project, @NotNull Editor editor, @NotNull PsiElement element) {
         if (element.getContainingFile() instanceof JetCodeFragment) {
             return false;
         }
@@ -181,43 +181,33 @@ public class SpecifyTypeExplicitlyAction extends PsiElementBaseIntentionAction {
             @NotNull JetProperty property,
             @NotNull JetType exprType
     ) {
-        if (property.getTypeRef() != null) {
+        if (property.getTypeRef() != null) return;
-            return;
-        }
 
-        PsiElement anchor = property.getNameIdentifier();
+        addTypeAnnotationWithTemplate(project, editor, property, exprType);
-        if (anchor == null) {
-            return;
-        }
-
-        addTypeAnnotation(project, editor, property, anchor, exprType);
     }
 
     public static void addTypeAnnotation(Project project, @Nullable Editor editor, JetFunction function, @NotNull JetType exprType) {
-        JetParameterList valueParameterList = function.getValueParameterList();
-        assert valueParameterList != null;
         if (editor != null) {
-            addTypeAnnotation(project, editor, function, valueParameterList, exprType);
+            addTypeAnnotationWithTemplate(project, editor, function, exprType);
         }
         else {
-            addTypeAnnotationSilently(project, function, valueParameterList);
+            function.setReturnTypeRef(anyTypeRef(project));
         }
     }
 
     public static void addTypeAnnotation(Project project, @Nullable Editor editor, JetParameter parameter, @NotNull JetType exprType) {
         if (editor != null) {
-            addTypeAnnotation(project, editor, parameter, parameter.getNameIdentifier(), exprType);
+            addTypeAnnotationWithTemplate(project, editor, parameter, exprType);
         }
         else {
-            addTypeAnnotationSilently(project, parameter, parameter.getNameIdentifier());
+            parameter.setTypeRef(anyTypeRef(project));
         }
     }
 
-    private static void addTypeAnnotation(
+    private static void addTypeAnnotationWithTemplate(
             @NotNull Project project,
             @NotNull Editor editor,
             @NotNull final JetNamedDeclaration namedDeclaration,
-            @NotNull PsiElement anchor,
             @NotNull JetType exprType
     ) {
         assert !exprType.isError() : "Unexpected error type: " + namedDeclaration.getText();
@@ -245,12 +235,13 @@ public class SpecifyTypeExplicitlyAction extends PsiElementBaseIntentionAction {
             }
         };
 
-        addTypeAnnotationSilently(project, namedDeclaration, anchor);
+        setTypeRef(namedDeclaration, anyTypeRef(project));
 
         PsiDocumentManager.getInstance(project).commitAllDocuments();
         PsiDocumentManager.getInstance(project).doPostponedOperationsAndUnblockDocument(editor.getDocument());
 
         JetTypeReference newTypeRef = getTypeRef(namedDeclaration);
+        assert newTypeRef != null;
         TemplateBuilderImpl builder = new TemplateBuilderImpl(newTypeRef);
         builder.replaceElement(newTypeRef, expression);
 
@@ -260,15 +251,15 @@ public class SpecifyTypeExplicitlyAction extends PsiElementBaseIntentionAction {
         manager.startTemplate(editor, builder.buildInlineTemplate(), new TemplateEditingAdapter() {
             @Override
             public void templateFinished(Template template, boolean brokenOff) {
-                ShortenReferences.INSTANCE$.process(getTypeRef(namedDeclaration));
+                JetTypeReference typeRef = getTypeRef(namedDeclaration);
+                assert typeRef != null;
+                ShortenReferences.INSTANCE$.process(typeRef);
             }
         });
     }
 
-    private static void addTypeAnnotationSilently(Project project, JetNamedDeclaration namedDeclaration, PsiElement anchor) {
+    private static JetTypeReference anyTypeRef(@NotNull Project project) {
-        JetPsiFactory psiFactory = JetPsiFactory(namedDeclaration);
+        return JetPsiFactory(project).createType("Any");
-        namedDeclaration.addAfter(psiFactory.createType("Any"), anchor);
-        namedDeclaration.addAfter(psiFactory.createColon(), anchor);
     }
 
     @Nullable
@@ -286,24 +277,18 @@ public class SpecifyTypeExplicitlyAction extends PsiElementBaseIntentionAction {
         return null;
     }
 
-    private static void removeTypeAnnotation(@Nullable PsiElement removeAfter, @Nullable JetTypeReference typeReference) {
+    @Nullable
-        if (removeAfter == null) return;
+    private static JetTypeReference setTypeRef(@NotNull JetNamedDeclaration namedDeclaration, @Nullable JetTypeReference typeRef) {
-        if (typeReference == null) return;
+        if (namedDeclaration instanceof JetProperty) {
-        PsiElement sibling = removeAfter.getNextSibling();
+            return ((JetProperty) namedDeclaration).setTypeRef(typeRef);
-        if (sibling == null) return;
+        }
-        PsiElement nextSibling = typeReference.getNextSibling();
+        else if (namedDeclaration instanceof JetParameter) {
-        sibling.getParent().getNode().removeRange(sibling.getNode(), nextSibling == null ? null : nextSibling.getNode());
+            return ((JetParameter) namedDeclaration).setTypeRef(typeRef);
-    }
+        }
-
+        else if (namedDeclaration instanceof JetFunction) {
-    public static void removeTypeAnnotation(JetVariableDeclaration property) {
+            return ((JetFunction) namedDeclaration).setReturnTypeRef(typeRef);
-        removeTypeAnnotation(property.getNameIdentifier(), property.getTypeRef());
+        }
-    }
+        assert false : "Wrong namedDeclaration: " + namedDeclaration.getText();
-
+        return null;
-    public static void removeTypeAnnotation(JetParameter parameter) {
-        removeTypeAnnotation(parameter.getNameIdentifier(), parameter.getTypeReference());
-    }
-
-    public static void removeTypeAnnotation(JetFunction function) {
-        removeTypeAnnotation(function.getValueParameterList(), function.getReturnTypeRef());
     }
 }

idea/src/org/jetbrains/jet/plugin/quickfix/ChangeFunctionReturnTypeFix.java

@@ -43,7 +43,6 @@ import org.jetbrains.jet.lang.types.checker.JetTypeChecker;
 import org.jetbrains.jet.lang.types.lang.KotlinBuiltIns;
 import org.jetbrains.jet.plugin.JetBundle;
 import org.jetbrains.jet.plugin.caches.resolve.ResolvePackage;
-import org.jetbrains.jet.plugin.intentions.SpecifyTypeExplicitlyAction;
 import org.jetbrains.jet.renderer.DescriptorRenderer;
 
 import java.util.LinkedList;
@@ -109,7 +108,7 @@ public class ChangeFunctionReturnTypeFix extends JetIntentionAction<JetFunction>
             changeFunctionLiteralReturnTypeFix.invoke(project, editor, file);
         }
         else {
-            SpecifyTypeExplicitlyAction.removeTypeAnnotation(element);
+            element.setReturnTypeRef(null);
             if (!(KotlinBuiltIns.getInstance().isUnit(type) && element.hasBlockBody())) {
                 addReturnTypeAnnotation(element, renderedType);
             }

idea/src/org/jetbrains/jet/plugin/quickfix/ChangeVariableTypeFix.java

@@ -39,7 +39,6 @@ import org.jetbrains.jet.lang.types.JetType;
 import org.jetbrains.jet.lang.types.checker.JetTypeChecker;
 import org.jetbrains.jet.plugin.JetBundle;
 import org.jetbrains.jet.plugin.caches.resolve.ResolvePackage;
-import org.jetbrains.jet.plugin.intentions.SpecifyTypeExplicitlyAction;
 import org.jetbrains.jet.renderer.DescriptorRenderer;
 
 import java.util.LinkedList;
@@ -82,7 +81,7 @@ public class ChangeVariableTypeFix extends JetIntentionAction<JetVariableDeclara
 
     @Override
     public void invoke(@NotNull Project project, Editor editor, JetFile file) throws IncorrectOperationException {
-        SpecifyTypeExplicitlyAction.removeTypeAnnotation(element);
+        element.setTypeRef(null);
         PsiElement nameIdentifier = element.getNameIdentifier();
         assert nameIdentifier != null : "ChangeVariableTypeFix applied to variable without name";
         JetPsiFactory psiFactory = JetPsiFactory(file);

idea/src/org/jetbrains/jet/plugin/refactoring/changeSignature/usages/JetFunctionDefinitionUsage.java

@@ -23,7 +23,6 @@ import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.types.lang.KotlinBuiltIns;
 import org.jetbrains.jet.lexer.JetKeywordToken;
 import org.jetbrains.jet.lexer.JetTokens;
-import org.jetbrains.jet.plugin.intentions.SpecifyTypeExplicitlyAction;
 import org.jetbrains.jet.plugin.quickfix.AddModifierFix;
 import org.jetbrains.jet.plugin.quickfix.ChangeFunctionReturnTypeFix;
 import org.jetbrains.jet.plugin.quickfix.ChangeVisibilityModifierFix;
@@ -59,7 +58,7 @@ public class JetFunctionDefinitionUsage extends JetUsageInfo<PsiElement> {
                 }
             }
             if (changeInfo.isReturnTypeChanged()) {
-                SpecifyTypeExplicitlyAction.removeTypeAnnotation(function);
+                function.setReturnTypeRef(null);
                 String returnTypeText = changeInfo.getNewReturnTypeText();
 
                 //TODO use ChangeFunctionReturnTypeFix.invoke when JetTypeCodeFragment.getType() is ready

idea/src/org/jetbrains/jet/plugin/refactoring/introduce/introduceVariable/KotlinInplaceVariableIntroducer.java

@@ -121,7 +121,7 @@ public class KotlinInplaceVariableIntroducer extends InplaceVariableIntroducer<J
                                 PsiDocumentManager.getInstance(myProject).commitDocument(myEditor.getDocument());
                             }
                             else {
-                                SpecifyTypeExplicitlyAction.removeTypeAnnotation(myProperty);
+                                myProperty.setTypeRef(null);
                             }
                         }
                     }.execute();