From 3f7c983f043365dd7b9fcf3a29be5a8e771a7a72 Mon Sep 17 00:00:00 2001
From: Dmitry Gridin <dmitry.gridin@jetbrains.com>
Date: Fri, 9 Aug 2019 19:29:38 +0300
Subject: [PATCH] CreateExpectFix: fix selection for primary constructor

---
 .../idea/quickfix/expectactual/CreateExpectedFix.kt  | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/quickfix/expectactual/CreateExpectedFix.kt b/idea/src/org/jetbrains/kotlin/idea/quickfix/expectactual/CreateExpectedFix.kt
index 32722a2ec43..899baa8aaf0 100644
--- a/idea/src/org/jetbrains/kotlin/idea/quickfix/expectactual/CreateExpectedFix.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/quickfix/expectactual/CreateExpectedFix.kt
@@ -138,7 +138,7 @@ class CreateExpectedClassFix(
             val prefix = klass.fqName?.asString()?.plus(".") ?: ""
             chooseMembers(project, membersForSelection, prefix) ?: return@block null
         }
-    }.plus(members.filter(KtNamedDeclaration::isAlwaysActual))
+    }.asSequence().plus(members.filter(KtNamedDeclaration::isAlwaysActual)).flatMap(KtNamedDeclaration::selected).toSet()
 
     val selectedClasses = findClasses(selectedElements + klass, commonModule)
     val resultDeclarations = if (selectedClasses != existingClasses) {
@@ -164,7 +164,7 @@ class CreateExpectedClassFix(
     generateClassOrObject(project, true, element, commonModule, listOfNotNull(outerExpectedClass), selectedClasses)
 })
 
-private tailrec fun findClasses(elements: List<KtNamedDeclaration>, module: Module): HashSet<String> {
+private tailrec fun findClasses(elements: Collection<KtNamedDeclaration>, module: Module): HashSet<String> {
     val classes = elements.filterIsInstance<KtClassOrObject>()
     val existingNames = classes.mapNotNull { it.fqName?.asString() }.toHashSet()
     val newExistingClasses = classes.filter { it.checkTypeAccessibilityInModule(module, existingNames) }
@@ -273,10 +273,16 @@ private fun repairActualModifiers(
 
 private tailrec fun KtDeclaration.makeActualWithParents() {
     makeActual()
-    safeAs<KtParameter>()?.parent?.parent?.safeAs<KtPrimaryConstructor>()?.takeUnless(KtDeclaration::hasActualModifier)?.makeActual()
     containingClassOrObject?.takeUnless(KtDeclaration::hasActualModifier)?.makeActualWithParents()
 }
 
+private fun KtNamedDeclaration.selected(): Sequence<KtNamedDeclaration> {
+    val additionalSequence = safeAs<KtParameter>()?.parent?.parent?.safeAs<KtPrimaryConstructor>()?.let {
+        sequenceOf(it)
+    } ?: emptySequence()
+    return sequenceOf(this) + additionalSequence + containingClassOrObject?.selected().orEmpty()
+}
+
 class CreateExpectedPropertyFix(
     property: KtNamedDeclaration,
     targetExpectedClass: KtClassOrObject?,