Override/Implement members: place members in the same order as super class members

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

@@ -218,7 +218,8 @@ fun <T : KtDeclaration> insertMembersAfter(
     editor: Editor?,
     classOrObject: KtClassOrObject,
     members: Collection<T>,
-    anchor: PsiElement? = null
+    anchor: PsiElement? = null,
+    getAnchor: (KtDeclaration) -> PsiElement? = { null }
 ): List<T> {
     members.ifEmpty { return emptyList() }
 
@@ -239,6 +240,8 @@ fun <T : KtDeclaration> insertMembersAfter(
 
             var afterAnchor = anchor ?: findInsertAfterAnchor(editor, body) ?: return@runWriteAction emptyList<T>()
             otherMembers.mapTo(insertedMembers) {
+                afterAnchor = getAnchor(it) ?: afterAnchor
+
                 if (classOrObject is KtClass && classOrObject.isEnum()) {
                     val enumEntries = classOrObject.declarations.filterIsInstance<KtEnumEntry>()
                     val bound = (enumEntries.lastOrNull() ?: classOrObject.allChildren.firstOrNull { element ->

idea/idea-core/src/org/jetbrains/kotlin/idea/core/overrideImplement/OverrideImplementMembersHandler.kt

@@ -25,13 +25,24 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.ui.DialogWrapper
 import com.intellij.psi.PsiDocumentManager
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
+import com.intellij.psi.util.PsiTreeUtil
+import org.jetbrains.kotlin.descriptors.CallableMemberDescriptor
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptorIfAny
 import org.jetbrains.kotlin.idea.core.insertMembersAfter
+import org.jetbrains.kotlin.idea.search.usagesSearch.descriptor
 import org.jetbrains.kotlin.psi.KtClassOrObject
+import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.prevSiblingOfSameType
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
+import org.jetbrains.kotlin.resolve.DescriptorUtils
+import org.jetbrains.kotlin.resolve.source.getPsi
+import org.jetbrains.kotlin.util.findCallableMemberBySignature
+import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
 
 abstract class OverrideImplementMembersHandler : LanguageCodeInsightActionHandler {
 
@@ -42,7 +53,10 @@ abstract class OverrideImplementMembersHandler : LanguageCodeInsightActionHandle
 
     protected abstract fun collectMembersToGenerate(descriptor: ClassDescriptor, project: Project): Collection<OverrideMemberChooserObject>
 
-    private fun showOverrideImplementChooser(project: Project, members: Array<OverrideMemberChooserObject>): MemberChooser<OverrideMemberChooserObject>? {
+    private fun showOverrideImplementChooser(
+        project: Project,
+        members: Array<OverrideMemberChooserObject>
+    ): MemberChooser<OverrideMemberChooserObject>? {
         val chooser = MemberChooser(members, true, true, project)
         chooser.title = getChooserTitle()
         chooser.show()
@@ -80,8 +94,7 @@ abstract class OverrideImplementMembersHandler : LanguageCodeInsightActionHandle
         if (implementAll) {
             selectedElements = members
             copyDoc = false
-        }
+        } else {
-        else {
             val chooser = showOverrideImplementChooser(project, members.toTypedArray()) ?: return
             selectedElements = chooser.selectedElements ?: return
             copyDoc = chooser.isCopyJavadoc
@@ -101,12 +114,65 @@ abstract class OverrideImplementMembersHandler : LanguageCodeInsightActionHandle
 
     companion object {
         fun generateMembers(
-                editor: Editor?,
+            editor: Editor?,
-                classOrObject: KtClassOrObject,
+            classOrObject: KtClassOrObject,
-                selectedElements: Collection<OverrideMemberChooserObject>,
+            selectedElements: Collection<OverrideMemberChooserObject>,
-                copyDoc: Boolean
+            copyDoc: Boolean
         ) {
-            insertMembersAfter(editor, classOrObject, selectedElements.map { it.generateMember(classOrObject, copyDoc) })
+            val selectedMemberDescriptors = selectedElements.associate { it.generateMember(classOrObject, copyDoc) to it.descriptor }
+
+            val classBody = classOrObject.body
+            if (classBody == null) {
+                insertMembersAfter(editor, classOrObject, selectedMemberDescriptors.keys)
+                return
+            }
+            val offset = editor?.caretModel?.offset ?: classBody.startOffset
+            val offsetCursorElement = PsiTreeUtil.findFirstParent(classBody.containingFile.findElementAt(offset)) {
+                it.parent == classBody
+            }
+            if (offsetCursorElement != null && offsetCursorElement != classBody.rBrace) {
+                insertMembersAfter(editor, classOrObject, selectedMemberDescriptors.keys)
+                return
+            }
+            val classLeftBrace = classBody.lBrace
+
+            val allSuperMemberDescriptors = selectedMemberDescriptors.values
+                .mapNotNull { it.containingDeclaration as? ClassDescriptor }
+                .associateWith {
+                    DescriptorUtils.getAllDescriptors(it.unsubstitutedMemberScope).filterIsInstance<CallableMemberDescriptor>()
+                }
+
+            val implementedElements = mutableMapOf<CallableMemberDescriptor, KtDeclaration>()
+
+            fun ClassDescriptor.findElement(memberDescriptor: CallableMemberDescriptor): KtDeclaration? {
+                return implementedElements[memberDescriptor]
+                    ?: (findCallableMemberBySignature(memberDescriptor)?.source?.getPsi() as? KtDeclaration)?.also {
+                        implementedElements[memberDescriptor] = it
+                    }
+            }
+
+            fun getAnchor(selectedElement: KtDeclaration): PsiElement? {
+                val lastElement = classOrObject.declarations.lastOrNull()
+                val selectedMemberDescriptor = selectedMemberDescriptors[selectedElement] ?: return lastElement
+                val superMemberDescriptors = allSuperMemberDescriptors[selectedMemberDescriptor.containingDeclaration] ?: return lastElement
+                val index = superMemberDescriptors.indexOf(selectedMemberDescriptor)
+                if (index == -1) return lastElement
+                val classDescriptor = classOrObject.descriptor as? ClassDescriptor ?: return lastElement
+
+                val upperElement = ((index - 1) downTo 0).firstNotNullResult {
+                    classDescriptor.findElement(superMemberDescriptors[it])
+                }
+                if (upperElement != null) return upperElement
+
+                val lowerElement = ((index + 1) until superMemberDescriptors.size).firstNotNullResult {
+                    classDescriptor.findElement(superMemberDescriptors[it])
+                }
+                if (lowerElement != null) return lowerElement.prevSiblingOfSameType() ?: classLeftBrace
+
+                return lastElement
+            }
+
+            insertMembersAfter(editor, classOrObject, selectedMemberDescriptors.keys) { getAnchor(it) }
         }
     }
 }

idea/testData/codeInsight/overrideImplement/implementFromClassName.kt

@@ -0,0 +1,8 @@
+interface Runnable {
+    fun run()
+}
+
+class <caret>A : Runnable {
+    fun foo() {
+    }
+}

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

@@ -0,0 +1,12 @@
+interface Runnable {
+    fun run()
+}
+
+class A : Runnable {
+    fun foo() {
+    }
+
+    override fun run() {
+        <selection><caret>TODO("not implemented") //To change body of created functions use File | Settings | File Templates.</selection>
+    }
+}

idea/testData/codeInsight/overrideImplement/implementFromClassName2.kt

@@ -0,0 +1,18 @@
+interface A {
+    fun a1()
+    fun a2()
+    fun a3()
+    fun a4()
+    fun a5()
+    fun a6()
+    fun a7()
+    fun a8()
+}
+
+class <caret>Test : A {
+    override fun a3() {
+    }
+
+    override fun a6() {
+    }
+}

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

@@ -0,0 +1,42 @@
+interface A {
+    fun a1()
+    fun a2()
+    fun a3()
+    fun a4()
+    fun a5()
+    fun a6()
+    fun a7()
+    fun a8()
+}
+
+class Test : A {
+    override fun a1() {
+        <selection><caret>TODO("not implemented") //To change body of created functions use File | Settings | File Templates.</selection>
+    }
+
+    override fun a2() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun a3() {
+    }
+
+    override fun a4() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun a5() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun a6() {
+    }
+
+    override fun a7() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun a8() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+}

idea/testData/codeInsight/overrideImplement/implementFromClassName3.kt

@@ -0,0 +1,18 @@
+interface A {
+    fun a1()
+    fun a2()
+    fun a3()
+    fun a4()
+    fun a5()
+    fun a6()
+    fun a7()
+    fun a8()
+}
+
+class <caret>Test : A {
+    override fun a6() {
+    }
+
+    override fun a3() {
+    }
+}

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

@@ -0,0 +1,42 @@
+interface A {
+    fun a1()
+    fun a2()
+    fun a3()
+    fun a4()
+    fun a5()
+    fun a6()
+    fun a7()
+    fun a8()
+}
+
+class Test : A {
+    override fun a6() {
+    }
+
+    override fun a7() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun a8() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun a1() {
+        <selection><caret>TODO("not implemented") //To change body of created functions use File | Settings | File Templates.</selection>
+    }
+
+    override fun a2() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun a3() {
+    }
+
+    override fun a4() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun a5() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+}

idea/testData/codeInsight/overrideImplement/implementFromClassName4.kt

@@ -0,0 +1,21 @@
+interface A {
+    fun a1()
+    fun a2()
+    fun a3()
+    fun a4()
+    fun a5()
+    fun a6()
+    fun a7()
+    fun a8()
+}
+
+class <caret>Test : A {
+    override fun a4() {
+    }
+
+    override fun a6() {
+    }
+
+    override fun a2() {
+    }
+}

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

@@ -0,0 +1,41 @@
+interface A {
+    fun a1()
+    fun a2()
+    fun a3()
+    fun a4()
+    fun a5()
+    fun a6()
+    fun a7()
+    fun a8()
+}
+
+class Test : A {
+    override fun a4() {
+    }
+
+    override fun a5() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun a6() {
+    }
+
+    override fun a7() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun a8() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun a1() {
+        <selection><caret>TODO("not implemented") //To change body of created functions use File | Settings | File Templates.</selection>
+    }
+
+    override fun a2() {
+    }
+
+    override fun a3() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+}

idea/testData/codeInsight/overrideImplement/implementFromClassName5.kt

@@ -0,0 +1,27 @@
+interface A {
+    fun a1()
+    fun a2()
+    fun a3()
+    fun a4()
+    fun a5()
+    fun a6()
+    fun a7()
+    fun a8()
+}
+
+class <caret>Test : A {
+    fun foo() {
+    }
+
+    override fun a3() {
+    }
+
+    fun bar() {
+    }
+
+    override fun a6() {
+    }
+
+    fun baz() {
+    }
+}

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

@@ -0,0 +1,51 @@
+interface A {
+    fun a1()
+    fun a2()
+    fun a3()
+    fun a4()
+    fun a5()
+    fun a6()
+    fun a7()
+    fun a8()
+}
+
+class Test : A {
+    fun foo() {
+    }
+
+    override fun a1() {
+        <selection><caret>TODO("not implemented") //To change body of created functions use File | Settings | File Templates.</selection>
+    }
+
+    override fun a2() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun a3() {
+    }
+
+    override fun a4() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun a5() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    fun bar() {
+    }
+
+    override fun a6() {
+    }
+
+    override fun a7() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun a8() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    fun baz() {
+    }
+}

idea/testData/codeInsight/overrideImplement/implementFromClassName6.kt

@@ -0,0 +1,27 @@
+interface A {
+    fun a1()
+    fun a2()
+    fun a3()
+    fun a4()
+    fun a5()
+    fun a6()
+    fun a7()
+    fun a8()
+}
+
+interface B {
+    fun b3()
+    fun b2()
+    fun b1()
+}
+
+class <caret>Test : A, B {
+    override fun a3() {
+    }
+
+    override fun b2() {
+    }
+
+    override fun a6() {
+    }
+}

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

@@ -0,0 +1,59 @@
+interface A {
+    fun a1()
+    fun a2()
+    fun a3()
+    fun a4()
+    fun a5()
+    fun a6()
+    fun a7()
+    fun a8()
+}
+
+interface B {
+    fun b3()
+    fun b2()
+    fun b1()
+}
+
+class Test : A, B {
+    override fun a1() {
+        <selection><caret>TODO("not implemented") //To change body of created functions use File | Settings | File Templates.</selection>
+    }
+
+    override fun a2() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun a3() {
+    }
+
+    override fun a4() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun a5() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun b3() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun b2() {
+    }
+
+    override fun b1() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun a6() {
+    }
+
+    override fun a7() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+
+    override fun a8() {
+        TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+    }
+}

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

@@ -3,6 +3,13 @@ interface Some {
 }
 
 class Other {
+    fun test() {
+        val a = 1
+    }
+    fun otherTest() {
+
+    }
+
     override fun equals(other: Any?): Boolean {
         <selection><caret>return super.equals(other)</selection>
     }
@@ -14,12 +21,5 @@ class Other {
     override fun toString(): String {
         return super.toString()
     }
-
-    fun test() {
-        val a = 1
-    }
-    fun otherTest() {
-
-    }
 }
 

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

@@ -7,11 +7,10 @@ interface T<X> {
 enum class E : T<Int> {
     A, B, C;
 
-    override fun foo(x: Int): Int {
-        <caret><selection>TODO("not implemented") //To change body of created functions use File | Settings | File Templates.</selection>
-    }
-
     val bar = 1
 
     fun baz() = 2
+    override fun foo(x: Int): Int {
+        <caret><selection>TODO("not implemented") //To change body of created functions use File | Settings | File Templates.</selection>
+    }
 }

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

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

idea/tests/org/jetbrains/kotlin/idea/codeInsight/OverrideImplementTest.kt

@@ -289,4 +289,28 @@ class OverrideImplementTest : AbstractOverrideImplementTest() {
     fun testUnresolvedType() {
         doOverrideFileTest()
     }
+
+    fun testImplementFromClassName() {
+        doMultiImplementFileTest()
+    }
+
+    fun testImplementFromClassName2() {
+        doMultiImplementFileTest()
+    }
+
+    fun testImplementFromClassName3() {
+        doMultiImplementFileTest()
+    }
+
+    fun testImplementFromClassName4() {
+        doMultiImplementFileTest()
+    }
+
+    fun testImplementFromClassName5() {
+        doMultiImplementFileTest()
+    }
+
+    fun testImplementFromClassName6() {
+        doMultiImplementFileTest()
+    }
 }