Create actual fix: handle companions correctly #KT-21114 Fixed

compiler/frontend/src/org/jetbrains/kotlin/psi/KtPsiFactory.kt

@@ -179,7 +179,11 @@ class KtPsiFactory @JvmOverloads constructor(private val project: Project, val m
     }
 
     fun createCompanionObject(): KtObjectDeclaration {
-        return createClass("class A {\n companion object{\n}\n}").companionObjects.first()
+        return createCompanionObject("companion object {\n}")
+    }
+
+    fun createCompanionObject(text: String): KtObjectDeclaration {
+        return createClass("class A {\n $text\n}").companionObjects.first()
     }
 
     fun createFileAnnotation(annotationText: String): KtAnnotationEntry {

idea/src/org/jetbrains/kotlin/idea/quickfix/expectactual/CreateActualFix.kt

@@ -225,7 +225,15 @@ internal fun KtPsiFactory.generateClassOrObjectByExpectedClass(
         }
 
     val expectedText = expectedClass.text
-    val actualClass = if (expectedClass is KtObjectDeclaration) createObject(expectedText) else createClass(expectedText)
+    val actualClass = if (expectedClass is KtObjectDeclaration) {
+        if (expectedClass.isCompanion()) {
+            createCompanionObject(expectedText)
+        } else {
+            createObject(expectedText)
+        }
+    } else {
+        createClass(expectedText)
+    }
     val isInterface = expectedClass is KtClass && expectedClass.isInterface()
     actualClass.declarations.forEach {
         if (it.exists()) {

idea/testData/multiModuleQuickFix/companionAbsence/header/header.kt

@@ -0,0 +1,7 @@
+// DISABLE-ERRORS
+
+expect class WithCompanion {
+    companion object {
+        fun foo()
+    }
+}

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

@@ -0,0 +1,7 @@
+// DISABLE-ERRORS
+
+expect class WithCompanion {
+    companion object {
+        fun foo()
+    }
+}

idea/testData/multiModuleQuickFix/companionAbsence/jvm/WithCompanion.kt

@@ -0,0 +1,6 @@
+// "Add missing actual members" "true"
+// DISABLE-ERRORS
+
+actual class <caret>WithCompanion {
+
+}

idea/testData/multiModuleQuickFix/companionAbsence/jvm/WithCompanion.kt.after

@@ -0,0 +1,9 @@
+// "Add missing actual members" "true"
+// DISABLE-ERRORS
+
+actual class WithCompanion {
+    actual companion object {
+        actual fun foo() {}
+    }
+
+}

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

@@ -74,6 +74,11 @@ class QuickFixMultiModuleTest : AbstractQuickFixMultiModuleTest() {
         doMultiPlatformTest()
     }
 
+    @Test
+    fun testCompanionAbsence() {
+        doMultiPlatformTest()
+    }
+
     @Test
     fun testDeprecatedHeader() {
         doMultiPlatformTest()