Move caret to generated element and add selection

idea/idea-analysis/src/org/jetbrains/jet/plugin/quickfix/generateUtil.kt

@@ -0,0 +1,87 @@
+/*
+ * 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.plugin.quickfix
+
+import com.intellij.openapi.editor.Editor
+import com.intellij.psi.PsiElement
+import org.jetbrains.jet.lang.psi.JetDeclarationWithBody
+import org.jetbrains.jet.lang.psi.JetBlockExpression
+import com.intellij.psi.PsiWhiteSpace
+import com.intellij.openapi.editor.ScrollType
+import org.jetbrains.jet.lang.psi.JetWithExpressionInitializer
+import org.jetbrains.jet.lang.psi.JetProperty
+import org.jetbrains.jet.lang.psi.psiUtil.siblings
+
+private fun moveCaretIntoGeneratedElement(editor: Editor, element: PsiElement): Boolean {
+    // Inspired by GenerateMembersUtils.positionCaret()
+
+    if (element is JetDeclarationWithBody && element.hasBody()) {
+        val expression = element.getBodyExpression()
+        if (expression is JetBlockExpression) {
+            val lBrace = expression.getLBrace()
+            val rBrace = expression.getRBrace()
+
+            if (lBrace != null && rBrace != null) {
+                val firstInBlock = lBrace.siblings(forward = true, withItself = false).first { it !is PsiWhiteSpace }
+                val lastInBlock = rBrace.siblings(forward = false, withItself = false).first { it !is PsiWhiteSpace }
+
+                val start = firstInBlock.getTextRange()!!.getStartOffset()
+                val end = lastInBlock.getTextRange()!!.getEndOffset()
+
+                editor.moveCaret(Math.min(start, end))
+
+                if (start < end) {
+                    editor.getSelectionModel().setSelection(start, end)
+                }
+
+                return true
+            }
+        }
+    }
+
+    if (element is JetWithExpressionInitializer && element.hasInitializer()) {
+        val expression = element.getInitializer()
+        if (expression == null) throw AssertionError()
+
+        val initializerRange = expression.getTextRange()
+
+        val offset = initializerRange?.getStartOffset() ?: element.getTextOffset()
+
+        editor.moveCaret(offset)
+
+        if (initializerRange != null) {
+            editor.getSelectionModel().setSelection(initializerRange.getStartOffset(), initializerRange.getEndOffset())
+        }
+
+        return true
+    }
+
+    if (element is JetProperty) {
+        for (accessor in element.getAccessors()) {
+            if (moveCaretIntoGeneratedElement(editor, accessor)) {
+                return true
+            }
+        }
+    }
+
+    return false
+}
+
+public fun Editor.moveCaret(offset: Int, scrollType: ScrollType = ScrollType.RELATIVE) {
+    getCaretModel().moveToOffset(offset)
+    getScrollingModel().scrollToCaret(scrollType)
+}

idea/src/org/jetbrains/jet/plugin/codeInsight/OverrideImplementMethodsHandler.java

@@ -25,6 +25,7 @@ import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.ui.DialogWrapper;
+import com.intellij.openapi.util.Computable;
 import com.intellij.openapi.util.Condition;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.PsiDocumentManager;
@@ -39,6 +40,7 @@ import org.jetbrains.jet.lang.psi.*;
 import org.jetbrains.jet.lang.types.JetType;
 import org.jetbrains.jet.lang.types.lang.KotlinBuiltIns;
 import org.jetbrains.jet.plugin.caches.resolve.ResolvePackage;
+import org.jetbrains.jet.plugin.quickfix.QuickfixPackage;
 import org.jetbrains.jet.renderer.DescriptorRenderer;
 import org.jetbrains.jet.renderer.DescriptorRendererBuilder;
 
@@ -78,32 +80,50 @@ public abstract class OverrideImplementMethodsHandler implements LanguageCodeIns
     }
 
     public static void generateMethods(
-            @NotNull Editor editor,
+            @NotNull final Editor editor,
-            @NotNull JetClassOrObject classOrObject,
+            @NotNull final JetClassOrObject classOrObject,
-            @NotNull List<DescriptorClassMember> selectedElements
+            @NotNull final List<DescriptorClassMember> selectedElements
     ) {
-        JetClassBody body = classOrObject.getBody();
+        PsiElement firstGenerated = ApplicationManager.getApplication().runWriteAction(new Computable<PsiElement>() {
-        if (body == null) {
+            @Override
-            JetPsiFactory psiFactory = JetPsiFactory(classOrObject);
+            public PsiElement compute() {
-            classOrObject.add(psiFactory.createWhiteSpace());
+                JetClassBody body = classOrObject.getBody();
-            body = (JetClassBody) classOrObject.add(psiFactory.createEmptyClassBody());
+                if (body == null) {
+                    JetPsiFactory psiFactory = JetPsiFactory(classOrObject);
+                    classOrObject.add(psiFactory.createWhiteSpace());
+                    body = (JetClassBody) classOrObject.add(psiFactory.createEmptyClassBody());
+                }
+
+                PsiElement afterAnchor = findInsertAfterAnchor(editor, body);
+
+                if (afterAnchor == null) {
+                    return null;
+                }
+
+                PsiElement firstGenerated = null;
+
+                List<JetElement> elementsToCompact = new ArrayList<JetElement>();
+                JetFile file = classOrObject.getContainingJetFile();
+                for (JetElement element : generateOverridingMembers(selectedElements, file)) {
+                    PsiElement added = body.addAfter(element, afterAnchor);
+
+                    if (firstGenerated == null) {
+                        firstGenerated = added;
+                    }
+
+                    afterAnchor = added;
+                    elementsToCompact.add((JetElement) added);
+                }
+
+                ShortenReferences.INSTANCE$.process(elementsToCompact);
+
+                return firstGenerated;
+            }
+        });
+
+        if (firstGenerated != null) {
+            QuickfixPackage.moveCaretIntoGeneratedElement(editor, firstGenerated);
         }
-
-        PsiElement afterAnchor = findInsertAfterAnchor(editor, body);
-
-        if (afterAnchor == null) {
-            return;
-        }
-
-        List<JetElement> elementsToCompact = new ArrayList<JetElement>();
-        JetFile file = classOrObject.getContainingJetFile();
-        for (JetElement element : generateOverridingMembers(selectedElements, file)) {
-            PsiElement added = body.addAfter(element, afterAnchor);
-            afterAnchor = added;
-            elementsToCompact.add((JetElement) added);
-        }
-
-        ShortenReferences.INSTANCE$.process(elementsToCompact);
     }
 
     @Nullable
@@ -270,9 +290,9 @@ public abstract class OverrideImplementMethodsHandler implements LanguageCodeIns
 
     protected abstract String getNoMethodsFoundHint();
 
-    public void invoke(@NotNull Project project, @NotNull final Editor editor, @NotNull PsiFile file, boolean implementAll) {
+    public void invoke(@NotNull Project project, @NotNull Editor editor, @NotNull PsiFile file, boolean implementAll) {
         PsiElement elementAtCaret = file.findElementAt(editor.getCaretModel().getOffset());
-        final JetClassOrObject classOrObject = PsiTreeUtil.getParentOfType(elementAtCaret, JetClassOrObject.class);
+        JetClassOrObject classOrObject = PsiTreeUtil.getParentOfType(elementAtCaret, JetClassOrObject.class);
 
         assert classOrObject != null : "ClassObject should be checked in isValidFor method";
 
@@ -283,7 +303,7 @@ public abstract class OverrideImplementMethodsHandler implements LanguageCodeIns
         }
         List<DescriptorClassMember> members = membersFromDescriptors((JetFile) file, missingImplementations);
 
-        final List<DescriptorClassMember> selectedElements;
+        List<DescriptorClassMember> selectedElements;
         if (implementAll) {
             selectedElements = members;
         }
@@ -302,12 +322,7 @@ public abstract class OverrideImplementMethodsHandler implements LanguageCodeIns
 
         PsiDocumentManager.getInstance(project).commitAllDocuments();
 
-        ApplicationManager.getApplication().runWriteAction(new Runnable() {
+        generateMethods(editor, classOrObject, selectedElements);
-            @Override
-            public void run() {
-                generateMethods(editor, classOrObject, selectedElements);
-            }
-        });
     }
 
     @Override

idea/testData/codeInsight/overrideImplement/checkNotImportedTypesFromJava/foo/Impl.kt.after

@@ -6,7 +6,7 @@ import bar.Bar
 
 class Impl: Foo() {
     override fun foo(list: ArrayList<Int>?, other: Other?): Bar? {
-        throw UnsupportedOperationException()
+        <selection><caret>throw UnsupportedOperationException()</selection>
     }
 }
 

idea/testData/codeInsight/overrideImplement/complexMultiOverride.kt.after

@@ -6,7 +6,7 @@ open class Base<A, B, C>() {
 
 class C : Base<String, C, Unit>() {
     override fun bar(value: () -> Unit): (String) -> Unit {
-        return super<Base>.bar(value)
+        <selection><caret>return super<Base>.bar(value)</selection>
     }
     override fun equals(other: Any?): Boolean {
         return super<Base>.equals(other)

idea/testData/codeInsight/overrideImplement/delegatedMembers.kt.after

@@ -5,7 +5,7 @@ trait T {
 
 class C(t :T) : T by t {
     override fun bar() {
-        throw UnsupportedOperationException()
+        <selection><caret>throw UnsupportedOperationException()</selection>
     }
     override fun equals(other: Any?): Boolean {
         return super<T>.equals(other)

idea/testData/codeInsight/overrideImplement/emptyClassBodyFunctionMethod.kt.after

@@ -5,6 +5,6 @@ trait T {
 
 class C : T {
     override fun Foo(): (String) -> Unit {
-        throw UnsupportedOperationException()
+        <selection><caret>throw UnsupportedOperationException()</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/functionMethod.kt.after

@@ -5,6 +5,6 @@ trait T {
 
 class C : T {
     override fun Foo(): (String) -> Unit {
-        throw UnsupportedOperationException()
+        <selection><caret>throw UnsupportedOperationException()</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/functionProperty.kt

@@ -6,3 +6,5 @@ trait A {
 fun some() : A {
     return object : A {<caret>}
 }
+
+// TODO: need better selection and caret

idea/testData/codeInsight/overrideImplement/functionProperty.kt.after

@@ -4,7 +4,9 @@ trait A {
 }
 
 fun some() : A {
-    return object : A {
+    return object : A {<caret>
         override val method: () -> Unit? = ?
     }
 }
+
+// TODO: need better selection and caret

idea/testData/codeInsight/overrideImplement/functionWithTypeParameters.kt.after

@@ -4,6 +4,6 @@ trait Trait {
 
 class TraitImpl : Trait {
     override fun <A, B : Runnable, E : Map.Entry<A, B>> foo() where B : Cloneable, B : Comparable<B> {
-        throw UnsupportedOperationException()
+        <selection><caret>throw UnsupportedOperationException()</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/generateMulti.kt.after

@@ -8,7 +8,7 @@ trait Some<T> {
 
 class SomeOther<S> : Some<S> {
     override fun someFoo() {
-        throw UnsupportedOperationException()
+        <selection><caret>throw UnsupportedOperationException()</selection>
     }
     override fun someGenericFoo(): S {
         throw UnsupportedOperationException()

idea/testData/codeInsight/overrideImplement/genericMethod.kt.after

@@ -4,6 +4,6 @@ trait G<T> {
 
 class GC() : G<Int> {
     override fun foo(t: Int): Int {
-        throw UnsupportedOperationException()
+        <selection><caret>throw UnsupportedOperationException()</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/genericTypesSeveralMethods.kt.after

@@ -2,6 +2,6 @@
 
 class MyClass<A: Comparable<A>> : Iterable<A> {
     override fun iterator(): Iterator<A> {
-        throw UnsupportedOperationException()
+        <selection><caret>throw UnsupportedOperationException()</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/implementSamAdapters/foo/Impl.kt.after

@@ -2,6 +2,6 @@ package foo
 
 class Impl: B {
     override fun foo(r: Runnable?) {
-        throw UnsupportedOperationException()
+        <selection><caret>throw UnsupportedOperationException()</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/inheritVisibilities.kt.after

@@ -12,7 +12,7 @@ trait B {
 
 class C : A(), B {
     override fun equals(other: Any?): Boolean {
-        return super<A>.equals(other)
+        <selection><caret>return super<A>.equals(other)</selection>
     }
     override fun hashCode(): Int {
         return super<A>.hashCode()

idea/testData/codeInsight/overrideImplement/javaInterfaceMethod/foo/Impl.kt.after

@@ -2,6 +2,6 @@ import foo.Intf
 
 class Impl(): Intf {
     override fun getFooBar(): String? {
-        throw UnsupportedOperationException()
+        <selection><caret>throw UnsupportedOperationException()</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/javaMethodWithPackageProtectedVisibility/foo/Impl.kt.after

@@ -4,6 +4,6 @@ import foo.Intf
 
 class Impl(): Intf() {
     override fun getFooBar(): String? {
-        return super<Intf>.getFooBar()
+        <selection><caret>return super<Intf>.getFooBar()</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/javaMethodWithPackageVisibility/foo/Impl.kt.after

@@ -4,6 +4,6 @@ import foo.Intf
 
 class Impl(): Intf() {
     override fun getFooBar(): String? {
-        return super<Intf>.getFooBar()
+        <selection><caret>return super<Intf>.getFooBar()</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/javaParameters/foo/Impl.kt.after

@@ -2,6 +2,6 @@ import foo.Intf
 
 class Impl(): Intf {
     override fun fooBar(i: Int, s: Array<out String>?, foo: Any?) {
-        throw UnsupportedOperationException()
+        <selection><caret>throw UnsupportedOperationException()</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/multiOverride.kt.after

@@ -5,7 +5,7 @@ trait A {
 
 class C : A {
     override fun bar(): String {
-        return super<A>.bar()
+        <selection><caret>return super<A>.bar()</selection>
     }
     override fun equals(other: Any?): Boolean {
         return super<A>.equals(other)

idea/testData/codeInsight/overrideImplement/overrideExplicitFunction.kt.after

@@ -4,6 +4,6 @@ trait A {
 
 class B : A {
     override fun String.foo() {
-        throw UnsupportedOperationException()
+        <selection><caret>throw UnsupportedOperationException()</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/overrideExtensionProperty.kt.after

@@ -4,5 +4,5 @@ trait A {
 
 class B : A {
     override val String.prop: Int
-        get() = 0
+        get() = <selection><caret>0</selection>
 }

idea/testData/codeInsight/overrideImplement/overrideFunctionProperty.kt

@@ -5,3 +5,5 @@ open class A() {
 fun some() : A {
     return object : A() {<caret>}
 }
+
+// TODO: need better selection and caret

idea/testData/codeInsight/overrideImplement/overrideFunctionProperty.kt.after

@@ -3,7 +3,9 @@ open class A() {
 }
 
 fun some() : A {
-    return object : A() {
+    return object : A() {<caret>
         override val method: () -> Unit? = ?
     }
 }
+
+// TODO: need better selection and caret

idea/testData/codeInsight/overrideImplement/overrideGenericFunction.kt.after

@@ -4,6 +4,6 @@ trait A<T> {
 
 class C : A<C> {
     override fun foo(value: C) {
-        super<A>.foo(value)
+        <selection><caret>super<A>.foo(value)</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/overrideJavaMethod/foo/Impl.kt.after

@@ -2,6 +2,6 @@ import foo.A
 
 class C : A() {
     override fun getAnswer(array: Array<out String>?, number: Int, value: Any?): Int {
-        return super<A>.getAnswer(array, number, value)
+        <selection><caret>return super<A>.getAnswer(array, number, value)</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/overrideMutableExtensionProperty.kt.after

@@ -4,7 +4,7 @@ trait A {
 
 class B : A {
     override var Int.foo: Double
-        get() = 0.0
+        get() = <selection><caret>0.0</selection>
         set(value) {
         }
 }

idea/testData/codeInsight/overrideImplement/overrideNonUnitFunction.kt.after

@@ -4,6 +4,6 @@ trait A {
 
 class C : A {
     override fun foo(value: String): Int {
-        return super<A>.foo(value)
+        <selection><caret>return super<A>.foo(value)</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/overridePrimitiveProperty.kt.after

@@ -10,7 +10,7 @@ trait T {
 }
 
 class C : T {
-    override val a1: Byte = 0
+    override val a1: Byte = <selection><caret>0</selection>
     override val a2: Short = 0
     override val a3: Int = 0
     override val a4: Long = 0

idea/testData/codeInsight/overrideImplement/overrideRespectCaretPosition.kt.after

@@ -6,7 +6,7 @@ open class A() {
 class C : A() {
     val constant = 42
     // Some comment
-    override val bar: Int = 0
+    override val bar: Int = <selection><caret>0</selection>
     override fun equals(other: Any?): Boolean {
         return super<A>.equals(other)
     }

idea/testData/codeInsight/overrideImplement/overrideSamAdapters/foo/Impl.kt.after

@@ -2,6 +2,6 @@ package foo
 
 class Impl: B() {
     override fun foo(r: Runnable?) {
-        super<B>.foo(r)
+        <selection><caret>super<B>.foo(r)</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/overrideUnitFunction.kt.after

@@ -4,6 +4,6 @@ trait A {
 
 class C : A {
     override fun foo(value: String) {
-        super<A>.foo(value)
+        <selection><caret>super<A>.foo(value)</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/propagationKJK/foo/Impl.kt.after

@@ -2,6 +2,6 @@ package foo
 
 class Impl : Bar() {
     override fun f(): Any {
-        return super<Bar>.f()
+        <selection><caret>return super<Bar>.f()</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/property.kt.after

@@ -3,5 +3,5 @@ trait T {
 }
 
 class GC() : T {
-    override val v: Int = 0
+    override val v: Int = <selection><caret>0</selection>
 }

idea/testData/codeInsight/overrideImplement/respectCaretPosition.kt.after

@@ -6,7 +6,7 @@ trait Test {
 class SomeTest : Test {
     val hello = 12
     override fun test() {
-        throw UnsupportedOperationException()
+        <selection><caret>throw UnsupportedOperationException()</selection>
     }
     override val testProp: Int = 0
     /**

idea/testData/codeInsight/overrideImplement/sameTypeName/foo/Impl.kt.after

@@ -3,6 +3,6 @@ import lib.ArrayFactory
 
 public class Impl : ArrayFactory {
     override fun create(): lib.Array {
-        throw UnsupportedOperationException()
+        <selection><caret>throw UnsupportedOperationException()</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/traitGenericImplement.kt.after

@@ -4,6 +4,6 @@ trait G<T> {
 
 class GC<T>() : G<T> {
     override fun foo(t: T): T {
-        throw UnsupportedOperationException()
+        <selection><caret>throw UnsupportedOperationException()</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/traitNullableFunction.kt.after

@@ -4,6 +4,6 @@ trait Some {
 
 class SomeOther : Some {
     override fun foo(some: Int?): Int {
-        throw UnsupportedOperationException()
+        <selection><caret>throw UnsupportedOperationException()</selection>
     }
 }

idea/testData/codeInsight/overrideImplement/withLib/fakeOverride.kt.after

@@ -8,6 +8,6 @@ public open class B() : A() {
 
 public open class C() : B() {
     override fun foo() {
-        super<B>.foo()
+        <selection><caret>super<B>.foo()</selection>
     }
 }