support creating companion objects (part of KTIJ-28926 K2: Create Function From Usage Fix)

GitOrigin-RevId: e5266b8bd296f2878eb2e6e2e448f037b34bf87e

plugins/kotlin/code-insight/kotlin.code-insight.k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/quickFixes/createFromUsage/CreateKotlinCallableAction.kt

@@ -87,7 +87,7 @@ internal class CreateKotlinCallableAction(
             require(call != null)
             CreateKotlinCallablePsiEditor(
                 project, pointerToContainer, callableInfo,
-            ).execute(call)
+            ).execute(call, request)
         }
     }
 
@@ -120,10 +120,9 @@ internal class CreateKotlinCallableAction(
 
     context (KtAnalysisSession)
     private fun renderTypeName(expectedType: ExpectedType, container: KtElement): String? {
-        val returnKtType = if (expectedType is ExpectedKotlinType) expectedType.ktType else expectedType.toKtTypeWithNullability(container)
-        if (returnKtType == null || returnKtType == builtinTypes.UNIT) return null
-        if (!isAccessibleInCreationPlace(returnKtType, (request as? CreateMethodFromKotlinUsageRequest)?.call)) return null
-        return returnKtType.render(renderer = WITH_TYPE_NAMES_FOR_CREATE_ELEMENTS, position = Variance.INVARIANT)
+        val ktType = if (expectedType is ExpectedKotlinType) expectedType.ktType else expectedType.toKtTypeWithNullability(container)
+        if (ktType == null || ktType == builtinTypes.UNIT) return null
+        return ktType.render(renderer = WITH_TYPE_NAMES_FOR_CREATE_ELEMENTS, position = Variance.INVARIANT)
     }
 
     private fun buildCallableAsString(): String? {

plugins/kotlin/code-insight/kotlin.code-insight.k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/quickFixes/createFromUsage/CreateKotlinCallablePsiEditor.kt

@@ -8,6 +8,7 @@ import com.intellij.codeInsight.template.TemplateEditingAdapter
 import com.intellij.codeInsight.template.TemplateManager
 import com.intellij.codeInsight.template.impl.TemplateImpl
 import com.intellij.ide.util.EditorHelper
+import com.intellij.lang.jvm.actions.CreateMethodRequest
 import com.intellij.openapi.command.WriteCommandAction
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.RangeMarker
@@ -54,7 +55,7 @@ internal class CreateKotlinCallablePsiEditor(
     private val pointerToContainer: SmartPsiElementPointer<*>,
     private val callableInfo: NewCallableInfo,
 ) {
-    fun execute(anchor: PsiElement) {
+    fun execute(anchor: PsiElement, request: CreateMethodRequest) {
         val factory = KtPsiFactory(project)
         var function = factory.createFunction(callableInfo.definitionAsString)
         val passedContainerElement = pointerToContainer.element
@@ -67,13 +68,26 @@ internal class CreateKotlinCallablePsiEditor(
             passedContainerElement
         }
 
-        val containerMaybeCompanion = if (insertContainer is KtClass && callableInfo.isForCompanion) {
-            insertContainer.getOrCreateCompanionObject()
-        } else insertContainer
+        val containerMaybeCompanion = if (callableInfo.isForCompanion) {
+            if (insertContainer is KtClass) {
+                insertContainer.getOrCreateCompanionObject()
+            } else {
+                val targetClass = (request as? CreateMethodFromKotlinUsageRequest)?.targetClass
+                val ktClass = targetClass as? KtClass
+                if (ktClass != null) {
+                    val hasCompanionObject = ktClass.companionObjects.isNotEmpty()
+                    val companion = ktClass.getOrCreateCompanionObject()
+                    if (!hasCompanionObject && request.isExtension) {
+                        companion.body?.delete()
+                    }
+                }
+                insertContainer
+            }
+        }
+        else insertContainer
 
         function = CreateFromUsageUtil.placeDeclarationInContainer(function, containerMaybeCompanion, anchor)
 
-        //function = function.addToContainer(container) as? KtNamedFunction ?: return
         function = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(function) ?: return
         runTemplate(function)
     }

plugins/kotlin/code-insight/kotlin.code-insight.k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/quickFixes/createFromUsage/CreateMethodFromKotlinUsageRequest.kt

@@ -6,7 +6,9 @@ import com.intellij.lang.jvm.actions.CreateMethodRequest
 import com.intellij.lang.jvm.actions.ExpectedType
 import com.intellij.lang.jvm.types.JvmReferenceType
 import com.intellij.openapi.util.NlsSafe
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.analysis.api.analyze
+import org.jetbrains.kotlin.analysis.api.symbols.KtClassLikeSymbol
 import org.jetbrains.kotlin.analysis.api.types.KtType
 import org.jetbrains.kotlin.psi.KtCallExpression
 import org.jetbrains.kotlin.psi.KtExpression
@@ -22,8 +24,16 @@ internal class CreateMethodFromKotlinUsageRequest (
     val receiverType: KtType?, // (in case receiverExpression is null) it can be notnull when there's implicit receiver: `blah { unknownFunc() }`
     val isExtension: Boolean,
     val isAbstractClassOrInterface: Boolean,
-    val isForCompanion: Boolean
+    val isForCompanion: Boolean,
 ) : CreateExecutableFromKotlinUsageRequest<KtCallExpression>(functionCall, modifiers), CreateMethodRequest {
+    internal val targetClass: PsiElement? = initializeTargetClass(receiverExpression, functionCall)
+
+    private fun initializeTargetClass(receiverExpression: KtExpression?, functionCall: KtCallExpression): PsiElement? {
+        return analyze(functionCall) {
+            (receiverExpression?.resolveExpression() as? KtClassLikeSymbol)?.psi
+        }
+    }
+
     private val returnType:List<ExpectedType> = initializeReturnType(functionCall)
 
     private fun initializeReturnType(functionCall: KtCallExpression): List<ExpectedType> {

plugins/kotlin/code-insight/kotlin.code-insight.k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/quickFixes/createFromUsage/KotlinElementActionsFactory.kt

@@ -7,29 +7,14 @@ import com.intellij.lang.jvm.JvmClass
 import com.intellij.lang.jvm.actions.CreateMethodRequest
 import com.intellij.lang.jvm.actions.JvmElementActionsFactory
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.asJava.classes.KtLightClassForFacade
-import org.jetbrains.kotlin.asJava.elements.KtLightElement
 import org.jetbrains.kotlin.idea.base.resources.KotlinBundle
 import org.jetbrains.kotlin.idea.refactoring.isAbstract
 import org.jetbrains.kotlin.idea.refactoring.isInterfaceClass
 import org.jetbrains.kotlin.psi.KtClassOrObject
 import org.jetbrains.kotlin.psi.KtElement
-import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.psiUtil.createSmartPointer
 
 class KotlinElementActionsFactory : JvmElementActionsFactory() {
-    private fun JvmClass.toKtClassOrFile(): KtElement? = if (this is JvmClassWrapperForKtClass<*>) {
-        ktClassOrFile
-    } else {
-        when (val psi = sourceElement) {
-            is KtClassOrObject -> psi
-            is KtLightClassForFacade -> psi.files.firstOrNull()
-            is KtLightElement<*, *> -> psi.kotlinOrigin
-            is KtFile -> psi
-            else -> null
-        }
-    }
-
     override fun createAddMethodActions(targetClass: JvmClass, request: CreateMethodRequest): List<IntentionAction> {
         if (targetClass is PsiElement && !BaseIntentionAction.canModify(targetClass)) return emptyList()
         var container = targetClass.toKtClassOrFile() ?: return emptyList()

plugins/kotlin/code-insight/kotlin.code-insight.k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/quickFixes/createFromUsage/Utils.kt

@@ -30,6 +30,8 @@ import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.analysis.api.types.*
 import org.jetbrains.kotlin.analysis.api.types.KtIntersectionType
 import org.jetbrains.kotlin.analysis.utils.printer.PrettyPrinter
+import org.jetbrains.kotlin.asJava.classes.KtLightClassForFacade
+import org.jetbrains.kotlin.asJava.elements.KtLightElement
 import org.jetbrains.kotlin.asJava.findFacadeClass
 import org.jetbrains.kotlin.asJava.toLightClass
 import org.jetbrains.kotlin.idea.base.codeInsight.KotlinNameSuggester
@@ -106,6 +108,9 @@ internal fun KtElement.getExpectedKotlinType(): ExpectedType? {
         }
     }
     if (expectedType == null) return null
+    if (!isAccessibleInCreationPlace(expectedType, this)) {
+        return null
+    }
     val jvmType = expectedType.convertToJvmType(this) ?: return null
     return ExpectedKotlinType(expectedType, jvmType)
 }
@@ -287,10 +292,24 @@ fun isAccessibleInCreationPlace(ktType: KtType, call: KtElement?): Boolean {
             // having `<T> caller(T t) { unknownMethod(t); }` the type `T` is not accessible in created method unknownMethod(t)
             val owner = (ktLeaf.symbol.psi as? KtTypeParameter)?.getOwningTypeParameterOwner()
             // todo must have been "insertion point" instead of `call`
-            owner == null || PsiTreeUtil.isAncestor(owner, call, false)
+            if (owner == null) true
+            else if (owner is KtCallableDeclaration) false
+            else PsiTreeUtil.isAncestor(owner, call, false)
         } else {
             // KtErrorType means this type is unresolved in the context of container
             ktLeaf !is KtErrorType
         }
     }
 }
+
+fun JvmClass.toKtClassOrFile(): KtElement? = if (this is JvmClassWrapperForKtClass<*>) {
+    ktClassOrFile
+} else {
+    when (val psi = sourceElement) {
+        is KtClassOrObject -> psi
+        is KtLightClassForFacade -> psi.files.firstOrNull()
+        is KtLightElement<*, *> -> psi.kotlinOrigin
+        is KtFile -> psi
+        else -> null
+    }
+}

plugins/kotlin/code-insight/kotlin.code-insight.k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/quickFixes/createFromUsage/createKotlinCallableFromUsage.kt

@@ -97,7 +97,7 @@ internal fun buildRequests(callExpression: KtCallExpression): List<Pair<JvmClass
                     receiverType = null,
                     isAbstractClassOrInterface = true,
                     isExtension = false,
-                    isForCompanion = false
+                    isForCompanion = false,
                 ))
             }
         }
@@ -107,7 +107,8 @@ internal fun buildRequests(callExpression: KtCallExpression): List<Pair<JvmClass
                 implicitReceiverType?.convertToClass() ?: calleeExpression.getNonStrictParentOfType<KtClassOrObject>()
                 ?: calleeExpression.containingKtFile
             val jvmClassWrapper = JvmClassWrapperForKtClass(containerClassForExtension)
-            val modifiers = computeModifiers(defaultContainerPsi?:calleeExpression.containingFile, calleeExpression, callExpression, false, true)
+            val shouldCreateCompanionClass = shouldCreateCompanionClass(calleeExpression)
+            val modifiers = computeModifiers(defaultContainerPsi?:calleeExpression.containingFile, calleeExpression, callExpression, shouldCreateCompanionClass, true)
             requests.add(jvmClassWrapper to CreateMethodFromKotlinUsageRequest(
                 callExpression,
                 modifiers,
@@ -115,7 +116,7 @@ internal fun buildRequests(callExpression: KtCallExpression): List<Pair<JvmClass
                 receiverType = implicitReceiverType,
                 isExtension = true,
                 isAbstractClassOrInterface = false,
-                isForCompanion = false
+                isForCompanion = shouldCreateCompanionClass,
             ))
         }
     return requests
@@ -146,15 +147,16 @@ private fun computeModifiers(
     isExtension: Boolean
 ): List<JvmModifier> {
     if (shouldCreateCompanionClass) {
-        return listOf(JvmModifier.PUBLIC, JvmModifier.STATIC) // methods in the companion class are typically public (and static in case of java)
+        // methods in the companion class are typically public (and static in case of java)
+        return if (samePackage(calleeExpression, callExpression)) if (isExtension) listOf(JvmModifier.PRIVATE, JvmModifier.STATIC) else listOf(JvmModifier.STATIC) else listOf(JvmModifier.PUBLIC, JvmModifier.STATIC)
     }
     val modifier = CreateFromUsageUtil.patchVisibilityForInlineFunction(callExpression)
     if (modifier != null) {
         return listOf(allModifiers[modifier]!!)
     }
     val modifierOwner = callExpression.getNonStrictParentOfType<KtModifierListOwner>()
-    val packageNameOfReceiver = calleeExpression.getReceiverOrContainerClassPackageName()
-    if (modifierOwner != null && packageNameOfReceiver != null && packageNameOfReceiver != callExpression.containingKtFile.packageFqName) {
+    val samePackage = samePackage(calleeExpression, callExpression)
+    if (modifierOwner != null && !samePackage) {
         for (entry in allModifiers) {
             if (modifierOwner.hasModifier(entry.key)) {
                 return listOf(entry.value)
@@ -175,6 +177,17 @@ private fun computeModifiers(
 
     return listOf(JvmModifier.PUBLIC)
 }
+
+context (KtAnalysisSession)
+private fun samePackage(
+    calleeExpression: KtSimpleNameExpression,
+    callExpression: KtCallExpression
+): Boolean {
+    val packageNameOfReceiver = calleeExpression.getReceiverOrContainerClassPackageName()
+    val samePackage = packageNameOfReceiver != null && packageNameOfReceiver == callExpression.containingKtFile.packageFqName
+    return samePackage
+}
+
 /**
  * Returns the type of the class containing this [KtSimpleNameExpression] if the class is abstract. Otherwise, returns null.
  */

plugins/kotlin/idea/tests/testData/quickfix/createFromUsage/createFunction/call/extension/approximateAnonymousObjectRuntime.kt

@@ -1,4 +1,4 @@
-// "Create extension function 'List<Int>.foo'" "true"
+// "/(Create extension function 'List<Int>.foo')|(Create extension function 'List.foo')/" "true"
 // FUS_QUICKFIX_NAME: org.jetbrains.kotlin.idea.quickfix.createFromUsage.createCallable.CreateExtensionCallableFromUsageFix
 // FUS_K2_QUICKFIX_NAME: org.jetbrains.kotlin.idea.k2.codeinsight.quickFixes.createFromUsage.CreateKotlinCallableAction
 

plugins/kotlin/idea/tests/testData/quickfix/createFromUsage/createFunction/call/extension/approximateAnonymousObjectRuntime.kt.after

@@ -1,4 +1,4 @@
-// "Create extension function 'List<Int>.foo'" "true"
+// "/(Create extension function 'List<Int>.foo')|(Create extension function 'List.foo')/" "true"
 // FUS_QUICKFIX_NAME: org.jetbrains.kotlin.idea.quickfix.createFromUsage.createCallable.CreateExtensionCallableFromUsageFix
 // FUS_K2_QUICKFIX_NAME: org.jetbrains.kotlin.idea.k2.codeinsight.quickFixes.createFromUsage.CreateKotlinCallableAction
 

plugins/kotlin/idea/tests/testData/quickfix/createFromUsage/createFunction/call/extension/approximateLocalClassRuntime.kt

@@ -1,4 +1,4 @@
-// "Create extension function 'List<Int>.foo'" "true"
+// "/(Create extension function 'List<Int>.foo')|(Create extension function 'List.foo')/" "true"
 // FUS_QUICKFIX_NAME: org.jetbrains.kotlin.idea.quickfix.createFromUsage.createCallable.CreateExtensionCallableFromUsageFix
 // FUS_K2_QUICKFIX_NAME: org.jetbrains.kotlin.idea.k2.codeinsight.quickFixes.createFromUsage.CreateKotlinCallableAction
 

plugins/kotlin/idea/tests/testData/quickfix/createFromUsage/createFunction/call/extension/approximateLocalClassRuntime.kt.after

@@ -1,4 +1,4 @@
-// "Create extension function 'List<Int>.foo'" "true"
+// "/(Create extension function 'List<Int>.foo')|(Create extension function 'List.foo')/" "true"
 // FUS_QUICKFIX_NAME: org.jetbrains.kotlin.idea.quickfix.createFromUsage.createCallable.CreateExtensionCallableFromUsageFix
 // FUS_K2_QUICKFIX_NAME: org.jetbrains.kotlin.idea.k2.codeinsight.quickFixes.createFromUsage.CreateKotlinCallableAction
 

plugins/kotlin/idea/tests/testData/quickfix/createFromUsage/createFunction/call/extension/extension.kt

@@ -1,4 +1,4 @@
-// "Create extension function 'List<T>.foo'" "true"
+// "/(Create extension function 'List<T>.foo')|(Create extension function 'List.foo')/" "true"
 // WITH_STDLIB
 // FUS_QUICKFIX_NAME: org.jetbrains.kotlin.idea.quickfix.createFromUsage.createCallable.CreateExtensionCallableFromUsageFix
 // FUS_K2_QUICKFIX_NAME: org.jetbrains.kotlin.idea.k2.codeinsight.quickFixes.createFromUsage.CreateKotlinCallableAction

plugins/kotlin/idea/tests/testData/quickfix/createFromUsage/createFunction/call/extension/extension.kt.after

@@ -1,4 +1,4 @@
-// "Create extension function 'List<T>.foo'" "true"
+// "/(Create extension function 'List<T>.foo')|(Create extension function 'List.foo')/" "true"
 // WITH_STDLIB
 // FUS_QUICKFIX_NAME: org.jetbrains.kotlin.idea.quickfix.createFromUsage.createCallable.CreateExtensionCallableFromUsageFix
 // FUS_K2_QUICKFIX_NAME: org.jetbrains.kotlin.idea.k2.codeinsight.quickFixes.createFromUsage.CreateKotlinCallableAction

plugins/kotlin/idea/tests/testData/quickfix/createFromUsage/createFunction/call/extension/extensionPartialSubstitution.kt

@@ -1,4 +1,4 @@
-// "Create extension function 'List<T>.foo'" "true"
+// "/(Create extension function 'List<T>.foo')|(Create extension function 'List.foo')/" "true"
 // WITH_STDLIB
 // FUS_QUICKFIX_NAME: org.jetbrains.kotlin.idea.quickfix.createFromUsage.createCallable.CreateExtensionCallableFromUsageFix
 // FUS_K2_QUICKFIX_NAME: org.jetbrains.kotlin.idea.k2.codeinsight.quickFixes.createFromUsage.CreateKotlinCallableAction

plugins/kotlin/idea/tests/testData/quickfix/createFromUsage/createFunction/call/extension/extensionPartialSubstitution.kt.after

@@ -1,4 +1,4 @@
-// "Create extension function 'List<T>.foo'" "true"
+// "/(Create extension function 'List<T>.foo')|(Create extension function 'List.foo')/" "true"
 // WITH_STDLIB
 // FUS_QUICKFIX_NAME: org.jetbrains.kotlin.idea.quickfix.createFromUsage.createCallable.CreateExtensionCallableFromUsageFix
 // FUS_K2_QUICKFIX_NAME: org.jetbrains.kotlin.idea.k2.codeinsight.quickFixes.createFromUsage.CreateKotlinCallableAction

plugins/kotlin/idea/tests/testData/quickfix/createFromUsage/createFunction/call/extension/extensionWithReceiverArg.kt

@@ -1,4 +1,4 @@
-// "Create extension function 'List<T>.foo'" "true"
+// "/(Create extension function 'List<T>.foo')|(Create extension function 'List.foo')/" "true"
 // WITH_STDLIB
 // FUS_QUICKFIX_NAME: org.jetbrains.kotlin.idea.quickfix.createFromUsage.createCallable.CreateExtensionCallableFromUsageFix
 // FUS_K2_QUICKFIX_NAME: org.jetbrains.kotlin.idea.k2.codeinsight.quickFixes.createFromUsage.CreateKotlinCallableAction

plugins/kotlin/idea/tests/testData/quickfix/createFromUsage/createFunction/call/extension/extensionWithReceiverArg.kt.after

@@ -1,4 +1,4 @@
-// "Create extension function 'List<T>.foo'" "true"
+// "/(Create extension function 'List<T>.foo')|(Create extension function 'List.foo')/" "true"
 // WITH_STDLIB
 // FUS_QUICKFIX_NAME: org.jetbrains.kotlin.idea.quickfix.createFromUsage.createCallable.CreateExtensionCallableFromUsageFix
 // FUS_K2_QUICKFIX_NAME: org.jetbrains.kotlin.idea.k2.codeinsight.quickFixes.createFromUsage.CreateKotlinCallableAction