Create header implementation fix: check import directive, change / add it if needed

idea/src/org/jetbrains/kotlin/idea/quickfix/createImpl/CreateHeaderImplementationFix.kt

@@ -30,13 +30,11 @@ import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.diagnostics.DiagnosticFactory
 import org.jetbrains.kotlin.diagnostics.Errors
-import org.jetbrains.kotlin.idea.core.ShortenReferences
+import org.jetbrains.kotlin.idea.core.*
-import org.jetbrains.kotlin.idea.core.TemplateKind
-import org.jetbrains.kotlin.idea.core.getFunctionBodyTextFromTemplate
-import org.jetbrains.kotlin.idea.core.toDescriptor
 import org.jetbrains.kotlin.idea.project.TargetPlatformDetector
 import org.jetbrains.kotlin.idea.quickfix.KotlinQuickFixAction
 import org.jetbrains.kotlin.idea.quickfix.KotlinSingleIntentionActionFactory
+import org.jetbrains.kotlin.idea.refactoring.createKotlinFile
 import org.jetbrains.kotlin.idea.refactoring.getOrCreateKotlinFile
 import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
 import org.jetbrains.kotlin.idea.util.application.runWriteAction
@@ -66,6 +64,20 @@ sealed class CreateHeaderImplementationFix<out D : KtNamedDeclaration>(
         val generated = factory.generateIt(project, element) ?: return
 
         runWriteAction {
+            if (implFile.packageDirective?.fqName != file.packageDirective?.fqName &&
+                implFile.declarations.isEmpty()) {
+                val packageDirective = file.packageDirective
+                packageDirective?.let {
+                    val oldPackageDirective = implFile.packageDirective
+                    val newPackageDirective = factory.createPackageDirective(it.fqName)
+                    if (oldPackageDirective != null) {
+                        oldPackageDirective.replace(newPackageDirective)
+                    }
+                    else {
+                        implFile.add(newPackageDirective)
+                    }
+                }
+            }
             val implDeclaration = implFile.add(generated) as KtElement
             val reformatted = CodeStyleManager.getInstance(project).reformat(implDeclaration)
             ShortenReferences.DEFAULT.process(reformatted as KtElement)
@@ -93,7 +105,28 @@ sealed class CreateHeaderImplementationFix<out D : KtNamedDeclaration>(
                 implModule, headerPackage?.qualifiedName ?: "", null, false
         ) ?: return null
         return runWriteAction {
-            getOrCreateKotlinFile("$name.kt", implDirectory)
+            val fileName = "$name.kt"
+            val existingFile = implDirectory.findFile(fileName)
+            val packageDirective = declaration.containingKtFile.packageDirective
+            val packageName =
+                    if (packageDirective?.packageNameExpression == null) implDirectory.getPackage()?.qualifiedName
+                    else packageDirective.fqName.asString()
+            if (existingFile is KtFile) {
+                val existingPackageDirective = existingFile.packageDirective
+                if (existingFile.declarations.isNotEmpty() &&
+                    existingPackageDirective?.fqName != packageDirective?.fqName) {
+                    val newName = KotlinNameSuggester.suggestNameByName(name) {
+                        implDirectory.findFile("$it.kt") == null
+                    } + ".kt"
+                    createKotlinFile(newName, implDirectory, packageName)
+                }
+                else {
+                    existingFile
+                }
+            }
+            else {
+                createKotlinFile(fileName, implDirectory, packageName)
+            }
         }
     }
 

idea/testData/multiModuleQuickFix/package/header/header.kt

@@ -0,0 +1,6 @@
+// "Create header class implementation for platform JVM" "true"
+// SHOULD_BE_AVAILABLE_AFTER_EXECUTION
+
+package test.inner
+
+header class <caret>My

idea/testData/multiModuleQuickFix/package/header/header.kt.after

@@ -0,0 +1,6 @@
+// "Create header class implementation for platform JVM" "true"
+// SHOULD_BE_AVAILABLE_AFTER_EXECUTION
+
+package test.inner
+
+header class My

idea/testData/multiModuleQuickFix/package/jvm/My.kt.after

@@ -0,0 +1,3 @@
+package test.inner
+
+impl class My

idea/testData/multiModuleQuickFix/packageIncorrect/header/header.kt

@@ -0,0 +1,6 @@
+// "Create header class implementation for platform JVM" "true"
+// SHOULD_BE_AVAILABLE_AFTER_EXECUTION
+
+package test.inner
+
+header class <caret>My

idea/testData/multiModuleQuickFix/packageIncorrect/header/header.kt.after

@@ -0,0 +1,6 @@
+// "Create header class implementation for platform JVM" "true"
+// SHOULD_BE_AVAILABLE_AFTER_EXECUTION
+
+package test.inner
+
+header class My

idea/testData/multiModuleQuickFix/packageIncorrect/jvm/My.kt

@@ -0,0 +1,3 @@
+package incorrect
+
+class Some

idea/testData/multiModuleQuickFix/packageIncorrect/jvm/My.kt.after

@@ -0,0 +1,3 @@
+package incorrect
+
+class Some

idea/testData/multiModuleQuickFix/packageIncorrect/jvm/My1.kt.after

@@ -0,0 +1,3 @@
+package test.inner
+
+impl class My

idea/testData/multiModuleQuickFix/packageIncorrectEmpty/header/header.kt

@@ -0,0 +1,6 @@
+// "Create header class implementation for platform JVM" "true"
+// SHOULD_BE_AVAILABLE_AFTER_EXECUTION
+
+package test.inner
+
+header class <caret>My

idea/testData/multiModuleQuickFix/packageIncorrectEmpty/header/header.kt.after

@@ -0,0 +1,6 @@
+// "Create header class implementation for platform JVM" "true"
+// SHOULD_BE_AVAILABLE_AFTER_EXECUTION
+
+package test.inner
+
+header class My

idea/testData/multiModuleQuickFix/packageIncorrectEmpty/jvm/My.kt

@@ -0,0 +1 @@
+package incorrect

idea/testData/multiModuleQuickFix/packageIncorrectEmpty/jvm/My.kt.after

@@ -0,0 +1,3 @@
+package test.inner
+
+impl class My

idea/tests/org/jetbrains/kotlin/idea/quickfix/AbstractQuickFixMultiModuleTest.kt

@@ -26,10 +26,12 @@ import junit.framework.ComparisonFailure
 import junit.framework.TestCase
 import org.jetbrains.kotlin.idea.inspections.findExistingEditor
 import org.jetbrains.kotlin.idea.project.PluginJetFilesProvider
+import org.jetbrains.kotlin.idea.refactoring.createKotlinFile
 import org.jetbrains.kotlin.idea.stubs.AbstractMultiModuleTest
 import org.jetbrains.kotlin.idea.test.DirectiveBasedActionUtils
 import org.jetbrains.kotlin.idea.test.PluginTestCaseBase
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.test.InTextDirectivesUtils
 import org.jetbrains.kotlin.test.KotlinTestUtils
 import java.io.File
 import java.util.regex.Pattern
@@ -40,7 +42,8 @@ abstract class AbstractQuickFixMultiModuleTest : AbstractMultiModuleTest() {
 
     override fun getTestDataPath() = testPath
 
-    protected fun shouldBeAvailableAfterExecution() = false
+    protected fun shouldBeAvailableAfterExecution(file: KtFile) =
+            InTextDirectivesUtils.isDirectiveDefined(file.text, "// SHOULD_BE_AVAILABLE_AFTER_EXECUTION")
 
     private fun getActionsTexts(availableActions: List<IntentionAction>) = availableActions.map { it.text }
 
@@ -68,7 +71,7 @@ abstract class AbstractQuickFixMultiModuleTest : AbstractMultiModuleTest() {
                     DirectiveBasedActionUtils.checkForUnexpectedErrors(psiFile)
                 }
 
-                doAction(text, actionShouldBeAvailable, actionFileName)
+                doAction(text, actionShouldBeAvailable, actionFileName, actionFile)
 
                 if (actionShouldBeAvailable) {
                     val testDirectory = File(testPath)
@@ -79,7 +82,9 @@ abstract class AbstractQuickFixMultiModuleTest : AbstractMultiModuleTest() {
                             try {
                                 val editedFile = allFilesInProject.find {
                                     it.name.toLowerCase() == file.name.removeSuffix(".after").toLowerCase()
-                                }!!
+                                } ?: allFilesInProject.mapNotNull {
+                                    it.containingDirectory?.findFile(file.name.removeSuffix(".after"))
+                                }.single()
                                 setActiveEditor(editedFile.findExistingEditor() ?: createEditor(editedFile.virtualFile))
                                 checkResultByFile(file.relativeTo(testDirectory).path)
                             }
@@ -105,7 +110,7 @@ abstract class AbstractQuickFixMultiModuleTest : AbstractMultiModuleTest() {
 
     // TODO: merge with AbstractQuickFixMultiFileTest
 
-    fun doAction(text: String, actionShouldBeAvailable: Boolean, testFilePath: String) {
+    fun doAction(text: String, actionShouldBeAvailable: Boolean, testFilePath: String, actionFile: KtFile) {
         val pattern = if (text.startsWith("/"))
             Pattern.compile(text.substring(1, text.length - 1))
         else
@@ -140,7 +145,7 @@ abstract class AbstractQuickFixMultiModuleTest : AbstractMultiModuleTest() {
             UIUtil.dispatchAllInvocationEvents()
 
 
-            if (!shouldBeAvailableAfterExecution()) {
+            if (!shouldBeAvailableAfterExecution(actionFile)) {
                 val afterAction = findActionByPattern(pattern, getAvailableActions())
 
                 if (afterAction != null) {

idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixMultiModuleTest.kt

@@ -65,6 +65,21 @@ class QuickFixMultiModuleTest : AbstractQuickFixMultiModuleTest() {
         doMultiPlatformTest()
     }
 
+    @Test
+    fun testPackage() {
+        doMultiPlatformTest()
+    }
+
+    @Test
+    fun testPackageIncorrect() {
+        doMultiPlatformTest()
+    }
+
+    @Test
+    fun testPackageIncorrectEmpty() {
+        doMultiPlatformTest()
+    }
+
     @Test
     fun testNested() {
         doMultiPlatformTest()