Allow "Create class" action in return statement

#KT-22137 Fixed

idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/createClass/CreateClassFromUsageFix.kt

@@ -49,6 +49,7 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.allChildren
 import org.jetbrains.kotlin.psi.psiUtil.getQualifiedElementSelector
+import org.jetbrains.kotlin.types.typeUtil.isUnit
 import org.jetbrains.kotlin.utils.SmartList
 import java.util.*
 import com.intellij.codeInsight.daemon.impl.quickfix.ClassKind as IdeaClassKind
@@ -221,7 +222,11 @@ open class CreateClassFromUsageFix<E : KtElement> protected constructor(
                         is PsiPackage -> createFileByPackage(selectedParent, editor, file)
                         else -> throw AssertionError("Unexpected element: " + selectedParent.text)
                     } ?: return@runWriteAction
-                val constructorInfo = ClassWithPrimaryConstructorInfo(classInfo, expectedTypeInfo)
+                val constructorInfo = ClassWithPrimaryConstructorInfo(
+                    classInfo,
+                    // Need for #KT-22137
+                    if (expectedTypeInfo.isUnit) TypeInfo.Empty else expectedTypeInfo
+                )
                 val builder = CallableBuilderConfiguration(
                     Collections.singletonList(constructorInfo),
                     element,
@@ -274,3 +279,5 @@ open class CreateClassFromUsageFix<E : KtElement> protected constructor(
         }
     }
 }
+
+private val TypeInfo.isUnit: Boolean get() = ((this as? TypeInfo.DelegatingTypeInfo)?.delegate as? TypeInfo.ByType)?.theType?.isUnit() == true

idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/createClass/createClassUtils.kt

@@ -43,6 +43,7 @@ import org.jetbrains.kotlin.resolve.scopes.receivers.Qualifier
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
 import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.substitutions.getTypeSubstitution
+import org.jetbrains.kotlin.types.typeUtil.isUnit
 import java.util.*
 import org.jetbrains.kotlin.descriptors.ClassKind as ClassDescriptorKind
 
@@ -99,7 +100,7 @@ internal fun KotlinType.toClassTypeInfo(): TypeInfo {
 internal fun getClassKindFilter(expectedType: KotlinType, containingDeclaration: PsiElement): (ClassKind) -> Boolean {
     val descriptor = expectedType.constructor.declarationDescriptor ?: return { _ -> false }
 
-    val canHaveSubtypes = !(expectedType.constructor.isFinal || expectedType.containsStarProjections())
+    val canHaveSubtypes = !(expectedType.constructor.isFinal || expectedType.containsStarProjections()) || expectedType.isUnit()
     val isEnum = DescriptorUtils.isEnumClass(descriptor)
 
     if (!(canHaveSubtypes || isEnum)

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

@@ -0,0 +1,14 @@
+// "Create class 'BookKeeper'" "true"
+// ERROR: Type mismatch: inferred type is BookKeeper but Unit was expected
+package pack
+
+import pack.Currrency.EUR
+
+enum class Currrency { EUR }
+class Item(val p1: Double, p2: Currrency)
+class Transaction(vararg val p: Item)
+
+fun place() {
+    val transactions = listOf(Transaction(Item(10.0, EUR), Item(10.0, EUR)))
+    return BookKee<caret>per(transactions)
+}

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

@@ -0,0 +1,18 @@
+// "Create class 'BookKeeper'" "true"
+// ERROR: Type mismatch: inferred type is BookKeeper but Unit was expected
+package pack
+
+import pack.Currrency.EUR
+
+enum class Currrency { EUR }
+class Item(val p1: Double, p2: Currrency)
+class Transaction(vararg val p: Item)
+
+fun place() {
+    val transactions = listOf(Transaction(Item(10.0, EUR), Item(10.0, EUR)))
+    return BookKeeper(transactions)
+}
+
+class BookKeeper(transactions: List<Transaction>) {
+
+}

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

@@ -0,0 +1,17 @@
+// "Create class 'BookKeeper'" "true"
+// ERROR: Type mismatch: inferred type is BookKeeper but Unit was expected
+// WITH_RUNTIME
+package pack
+
+import pack.Currrency.EUR
+
+enum class Currrency { EUR }
+class Item(val p1: Double, p2: Currrency)
+class Transaction(vararg val p: Item)
+
+fun place() {
+    val transactions = listOf(Transaction(Item(10.0, EUR), Item(10.0, EUR)))
+    if (transactions.isNotEmpty()) {
+        return BookKee<caret>per(transactions)
+    }
+}

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

@@ -0,0 +1,21 @@
+// "Create class 'BookKeeper'" "true"
+// ERROR: Type mismatch: inferred type is BookKeeper but Unit was expected
+// WITH_RUNTIME
+package pack
+
+import pack.Currrency.EUR
+
+enum class Currrency { EUR }
+class Item(val p1: Double, p2: Currrency)
+class Transaction(vararg val p: Item)
+
+fun place() {
+    val transactions = listOf(Transaction(Item(10.0, EUR), Item(10.0, EUR)))
+    if (transactions.isNotEmpty()) {
+        return BookKeeper(transactions)
+    }
+}
+
+class BookKeeper(transactions: List<Transaction>) {
+
+}

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

@@ -0,0 +1,17 @@
+// "Create class 'BookKeeper'" "false"
+// ACTION: Create function 'BookKeeper'
+// ACTION: Rename reference
+// ERROR: Unresolved reference: BookKeeper
+package pack
+
+import pack.Currrency.EUR
+
+enum class Currrency { EUR }
+class Item(val p1: Double, p2: Currrency)
+class Transaction(vararg val p: Item)
+class Man
+
+fun place(): Man {
+    val transactions = listOf(Transaction(Item(10.0, EUR), Item(10.0, EUR)))
+    return BookKee<caret>per(transactions)
+}

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

@@ -2656,6 +2656,21 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
                     runTest("idea/testData/quickfix/createFromUsage/createClass/callExpression/expectedTypeBySuperFunction.kt");
                 }
 
+                @TestMetadata("inReturn.kt")
+                public void testInReturn() throws Exception {
+                    runTest("idea/testData/quickfix/createFromUsage/createClass/callExpression/inReturn.kt");
+                }
+
+                @TestMetadata("inReturn2.kt")
+                public void testInReturn2() throws Exception {
+                    runTest("idea/testData/quickfix/createFromUsage/createClass/callExpression/inReturn2.kt");
+                }
+
+                @TestMetadata("notApplicableInReturn.kt")
+                public void testNotApplicableInReturn() throws Exception {
+                    runTest("idea/testData/quickfix/createFromUsage/createClass/callExpression/notApplicableInReturn.kt");
+                }
+
                 @TestMetadata("quotedName.kt")
                 public void testQuotedName() throws Exception {
                     runTest("idea/testData/quickfix/createFromUsage/createClass/callExpression/quotedName.kt");