Create from Usage: Fix "Create class" applicability check for when entries

#KT-22329 Fixed

idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/callableBuilder/typeUtils.kt

@@ -32,6 +32,7 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getAssignmentByLHS
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.bindingContextUtil.getDataFlowInfoAfter
 import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsStatement
@@ -148,22 +149,32 @@ fun KtExpression.guessTypes(
         && isUsedAsStatement(context)
         && getNonStrictParentOfType<KtAnnotationEntry>() == null) return arrayOf(module.builtIns.unitType)
 
-    // if we know the actual type of the expression
-    val theType1 = context.getType(this)
-    if (theType1 != null && isAcceptable(theType1)) {
-        val dataFlowInfo = context.getDataFlowInfoAfter(this)
-        val dataFlowValueFactory = this.getResolutionFacade().frontendService<DataFlowValueFactory>()
-        val dataFlowValue = dataFlowValueFactory.createDataFlowValue(this, theType1, context, module)
+    val parent = parent
 
-        val possibleTypes = dataFlowInfo.getCollectedTypes(dataFlowValue, languageVersionSettings)
-        return if (possibleTypes.isNotEmpty()) possibleTypes.toTypedArray() else arrayOf(theType1)
+    // Type/Expected type may be wrong for the expression of KtWhenEntry when some branches have unresolved expressions
+    if (parent is KtWhenEntry && parent.expression == this) {
+        return parent
+            .getStrictParentOfType<KtWhenExpression>()
+            ?.guessTypes(context, module, pseudocode, coerceUnusedToUnit, allowErrorTypes) ?: arrayOf()
+    }
+
+    if (this !is KtWhenExpression) {
+        // if we know the actual type of the expression
+        val theType1 = context.getType(this)
+        if (theType1 != null && isAcceptable(theType1)) {
+            val dataFlowInfo = context.getDataFlowInfoAfter(this)
+            val dataFlowValueFactory = this.getResolutionFacade().frontendService<DataFlowValueFactory>()
+            val dataFlowValue = dataFlowValueFactory.createDataFlowValue(this, theType1, context, module)
+
+            val possibleTypes = dataFlowInfo.getCollectedTypes(dataFlowValue, languageVersionSettings)
+            return if (possibleTypes.isNotEmpty()) possibleTypes.toTypedArray() else arrayOf(theType1)
+        }
     }
 
     // expression has an expected type
     val theType2 = context[BindingContext.EXPECTED_EXPRESSION_TYPE, this]
     if (theType2 != null && isAcceptable(theType2)) return arrayOf(theType2)
 
-    val parent = parent
     return when {
         this is KtTypeConstraint -> {
             // expression itself is a type assertion

idea/src/org/jetbrains/kotlin/idea/slicer/KotlinSliceProvider.kt

@@ -23,6 +23,7 @@ import com.intellij.psi.PsiElement
 import com.intellij.slicer.*
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.analyzeAndGetResult
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptorIfAny
 import org.jetbrains.kotlin.idea.quickfix.createFromUsage.callableBuilder.guessTypes
@@ -31,6 +32,7 @@ import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.isPlainWithEscapes
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.types.isError
 import org.jetbrains.kotlin.types.isNullabilityFlexible
@@ -53,10 +55,7 @@ class KotlinSliceProvider : SliceLanguageSupportProvider, SliceUsageTransformer
                 val types = when (element) {
                     is KtCallableDeclaration -> listOfNotNull((element.resolveToDescriptorIfAny() as? CallableDescriptor)?.returnType)
                     is KtDeclaration -> emptyList()
-                    is KtExpression -> {
-                        val (bindingContext, moduleDescriptor) = element.analyzeAndGetResult()
-                        element.guessTypes(bindingContext, moduleDescriptor).toList()
-                    }
+                    is KtExpression -> listOfNotNull(element.analyze(BodyResolveMode.PARTIAL).getType(element))
                     else -> emptyList()
                 }
                 return when {

idea/testData/quickfix/createFromUsage/createClass/callExpression/callInWhenEntry.kt

@@ -0,0 +1,15 @@
+// "Create class 'AAA'" "true"
+// ERROR: Unresolved reference: BBB
+
+abstract class I
+
+fun test(n: Int): I? {
+    return if (n > 0) {
+        when (n) {
+            1 -> <caret>AAA("1")
+            2 -> BBB("2")
+            else -> null
+        }
+    }
+    else null
+}

idea/testData/quickfix/createFromUsage/createClass/callExpression/callInWhenEntry.kt.after

@@ -0,0 +1,19 @@
+// "Create class 'AAA'" "true"
+// ERROR: Unresolved reference: BBB
+
+abstract class I
+
+fun test(n: Int): I? {
+    return if (n > 0) {
+        when (n) {
+            1 -> AAA("1")
+            2 -> BBB("2")
+            else -> null
+        }
+    }
+    else null
+}
+
+class AAA(s: String) : I() {
+
+}

idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java

@@ -1933,6 +1933,12 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
                     doTest(fileName);
                 }
 
+                @TestMetadata("callInWhenEntry.kt")
+                public void testCallInWhenEntry() throws Exception {
+                    String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/createFromUsage/createClass/callExpression/callInWhenEntry.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("callNoReceiver.kt")
                 public void testCallNoReceiver() throws Exception {
                     String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/createFromUsage/createClass/callExpression/callNoReceiver.kt");