KT-18158 Expand selection should select the comment after expression getter on the same line (#1122)

* Expand selection should select the comment after expression getter on the same line #KT-18158 Fixed

* Fixed expand selection behavior for the declaration with comment #KT-18158

* Remove redundant code #KT-18158

idea/idea-completion/testData/handlers/basic/override/AfterValKeyword.kt.after

@@ -4,7 +4,7 @@ interface I {
 
 class A : I {
     override val someVal: String?
-        get() = <caret><selection>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+        get() = <caret><selection>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>
 }
 
 // ELEMENT_TEXT: "override val someVal: String?"

idea/idea-completion/testData/handlers/basic/override/ImplementVal.kt.after

@@ -4,7 +4,7 @@ interface I {
 
 class A : I {
     override val someVal: String?
-        get() = <caret><selection>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+        get() = <caret><selection>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>
 }
 
 // ELEMENT_TEXT: "override val someVal: String?"

idea/idea-completion/testData/handlers/basic/override/ImplementVar.kt.after

@@ -4,7 +4,7 @@ interface I {
 
 class A : I {
     override var someVar: String
-        get() = <caret><selection>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+        get() = <caret><selection>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>
         set(value) {}
 }
 

idea/idea-core/src/org/jetbrains/kotlin/idea/core/generateUtil.kt

@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.idea.util.application.runWriteAction
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.allChildren
+import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.siblings
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.utils.ifEmpty
@@ -71,8 +72,7 @@ private fun moveCaretIntoGeneratedElementDocumentUnblocked(editor: Editor, eleme
     }
 
     if (element is KtDeclarationWithInitializer && element.hasInitializer()) {
-        val expression = element.initializer
+        val expression = element.initializer ?: throw AssertionError()
-        if (expression == null) throw AssertionError()
 
         val initializerRange = expression.textRange
 
@@ -81,7 +81,8 @@ private fun moveCaretIntoGeneratedElementDocumentUnblocked(editor: Editor, eleme
         editor.moveCaret(offset)
 
         if (initializerRange != null) {
-            editor.selectionModel.setSelection(initializerRange.startOffset, initializerRange.endOffset)
+            val endOffset = expression.siblings(forward = true, withItself = false).lastOrNull()?.endOffset ?: initializerRange.endOffset
+            editor.selectionModel.setSelection(initializerRange.startOffset, endOffset)
         }
 
         return true

idea/src/org/jetbrains/kotlin/idea/editor/wordSelection/KotlinDeclarationSelectioner.kt

@@ -16,18 +16,18 @@
 
 package org.jetbrains.kotlin.idea.editor.wordSelection
 
-import com.intellij.psi.PsiElement
+import com.intellij.codeInsight.editorActions.ExtendWordSelectionHandlerBase
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.util.TextRange
-import com.intellij.codeInsight.editorActions.ExtendWordSelectionHandlerBase
-import org.jetbrains.kotlin.psi.KtDeclaration
-import java.util.ArrayList
-import org.jetbrains.kotlin.psi.psiUtil.siblings
 import com.intellij.psi.PsiComment
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiWhiteSpace
+import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtDestructuringDeclaration
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import org.jetbrains.kotlin.psi.psiUtil.siblings
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
+import java.util.*
 
 class KotlinDeclarationSelectioner : ExtendWordSelectionHandlerBase() {
     override fun canSelect(e: PsiElement)
@@ -39,16 +39,33 @@ class KotlinDeclarationSelectioner : ExtendWordSelectionHandlerBase() {
         }
 
         val result = ArrayList<TextRange>()
+
         val firstChild = e.firstChild
+        val firstComment = firstChild
+                .siblings(forward = true, withItself = true)
+                .firstOrNull { it is PsiComment }
         val firstNonComment = firstChild
                 .siblings(forward = true, withItself = true)
                 .first { it !is PsiComment && it !is PsiWhiteSpace }
 
         val lastChild = e.lastChild
+        val lastComment = lastChild
+                .siblings(forward = false, withItself = true)
+                .firstOrNull { it is PsiComment }
         val lastNonComment = lastChild
                 .siblings(forward = false, withItself = true)
                 .first { it !is PsiComment && it !is PsiWhiteSpace }
 
+        if (firstComment != null || lastComment != null) {
+            val startOffset = minOf(
+                    firstComment?.startOffset ?: Int.MAX_VALUE,
+                    lastNonComment.startOffset)
+            val endOffset = maxOf(
+                    lastComment?.endOffset ?: Int.MIN_VALUE,
+                    lastNonComment.endOffset)
+            result.addRange(editorText, TextRange(startOffset, endOffset))
+        }
+
         if (firstNonComment != firstChild || lastNonComment != lastChild) {
             result.addRange(editorText, TextRange(firstNonComment.startOffset, lastNonComment.endOffset))
         }

idea/src/org/jetbrains/kotlin/idea/intentions/declarations/ConvertMemberToExtensionIntention.kt

@@ -22,6 +22,7 @@ import com.intellij.openapi.editor.ScrollType
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiMethodCallExpression
 import com.intellij.psi.PsiReferenceExpression
+import com.intellij.psi.PsiWhiteSpace
 import com.intellij.psi.codeStyle.JavaCodeStyleManager
 import com.intellij.psi.search.searches.ReferencesSearch
 import com.intellij.util.SmartList
@@ -36,6 +37,8 @@ import org.jetbrains.kotlin.idea.util.ImportInsertHelper
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
+import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import org.jetbrains.kotlin.psi.psiUtil.siblings
 import org.jetbrains.kotlin.utils.addIfNotNull
 
 class ConvertMemberToExtensionIntention : SelfTargetingRangeIntention<KtCallableDeclaration>(KtCallableDeclaration::class.java, "Convert member to extension"), LowPriorityAction {
@@ -60,9 +63,17 @@ class ConvertMemberToExtensionIntention : SelfTargetingRangeIntention<KtCallable
             unblockDocument()
 
             if (bodyToSelect != null) {
-                val range = bodyToSelect!!.textRange
+                val range = bodyToSelect.textRange
                 moveCaret(range.startOffset, ScrollType.CENTER)
-                selectionModel.setSelection(range.startOffset, range.endOffset)
+
+                val parent = bodyToSelect.parent
+                val lastSibling =
+                        if (parent is KtBlockExpression)
+                            parent.rBrace?.siblings(forward = false, withItself = false)?.first { it !is PsiWhiteSpace }
+                        else
+                            bodyToSelect.siblings(forward = true, withItself = false).lastOrNull()
+                val endOffset = lastSibling?.endOffset ?: range.endOffset
+                selectionModel.setSelection(range.startOffset, endOffset)
             }
             else {
                 moveCaret(extension.textOffset, ScrollType.CENTER)

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

@@ -6,7 +6,7 @@ interface A {
 fun some() : A {
     return object : A {
         override val method: () -> Unit?
-            get() = <selection><caret>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+            get() = <selection><caret>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>
     }
 }
 

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

@@ -4,5 +4,5 @@ interface A {
 
 class B : A {
     override val String.prop: Int
-        get() = <selection><caret>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+        get() = <selection><caret>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>
 }

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

@@ -4,6 +4,6 @@ interface A {
 
 class B : A {
     override var Int.foo: Double
-        get() = <selection><caret>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+        get() = <selection><caret>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>
         set(value) {}
 }

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

@@ -11,7 +11,7 @@ interface T {
 
 class C : T {
     override val a1: Byte
-        get() = <selection><caret>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+        get() = <selection><caret>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>
     override val a2: Short
         get() = TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.
     override val a3: Int

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

@@ -4,5 +4,5 @@ interface T {
 
 class GC() : T {
     override val v: Int
-        get() = <selection><caret>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+        get() = <selection><caret>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>
 }

idea/testData/intentions/declarations/convertMemberToExtension/abstract.kt.after

@@ -3,5 +3,5 @@ abstract class Owner {
 }
 
 fun Owner.f() {
-    <caret><selection>TODO("not implemented")</selection> //To change body of created functions use File | Settings | File Templates.
+    <caret><selection>TODO("not implemented") //To change body of created functions use File | Settings | File Templates.</selection>
 }

idea/testData/intentions/declarations/convertMemberToExtension/abstractProperty.kt.after

@@ -3,4 +3,4 @@ abstract class Owner {
 }
 
 val Owner.p: Int
-    get() = <caret><selection>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+    get() = <caret><selection>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>

idea/testData/intentions/declarations/convertMemberToExtension/valWithDefaultGetter.kt.after

@@ -4,4 +4,4 @@ class Owner {
 }
 
 val Owner.p: Int
-    get() = <caret><selection>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+    get() = <caret><selection>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>

idea/testData/intentions/declarations/convertMemberToExtension/varWithDefaultGetterAndSetter.kt.after

@@ -4,7 +4,7 @@ class Owner {
 }
 
 var Owner.p: Int
-    get() = <caret><selection>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+    get() = <caret><selection>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>
     set(value) {
         TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.
     }

idea/testData/intentions/declarations/convertMemberToExtension/varWithDefaultSetter.kt.after

@@ -6,5 +6,5 @@ class Owner {
 var Owner.p: Int
     get() = 1
     set(value) {
-        <caret><selection>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+        <caret><selection>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>
     }

idea/testData/intentions/declarations/convertMemberToExtension/varWithGetter.kt.after

@@ -6,5 +6,5 @@ class Owner {
 var Owner.p: Int
     get() { return 1 }
     set(value) {
-        <caret><selection>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+        <caret><selection>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>
     }

idea/testData/intentions/declarations/convertMemberToExtension/varWithSetter.kt.after

@@ -5,5 +5,5 @@ class Owner {
 }
 
 var Owner.p: Int
-    get() = <caret><selection>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+    get() = <caret><selection>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>
     set(v) {}

idea/testData/intentions/implementAbstractMember/property/enumClass.kt.after

@@ -8,5 +8,5 @@ enum class E : T<Int> {
     A, B, C;
 
     override val foo: Int
-        get() = <caret><selection>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+        get() = <caret><selection>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>
 }

idea/testData/intentions/implementAbstractMember/property/enumClassWithSemicolon.kt.after

@@ -8,5 +8,5 @@ enum class E : T<Int> {
     A, B, C;
 
     override val foo: Int
-        get() = <caret><selection>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+        get() = <caret><selection>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>
 }

idea/testData/intentions/implementAbstractMember/property/enumClassWithSemicolonAndMembers.kt.after

@@ -8,7 +8,7 @@ enum class E : T<Int> {
     A, B, C;
 
     override val foo: Int
-        get() = <caret><selection>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+        get() = <caret><selection>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>
 
     val bar = 1
 

idea/testData/intentions/implementAbstractMember/property/enumEntriesWithArgs.kt.after

@@ -9,7 +9,7 @@ enum class E(n: Int) {
             get() = TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.
     }, C(3) {
         override val foo: Int
-            get() = <caret><selection>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+            get() = <caret><selection>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>
     };
 
     abstract val foo: Int

idea/testData/intentions/implementAbstractMember/property/implementAll.kt.after

@@ -12,7 +12,7 @@ class U : T<String> {
 
 class V : T<Int> {
     override val foo: Int
-        get() = <caret><selection>TODO("not implemented")</selection> //To change initializer of created properties use File | Settings | File Templates.
+        get() = <caret><selection>TODO("not implemented") //To change initializer of created properties use File | Settings | File Templates.</selection>
 
 }
 

idea/testData/wordSelection/DeclarationWithComment1/0.kt

@@ -0,0 +1,2 @@
+val s1: String
+    get() = <caret>TODO("not implemented") // after

idea/testData/wordSelection/DeclarationWithComment1/1.kt

@@ -0,0 +1,2 @@
+val s1: String
+    get() = <selection><caret>TODO</selection>("not implemented") // after

idea/testData/wordSelection/DeclarationWithComment1/2.kt

@@ -0,0 +1,2 @@
+val s1: String
+    get() = <selection><caret>TODO("not implemented")</selection> // after

idea/testData/wordSelection/DeclarationWithComment1/3.kt

@@ -0,0 +1,2 @@
+val s1: String
+    get() = <selection><caret>TODO("not implemented") // after</selection>

idea/testData/wordSelection/DeclarationWithComment1/4.kt

@@ -0,0 +1,2 @@
+val s1: String
+    <selection>get() = <caret>TODO("not implemented") // after</selection>

idea/testData/wordSelection/DeclarationWithComment2/0.kt

@@ -0,0 +1,2 @@
+val s2: String
+    get() = /* before */ <caret>TODO("not implemented")

idea/testData/wordSelection/DeclarationWithComment2/1.kt

@@ -0,0 +1,2 @@
+val s2: String
+    get() = /* before */ <selection><caret>TODO</selection>("not implemented")

idea/testData/wordSelection/DeclarationWithComment2/2.kt

@@ -0,0 +1,2 @@
+val s2: String
+    get() = /* before */ <selection><caret>TODO("not implemented")</selection>

idea/testData/wordSelection/DeclarationWithComment2/3.kt

@@ -0,0 +1,2 @@
+val s2: String
+    get() = <selection>/* before */ <caret>TODO("not implemented")</selection>

idea/testData/wordSelection/DeclarationWithComment2/4.kt

@@ -0,0 +1,2 @@
+val s2: String
+    <selection>get() = /* before */ <caret>TODO("not implemented")</selection>

idea/testData/wordSelection/DeclarationWithComment3/0.kt

@@ -0,0 +1,2 @@
+val s3: String
+    get() = /* before */ <caret>TODO("not implemented") // after

idea/testData/wordSelection/DeclarationWithComment3/1.kt

@@ -0,0 +1,2 @@
+val s3: String
+    get() = /* before */ <selection><caret>TODO</selection>("not implemented") // after

idea/testData/wordSelection/DeclarationWithComment3/2.kt

@@ -0,0 +1,2 @@
+val s3: String
+    get() = /* before */ <selection><caret>TODO("not implemented")</selection> // after

idea/testData/wordSelection/DeclarationWithComment3/3.kt

@@ -0,0 +1,2 @@
+val s3: String
+    get() = <selection>/* before */ <caret>TODO("not implemented") // after</selection>

idea/testData/wordSelection/DeclarationWithComment3/4.kt

@@ -0,0 +1,2 @@
+val s3: String
+    <selection>get() = /* before */ <caret>TODO("not implemented") // after</selection>

idea/testData/wordSelection/DeclarationWithComment4/0.kt

@@ -0,0 +1,2 @@
+val s4: String
+    get() = /* before1 */ /* before2 */ <caret>TODO("not implemented") /* after1 */ /* after2 */

idea/testData/wordSelection/DeclarationWithComment4/1.kt

@@ -0,0 +1,2 @@
+val s4: String
+    get() = /* before1 */ /* before2 */ <selection><caret>TODO</selection>("not implemented") /* after1 */ /* after2 */

idea/testData/wordSelection/DeclarationWithComment4/2.kt

@@ -0,0 +1,2 @@
+val s4: String
+    get() = /* before1 */ /* before2 */ <selection><caret>TODO("not implemented")</selection> /* after1 */ /* after2 */

idea/testData/wordSelection/DeclarationWithComment4/3.kt

@@ -0,0 +1,2 @@
+val s4: String
+    get() = <selection>/* before1 */ /* before2 */ <caret>TODO("not implemented") /* after1 */ /* after2 */</selection>

idea/testData/wordSelection/DeclarationWithComment4/4.kt

@@ -0,0 +1,2 @@
+val s4: String
+    <selection>get() = /* before1 */ /* before2 */ <caret>TODO("not implemented") /* after1 */ /* after2 */</selection>

idea/tests/org/jetbrains/kotlin/idea/WordSelectionTest.java

@@ -112,6 +112,11 @@ public class WordSelectionTest extends KotlinLightCodeInsightFixtureTestCase {
         doTest();
     }
 
+    public void testDeclarationWithComment1() { doTest(); }
+    public void testDeclarationWithComment2() { doTest(); }
+    public void testDeclarationWithComment3() { doTest(); }
+    public void testDeclarationWithComment4() { doTest(); }
+
     private void doTest() {
         String dirName = getTestName(false);