[kotlin] k2 extract function: let's show scope chooser with predefined selection

- setup selection for extract to scope as well

After the change, nearly each extract function would show scope chooser with preselection corresponding to the default scope. Though it requires additional enter, it highlights possible targets and provides more consistent behavior with java.

With registry key on, extract function and extract function to scope are equivalent. So the last one could be deleted.

^KTIJ-5740 fixed

GitOrigin-RevId: 4c86fabe56faf5b531f5218ba44670531eee2e18

plugins/kotlin/compiler-plugins/kotlinx-serialization/common/src/org/jetbrains/kotlin/idea/compilerPlugin/kotlinxSerialization/quickfixes/JsonRedundantQuickFix.kt

@@ -69,6 +69,7 @@ internal class JsonRedundantQuickFix(expression: KtCallExpression) : KotlinQuick
             editor,
             KotlinBundle.message("title.select.target.code.block"),
             true,
+            null,
             { it },
             { onSelect(it) }
         )

plugins/kotlin/idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/createCallable/CreateCallableFromUsageFix.kt

@@ -23,7 +23,7 @@ import org.jetbrains.kotlin.idea.base.resources.KotlinBundle
 import org.jetbrains.kotlin.idea.codeInsight.DescriptorToSourceUtilsIde
 import org.jetbrains.kotlin.idea.quickfix.KotlinCrossLanguageQuickFixAction
 import org.jetbrains.kotlin.idea.quickfix.createFromUsage.callableBuilder.*
-import org.jetbrains.kotlin.idea.refactoring.canRefactor
+import org.jetbrains.kotlin.idea.refactoring.canRefactorElement
 import org.jetbrains.kotlin.idea.refactoring.chooseContainer.chooseContainerElementIfNecessary
 import org.jetbrains.kotlin.idea.search.usagesSearch.descriptor
 import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
@@ -237,7 +237,7 @@ abstract class CreateCallableFromUsageFixBase<E : KtElement>(
         if (isExtension && staticContextRequired && descriptor is JavaClassDescriptor) return null
         val declaration = getDeclaration(descriptor, project) ?: return null
         if (declaration !is KtClassOrObject && declaration !is KtTypeParameter && declaration !is PsiClass) return null
-        return if ((isExtension && !staticContextRequired) || declaration.canRefactor()) declaration else null
+        return if ((isExtension && !staticContextRequired) || declaration.canRefactorElement()) declaration else null
     }
 
     private fun checkIsInitialized() {
@@ -306,7 +306,7 @@ abstract class CreateCallableFromUsageFixBase<E : KtElement>(
             val containers = receiverTypeCandidates
                 .mapNotNull { candidate -> getDeclarationIfApplicable(project, candidate, staticContextRequired)?.let { candidate to it } }
 
-            chooseContainerElementIfNecessary(containers, editorForBuilder, popupTitle, false, { it.second }) {
+            chooseContainerElementIfNecessary(containers, editorForBuilder, popupTitle, false, null, { it.second }) {
                 runBuilder {
                     val receiverClass = it.second as? KtClass
                     if (staticContextRequired && receiverClass?.isWritable == true) {

plugins/kotlin/refactorings/kotlin.refactorings.common/src/org/jetbrains/kotlin/idea/refactoring/chooseContainer/chooseContainerUtil.kt

@@ -38,33 +38,36 @@ fun <T> chooseContainerElementIfNecessary(
     editor: Editor,
     @NlsContexts.PopupTitle title: String,
     highlightSelection: Boolean,
+    selection: T? = null,
     toPsi: (T) -> PsiElement,
     onSelect: (T) -> Unit
-): Unit = chooseContainerElementIfNecessaryImpl(containers, editor, title, highlightSelection, toPsi, onSelect)
+): Unit = chooseContainerElementIfNecessaryImpl(containers, editor, title, highlightSelection, selection, toPsi, onSelect)
 
 fun <T : PsiElement> chooseContainerElementIfNecessary(
     containers: List<T>,
     editor: Editor,
     @NlsContexts.PopupTitle title: String,
     highlightSelection: Boolean,
+    selection: T? = null,
     onSelect: (T) -> Unit
-): Unit = chooseContainerElementIfNecessaryImpl(containers, editor, title, highlightSelection, null, onSelect)
+): Unit = chooseContainerElementIfNecessaryImpl(containers, editor, title, highlightSelection, selection, null, onSelect)
 
 private fun <T> chooseContainerElementIfNecessaryImpl(
     containers: List<T>,
     editor: Editor,
     @NlsContexts.PopupTitle title: String,
     highlightSelection: Boolean,
+    selection: T? = null,
     toPsi: ((T) -> PsiElement)?,
     onSelect: (T) -> Unit
 ) {
     when {
         containers.isEmpty() -> return
         containers.size == 1 || isUnitTestMode() -> onSelect(containers.first())
-        toPsi != null -> chooseContainerElement(containers, editor, title, highlightSelection, toPsi, onSelect)
+        toPsi != null -> chooseContainerElement(containers, editor, title, highlightSelection, selection, toPsi, onSelect)
         else -> {
             @Suppress("UNCHECKED_CAST")
-            chooseContainerElement(containers as List<PsiElement>, editor, title, highlightSelection, onSelect as (PsiElement) -> Unit)
+            chooseContainerElement(containers as List<PsiElement>, editor, title, highlightSelection, selection as PsiElement?, onSelect as (PsiElement) -> Unit)
         }
     }
 }
@@ -74,6 +77,7 @@ private fun <T> chooseContainerElement(
     editor: Editor,
     @NlsContexts.PopupTitle title: String,
     highlightSelection: Boolean,
+    selection: T? = null,
     toPsi: (T) -> PsiElement,
     onSelect: (T) -> Unit
 ) {
@@ -93,12 +97,14 @@ private fun <T : PsiElement> chooseContainerElement(
     editor: Editor,
     @NlsContexts.PopupTitle title: String,
     highlightSelection: Boolean,
+    selection: T? = null,
     onSelect: (T) -> Unit
 ): Unit = choosePsiContainerElement(
     elements = elements,
     editor = editor,
     title = title,
     highlightSelection = highlightSelection,
+    selection = selection,
     psi2Container = { it },
     onSelect = onSelect,
 )
@@ -108,6 +114,7 @@ private fun <T, E : PsiElement> choosePsiContainerElement(
     editor: Editor,
     @NlsContexts.PopupTitle title: String,
     highlightSelection: Boolean,
+    selection: E? = null,
     psi2Container: (E) -> T,
     onSelect: (T) -> Unit,
 ) {
@@ -117,6 +124,7 @@ private fun <T, E : PsiElement> choosePsiContainerElement(
         popupPresentationProvider(),
         title,
         highlightSelection,
+        selection,
     ) { psiElement ->
         @Suppress("UNCHECKED_CAST")
         onSelect(psi2Container(psiElement as E))
@@ -134,12 +142,14 @@ private fun <T : PsiElement> getPsiElementPopup(
     presentationProvider: TargetPresentationProvider<T>,
     @NlsContexts.PopupTitle title: String?,
     highlightSelection: Boolean,
+    selection: T? = null,
     processor: (T) -> Boolean
 ): JBPopup {
     val project = elements.firstOrNull()?.project ?: throw IllegalArgumentException("Can't create popup because no elements are provided")
     val highlighter = if (highlightSelection) SelectionAwareScopeHighlighter(editor) else null
     return PsiTargetNavigator(elements)
         .presentationProvider(presentationProvider)
+        .selection(selection)
         .builderConsumer { builder ->
             builder
                 .setItemSelectedCallback { presentation ->

plugins/kotlin/refactorings/kotlin.refactorings.common/src/org/jetbrains/kotlin/idea/refactoring/introduce/KotlinIntroduceVariableHandler.kt

@@ -194,7 +194,7 @@ abstract class KotlinIntroduceVariableHandler : RefactoringActionHandler {
         } else {
             chooseContainerElementIfNecessary(
                 candidateContainers, editor,
-                KotlinBundle.message("text.select.target.code.block"), true, { it.targetContainer },
+                KotlinBundle.message("text.select.target.code.block"), true, null, { it.targetContainer },
                 doRefactoring
             )
         }

plugins/kotlin/refactorings/kotlin.refactorings.common/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractFunction/AbstractExtractKotlinFunctionHandler.kt

@@ -28,6 +28,9 @@ abstract class AbstractExtractKotlinFunctionHandler(
     )
 
     fun selectElements(editor: Editor, file: KtFile, continuation: (elements: List<PsiElement>, targetSibling: PsiElement) -> Unit) {
+        val selection: ((elements: List<PsiElement>, commonParent: PsiElement) -> PsiElement?)? = if (allContainersEnabled) {
+            { elements, parent -> parent.getExtractionContainers(elements.size == 1, false, acceptScripts).firstOrNull() }
+        } else null
         selectElementsWithTargetSibling(
             EXTRACT_FUNCTION,
             editor,
@@ -36,7 +39,8 @@ abstract class AbstractExtractKotlinFunctionHandler(
             listOf(ElementKind.EXPRESSION),
             ::validateExpressionElements,
             { elements, parent -> parent.getExtractionContainers(elements.size == 1, allContainersEnabled, acceptScripts) },
-            continuation
+            continuation,
+            selection
         )
     }
 

plugins/kotlin/refactorings/kotlin.refactorings.common/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceUtils.kt

@@ -127,7 +127,8 @@ fun selectElementsWithTargetSibling(
     elementKinds: Collection<ElementKind>,
     elementValidator: (List<PsiElement>) -> String?,
     getContainers: (elements: List<PsiElement>, commonParent: PsiElement) -> List<PsiElement>,
-    continuation: (elements: List<PsiElement>, targetSibling: PsiElement) -> Unit
+    continuation: (elements: List<PsiElement>, targetSibling: PsiElement) -> Unit,
+    selection: ((elements: List<PsiElement>, commonParent: PsiElement) -> PsiElement?)? = null
 ) {
     fun onSelectionComplete(elements: List<PsiElement>, targetContainer: PsiElement) {
         val physicalElements = elements.map { it.substringContextOrThis }
@@ -148,7 +149,7 @@ fun selectElementsWithTargetSibling(
         continuation(elements, outermostParent)
     }
 
-    selectElementsWithTargetParent(operationName, editor, file, title, elementKinds, elementValidator, getContainers, ::onSelectionComplete)
+    selectElementsWithTargetParent(operationName, editor, file, title, elementKinds, elementValidator, getContainers, ::onSelectionComplete, selection)
 }
 
 fun selectElementsWithTargetParent(
@@ -158,8 +159,9 @@ fun selectElementsWithTargetParent(
     @NlsContexts.DialogTitle title: String,
     elementKinds: Collection<ElementKind>,
     elementValidator: (List<PsiElement>) -> @NlsContexts.DialogMessage String?,
-    getContainers: (elements: List<PsiElement>, commonParent: PsiElement) -> List<PsiElement>,
-    continuation: (elements: List<PsiElement>, targetParent: PsiElement) -> Unit
+    getContainers: (List<PsiElement>, PsiElement) -> List<PsiElement>,
+    continuation: (List<PsiElement>, PsiElement) -> Unit,
+    selection: ((List<PsiElement>, PsiElement) -> PsiElement?)? = null
 ) {
     fun showErrorHintByKey(key: String) {
         showErrorHintByKey(file.project, editor, key, operationName)
@@ -185,7 +187,8 @@ fun selectElementsWithTargetParent(
             containers,
             editor,
             title,
-            true
+            true,
+            selection?.invoke(physicalElements, parent)
         ) {
             continuation(elements, it)
         }

plugins/kotlin/refactorings/kotlin.refactorings.k2/resources/kotlin.refactorings.k2.xml

@@ -38,6 +38,8 @@
             order="first" />
     <updateAddedFileProcessor implementation="org.jetbrains.kotlin.idea.k2.refactoring.copy.KotlinUpdateAddedFileProcessor" order="first"/>
 
+    <registryKey defaultValue="true" description="Show scope chooser for kotlin extract function refactoring"
+                 key="k2.extract.function.scope.chooser"/>
   </extensions>
 
   <extensions defaultExtensionNs="org.jetbrains.kotlin">

plugins/kotlin/refactorings/kotlin.refactorings.k2/src/org/jetbrains/kotlin/idea/k2/refactoring/extractFunction/ExtractK2FunctionAction.kt

@@ -3,6 +3,7 @@
 package org.jetbrains.kotlin.idea.k2.refactoring.extractFunction
 
 import com.intellij.lang.refactoring.RefactoringSupportProvider
+import com.intellij.openapi.util.registry.Registry
 import com.intellij.psi.PsiElement
 import com.intellij.refactoring.RefactoringActionHandler
 import com.intellij.refactoring.actions.BasePlatformRefactoringAction
@@ -11,7 +12,7 @@ import org.jetbrains.kotlin.psi.KtElement
 
 class ExtractK2FunctionAction : BasePlatformRefactoringAction() {
     override fun getRefactoringHandler(provider: RefactoringSupportProvider): RefactoringActionHandler? =
-        KotlinFirExtractFunctionHandler(false)
+        KotlinFirExtractFunctionHandler(Registry.`is`("k2.extract.function.scope.chooser", true))
 
     override fun isAvailableInEditorOnly(): Boolean {
         return true