From 5abb9dd23ecde7836ab27064949643156856e878 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <Mikhail.Glukhikh@jetbrains.com>
Date: Fri, 29 Dec 2017 20:06:44 +0300
Subject: [PATCH] Create actual fix: handle companions correctly #KT-21114
 Fixed

---
 .../src/org/jetbrains/kotlin/psi/KtPsiFactory.kt       |  6 +++++-
 .../idea/quickfix/expectactual/CreateActualFix.kt      | 10 +++++++++-
 .../companionAbsence/header/header.kt                  |  7 +++++++
 .../companionAbsence/header/header.kt.after            |  7 +++++++
 .../companionAbsence/jvm/WithCompanion.kt              |  6 ++++++
 .../companionAbsence/jvm/WithCompanion.kt.after        |  9 +++++++++
 .../kotlin/idea/quickfix/QuickFixMultiModuleTest.kt    |  5 +++++
 7 files changed, 48 insertions(+), 2 deletions(-)
 create mode 100644 idea/testData/multiModuleQuickFix/companionAbsence/header/header.kt
 create mode 100644 idea/testData/multiModuleQuickFix/companionAbsence/header/header.kt.after
 create mode 100644 idea/testData/multiModuleQuickFix/companionAbsence/jvm/WithCompanion.kt
 create mode 100644 idea/testData/multiModuleQuickFix/companionAbsence/jvm/WithCompanion.kt.after

diff --git a/compiler/frontend/src/org/jetbrains/kotlin/psi/KtPsiFactory.kt b/compiler/frontend/src/org/jetbrains/kotlin/psi/KtPsiFactory.kt
index 02acc7f070b..3ff35c3a165 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/psi/KtPsiFactory.kt
+++ b/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 {
diff --git a/idea/src/org/jetbrains/kotlin/idea/quickfix/expectactual/CreateActualFix.kt b/idea/src/org/jetbrains/kotlin/idea/quickfix/expectactual/CreateActualFix.kt
index 6b5597b9f09..c662a48213d 100644
--- a/idea/src/org/jetbrains/kotlin/idea/quickfix/expectactual/CreateActualFix.kt
+++ b/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()) {
diff --git a/idea/testData/multiModuleQuickFix/companionAbsence/header/header.kt b/idea/testData/multiModuleQuickFix/companionAbsence/header/header.kt
new file mode 100644
index 00000000000..9e928b14a15
--- /dev/null
+++ b/idea/testData/multiModuleQuickFix/companionAbsence/header/header.kt
@@ -0,0 +1,7 @@
+// DISABLE-ERRORS
+
+expect class WithCompanion {
+    companion object {
+        fun foo()
+    }
+}
\ No newline at end of file
diff --git a/idea/testData/multiModuleQuickFix/companionAbsence/header/header.kt.after b/idea/testData/multiModuleQuickFix/companionAbsence/header/header.kt.after
new file mode 100644
index 00000000000..9e928b14a15
--- /dev/null
+++ b/idea/testData/multiModuleQuickFix/companionAbsence/header/header.kt.after
@@ -0,0 +1,7 @@
+// DISABLE-ERRORS
+
+expect class WithCompanion {
+    companion object {
+        fun foo()
+    }
+}
\ No newline at end of file
diff --git a/idea/testData/multiModuleQuickFix/companionAbsence/jvm/WithCompanion.kt b/idea/testData/multiModuleQuickFix/companionAbsence/jvm/WithCompanion.kt
new file mode 100644
index 00000000000..93e15ae5d1a
--- /dev/null
+++ b/idea/testData/multiModuleQuickFix/companionAbsence/jvm/WithCompanion.kt
@@ -0,0 +1,6 @@
+// "Add missing actual members" "true"
+// DISABLE-ERRORS
+
+actual class <caret>WithCompanion {
+
+}
\ No newline at end of file
diff --git a/idea/testData/multiModuleQuickFix/companionAbsence/jvm/WithCompanion.kt.after b/idea/testData/multiModuleQuickFix/companionAbsence/jvm/WithCompanion.kt.after
new file mode 100644
index 00000000000..26e5f6947dc
--- /dev/null
+++ b/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() {}
+    }
+
+}
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixMultiModuleTest.kt b/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixMultiModuleTest.kt
index ee3d7872c3a..05d8db98e80 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixMultiModuleTest.kt
+++ b/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()