Copy: Support multiple classes in the same file

idea/src/META-INF/plugin.xml

@@ -398,7 +398,7 @@
         language="kotlin" />
     <refactoring.copyHandler
         id="kotlinClass"
-        implementation="org.jetbrains.kotlin.idea.refactoring.copy.CopyKotlinClassHandler"
+        implementation="org.jetbrains.kotlin.idea.refactoring.copy.CopyKotlinClassesHandler"
         order="first" />
     <refactoring.copyHandler
         id="kotlinFile"

idea/src/org/jetbrains/kotlin/idea/refactoring/copy/CopyKotlinClassHandler.kt → idea/src/org/jetbrains/kotlin/idea/refactoring/copy/CopyKotlinClassesHandler.kt

@@ -25,16 +25,15 @@ import com.intellij.openapi.util.Key
 import com.intellij.psi.PsiDirectory
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
-import com.intellij.refactoring.MoveDestination
 import com.intellij.refactoring.RefactoringBundle
+import com.intellij.refactoring.copy.CopyFilesOrDirectoriesDialog
 import com.intellij.refactoring.copy.CopyHandlerDelegateBase
 import com.intellij.refactoring.rename.RenameProcessor
 import com.intellij.refactoring.util.MoveRenameUsageInfo
 import com.intellij.usageView.UsageInfo
 import com.intellij.util.IncorrectOperationException
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.idea.codeInsight.shorten.performDelayedRefactoringRequests
-import org.jetbrains.kotlin.idea.core.getPackage
 import org.jetbrains.kotlin.idea.core.quoteIfNeeded
 import org.jetbrains.kotlin.idea.refactoring.createKotlinFile
 import org.jetbrains.kotlin.idea.refactoring.move.*
@@ -47,25 +46,26 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.UserDataProperty
 import org.jetbrains.kotlin.psi.psiUtil.isAncestor
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
+import org.jetbrains.kotlin.utils.ifEmpty
 
-class CopyKotlinClassHandler : CopyHandlerDelegateBase() {
+class CopyKotlinClassesHandler : CopyHandlerDelegateBase() {
     companion object {
         @set:TestOnly
         var Project.newName: String? by UserDataProperty(Key.create("NEW_NAME"))
 
-        private fun PsiElement.getTopLevelClass(): KtClassOrObject? {
+        private fun PsiElement.getTopLevelClasses(): List<KtClassOrObject> {
             val classOrFile = parentsWithSelf.firstOrNull { it is KtFile || (it is KtClassOrObject && it.isTopLevel()) }
             return when (classOrFile) {
-                is KtFile -> classOrFile.declarations.singleOrNull() as? KtClassOrObject
-                is KtClassOrObject -> classOrFile
-                else -> null
+                is KtFile -> classOrFile.declarations.filterIsInstance<KtClassOrObject>()
+                is KtClassOrObject -> listOf(classOrFile)
+                else -> emptyList()
             }
         }
-
-        private fun getClassToCopy(elements: Array<out PsiElement>) = elements.singleOrNull()?.getTopLevelClass()
     }
 
-    override fun canCopy(elements: Array<out PsiElement>, fromUpdate: Boolean) = getClassToCopy(elements) != null
+    override fun canCopy(elements: Array<out PsiElement>, fromUpdate: Boolean): Boolean {
+        return elements.flatMap { it.getTopLevelClasses().ifEmpty { return false } }.distinctBy { it.containingFile }.size == 1
+    }
 
     enum class ExistingFilePolicy {
         APPEND, OVERWRITE, SKIP
@@ -135,9 +135,12 @@ class CopyKotlinClassHandler : CopyHandlerDelegateBase() {
     }
 
     override fun doCopy(elements: Array<out PsiElement>, defaultTargetDirectory: PsiDirectory?) {
-        val classToCopy = getClassToCopy(elements) ?: return
+        val classesToCopy = elements.flatMap { it.getTopLevelClasses() }
+        if (classesToCopy.isEmpty()) return
 
-        val originalFile = classToCopy.containingKtFile
+        val singleClassToCopy = classesToCopy.singleOrNull()
+
+        val originalFile = classesToCopy.first().containingKtFile
         val initialTargetDirectory = defaultTargetDirectory ?: originalFile.containingDirectory ?: return
 
         val project = initialTargetDirectory.project
@@ -147,61 +150,74 @@ class CopyKotlinClassHandler : CopyHandlerDelegateBase() {
         val commandName = RefactoringBundle.message("copy.handler.copy.class")
 
         var openInEditor = false
-        var newClassName: String? = classToCopy.name
-        var moveDestination: MoveDestination? = null
+        var newName: String? = singleClassToCopy?.name ?: originalFile.name
+        var targetDirWrapper: AutocreatingPsiDirectoryWrapper = initialTargetDirectory.toDirectoryWrapper()
 
         if (!ApplicationManager.getApplication().isUnitTestMode) {
-            val dialog = CopyKotlinClassDialog(classToCopy, initialTargetDirectory, project)
-            dialog.title = commandName
-            if (!dialog.showAndGet()) return
-
-            openInEditor = dialog.openInEditor
-            newClassName = dialog.className
-            moveDestination = dialog.targetDirectory ?: return
+            if (singleClassToCopy != null) {
+                val dialog = CopyKotlinClassDialog(singleClassToCopy, initialTargetDirectory, project)
+                dialog.title = commandName
+                if (!dialog.showAndGet()) return
+
+                openInEditor = dialog.openInEditor
+                newName = dialog.className ?: singleClassToCopy.name
+                targetDirWrapper = dialog.targetDirectory?.toDirectoryWrapper() ?: return
+            }
+            else {
+                val dialog = CopyFilesOrDirectoriesDialog(arrayOf(originalFile), initialTargetDirectory, project, false)
+                if (!dialog.showAndGet()) return
+                openInEditor = dialog.openInEditor()
+                newName = dialog.newName
+                targetDirWrapper = dialog.targetDirectory?.toDirectoryWrapper() ?: return
+            }
         }
         else {
-            project.newName?.let { newClassName = it }
+            project.newName?.let { newName = it }
         }
 
-        if (newClassName.isNullOrEmpty()) return
+        if (singleClassToCopy != null && newName.isNullOrEmpty()) return
 
         val internalUsages = runReadAction {
-            val targetPackageName = moveDestination?.targetPackage?.qualifiedName
-                                    ?: initialTargetDirectory.getPackage()?.qualifiedName
-                                    ?: ""
+            val targetPackageName = targetDirWrapper.getPackageName()
             val changeInfo = ContainerChangeInfo(
                     ContainerInfo.Package(originalFile.packageFqName),
                     ContainerInfo.Package(FqName(targetPackageName))
             )
-            classToCopy
-                    .getInternalReferencesToUpdateOnPackageNameChange(changeInfo)
-                    .filter {
-                        val referencedElement = (it as? MoveRenameUsageInfo)?.referencedElement
-                        referencedElement == null || !classToCopy.isAncestor(referencedElement)
-                    }
+            classesToCopy.flatMap { classToCopy ->
+                classToCopy.getInternalReferencesToUpdateOnPackageNameChange(changeInfo).filter {
+                    val referencedElement = (it as? MoveRenameUsageInfo)?.referencedElement
+                    referencedElement == null || !classToCopy.isAncestor(referencedElement)
+                }
+            }
         }
         markInternalUsages(internalUsages)
 
         val restoredInternalUsages = ArrayList<UsageInfo>()
 
         project.executeCommand(commandName) {
             try {
-                val targetDirectory = runWriteAction { moveDestination?.getTargetDirectory(initialTargetDirectory) ?: initialTargetDirectory }
-                val targetFileName = newClassName + "." + originalFile.virtualFile.extension
+                val targetDirectory = runWriteAction { targetDirWrapper.getOrCreateDirectory(initialTargetDirectory) }
+                val targetFileName = if (newName?.contains(".") ?: false) newName!! else newName + "." + originalFile.virtualFile.extension
 
                 val targetFile = getOrCreateTargetFile(originalFile, targetDirectory, targetFileName, commandName) ?: return@executeCommand
 
-                val newClass = runWriteAction {
-                    val newClass = targetFile.add(classToCopy.copy()) as KtClassOrObject
-                    val oldToNewElementsMapping: Map<PsiElement, PsiElement> = mapOf(classToCopy to newClass)
-                    restoredInternalUsages += restoreInternalUsages(newClass, oldToNewElementsMapping, true)
-                    postProcessMoveUsages(restoredInternalUsages, oldToNewElementsMapping)
+                val newClasses = runWriteAction {
+                    val newClasses = classesToCopy.map { targetFile.add(it.copy()) as KtClassOrObject }
+                    val oldToNewElementsMapping = classesToCopy.zip(newClasses).toMap<PsiElement, PsiElement>()
+
+                    for (newClass in newClasses) {
+                        restoredInternalUsages += restoreInternalUsages(newClass, oldToNewElementsMapping, true)
+                        postProcessMoveUsages(restoredInternalUsages, oldToNewElementsMapping)
+                    }
+
                     performDelayedRefactoringRequests(project)
 
-                    newClass
+                    newClasses
                 }
 
-                RenameProcessor(project, newClass, newClassName!!.quoteIfNeeded(), false, false).run()
+                newClasses.singleOrNull()?.let {
+                    RenameProcessor(project, it, newName!!.quoteIfNeeded(), false, false).run()
+                }
 
                 if (openInEditor) {
                     EditorHelper.openFilesInEditor(arrayOf(targetFile))

idea/src/org/jetbrains/kotlin/idea/refactoring/move/AutocreatingPsiDirectoryWrapper.kt

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2010-2017 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.kotlin.idea.refactoring.move
+
+import com.intellij.psi.PsiDirectory
+import com.intellij.refactoring.MoveDestination
+import org.jetbrains.kotlin.idea.core.getPackage
+
+sealed class AutocreatingPsiDirectoryWrapper {
+    class ByPsiDirectory(val psiDirectory: PsiDirectory) : AutocreatingPsiDirectoryWrapper() {
+        override fun getPackageName(): String = psiDirectory.getPackage()?.qualifiedName ?: ""
+        override fun getOrCreateDirectory(source: PsiDirectory) = psiDirectory
+    }
+
+    class ByMoveDestination(val moveDestination: MoveDestination) : AutocreatingPsiDirectoryWrapper() {
+        override fun getPackageName() = moveDestination.targetPackage.qualifiedName
+        override fun getOrCreateDirectory(source: PsiDirectory) = moveDestination.getTargetDirectory(source)
+    }
+
+    abstract fun getPackageName(): String
+    abstract fun getOrCreateDirectory(source: PsiDirectory): PsiDirectory
+}
+
+fun MoveDestination.toDirectoryWrapper() = AutocreatingPsiDirectoryWrapper.ByMoveDestination(this)
+fun PsiDirectory.toDirectoryWrapper() = AutocreatingPsiDirectoryWrapper.ByPsiDirectory(this)

idea/testData/refactoring/copy/copyMultiClassFile/after/bar/test.kt

@@ -0,0 +1,11 @@
+package bar
+
+class A {
+    val a: A = A()
+    val b: B = B()
+}
+
+class B {
+    val a: A = A()
+    val b: B = B()
+}

idea/testData/refactoring/copy/copyMultiClassFile/after/foo/test.kt

@@ -0,0 +1,11 @@
+package foo
+
+class A {
+    val a: A = A()
+    val b: B = B()
+}
+
+class B {
+    val a: A = A()
+    val b: B = B()
+}

idea/testData/refactoring/copy/copyMultiClassFile/before/foo/test.kt

@@ -0,0 +1,11 @@
+package foo
+
+class A {
+    val a: A = A()
+    val b: B = B()
+}
+
+class B {
+    val a: A = A()
+    val b: B = B()
+}

idea/testData/refactoring/copy/copyMultiClassFile/copyMultiClassFile.test

@@ -0,0 +1,4 @@
+{
+  "mainFile": "foo/test.kt",
+  "targetPackage": "bar"
+}

idea/testData/refactoring/copy/copyMultipleClassesToExistingFile/after/bar/test.kt

@@ -0,0 +1,19 @@
+package bar
+
+import foo.B
+
+fun test() {
+
+}
+
+class A {
+    val a: A = A()
+    val b: B = B()
+    val c: C = C()
+}
+
+class C {
+    val a: A = A()
+    val b: B = B()
+    val c: C = C()
+}

idea/testData/refactoring/copy/copyMultipleClassesToExistingFile/after/foo/test.kt

@@ -0,0 +1,19 @@
+package foo
+
+class A {
+    val a: A = A()
+    val b: B = B()
+    val c: C = C()
+}
+
+class B {
+    val a: A = A()
+    val b: B = B()
+    val c: C = C()
+}
+
+class C {
+    val a: A = A()
+    val b: B = B()
+    val c: C = C()
+}

idea/testData/refactoring/copy/copyMultipleClassesToExistingFile/before/bar/test.kt

@@ -0,0 +1,5 @@
+package bar
+
+fun test() {
+
+}

idea/testData/refactoring/copy/copyMultipleClassesToExistingFile/before/foo/test.kt

@@ -0,0 +1,19 @@
+package foo
+
+class <caret>A {
+    val a: A = A()
+    val b: B = B()
+    val c: C = C()
+}
+
+class B {
+    val a: A = A()
+    val b: B = B()
+    val c: C = C()
+}
+
+class <caret>C {
+    val a: A = A()
+    val b: B = B()
+    val c: C = C()
+}

idea/testData/refactoring/copy/copyMultipleClassesToExistingFile/copyMultipleClassesToExistingFile.test

@@ -0,0 +1,4 @@
+{
+  "mainFile": "foo/test.kt",
+  "targetPackage": "bar"
+}

idea/testData/refactoring/copy/copyMultipleClassesToNewFile/after/bar/test.kt

@@ -0,0 +1,15 @@
+package bar
+
+import foo.B
+
+class A {
+    val a: A = A()
+    val b: B = B()
+    val c: C = C()
+}
+
+class C {
+    val a: A = A()
+    val b: B = B()
+    val c: C = C()
+}

idea/testData/refactoring/copy/copyMultipleClassesToNewFile/after/foo/test.kt

@@ -0,0 +1,19 @@
+package foo
+
+class A {
+    val a: A = A()
+    val b: B = B()
+    val c: C = C()
+}
+
+class B {
+    val a: A = A()
+    val b: B = B()
+    val c: C = C()
+}
+
+class C {
+    val a: A = A()
+    val b: B = B()
+    val c: C = C()
+}

idea/testData/refactoring/copy/copyMultipleClassesToNewFile/before/foo/test.kt

@@ -0,0 +1,19 @@
+package foo
+
+class <caret>A {
+    val a: A = A()
+    val b: B = B()
+    val c: C = C()
+}
+
+class B {
+    val a: A = A()
+    val b: B = B()
+    val c: C = C()
+}
+
+class <caret>C {
+    val a: A = A()
+    val b: B = B()
+    val c: C = C()
+}