[Java. Code completion] IDEA-379114 Add tests for method completion for blocks

GitOrigin-RevId: 56a4d45b99289e03698c04d3a66c7b881f3d8d8f

java/java-impl/src/com/intellij/codeInsight/completion/commands/impl/JavaExtractConstantFieldParameterVariableMethodCompletionCommand.kt

@@ -45,7 +45,7 @@ internal class JavaExtractMethodCompletionCommandProvider : AbstractExtractMetho
   previewText = ActionsBundle.message("action.ExtractMethod.description"),
   synonyms = listOf("Extract method", "Introduce method")
 ) {
-  override fun findControlFlowStatement(offset: Int, psiFile: PsiFile): PsiStatement? {
+  override fun findControlFlowStatement(offset: Int, psiFile: PsiFile): PsiElement? {
     val element = getCommandContext(offset, psiFile) ?: return null
     val elementType = element.elementType
     if (elementType != JavaTokenType.RBRACE && elementType != JavaTokenType.LBRACE) return null
@@ -57,7 +57,14 @@ internal class JavaExtractMethodCompletionCommandProvider : AbstractExtractMetho
     if (blockParent !is PsiBlockStatement) return null
 
     val controlFlowStatement = blockParent.parent
-    if (controlFlowStatement is PsiLoopStatement || controlFlowStatement is PsiIfStatement) return controlFlowStatement
+    if (controlFlowStatement is PsiLoopStatement) return controlFlowStatement
+    else if (controlFlowStatement is PsiIfStatement) {
+      var expression = controlFlowStatement
+      while(true) {
+        val parent = expression.parent
+        if (parent is PsiIfStatement) expression = parent else return expression
+      }
+    }
     return null
   }
 

java/java-tests/testSrc/com/intellij/java/codeInsight/completion/commands/JavaCommandsCompletionTest.kt

@@ -1392,6 +1392,78 @@ class JavaCommandsCompletionTest : LightFixtureCompletionTestCase() {
     assertTrue(elements.any { element -> element.lookupString.contains("Extract method", ignoreCase = true) })
   }
 
+  fun testExtractMethodControlFlowForStatementLBrace() {
+    Registry.get("ide.completion.command.force.enabled").setValue(true, getTestRootDisposable())
+    myFixture.configureByText(JavaFileType.INSTANCE, """
+      class A {
+          void foo() {
+              for(;;) {.<caret>
+              
+              }
+          }
+      }""".trimIndent())
+
+    val elements = myFixture.completeBasic()
+    assertTrue(elements.any { element -> element.lookupString.contains("Extract method", ignoreCase = true) })
+  }
+
+  fun testExtractMethodControlFlowForStatementRBrace() {
+    Registry.get("ide.completion.command.force.enabled").setValue(true, getTestRootDisposable())
+    myFixture.configureByText(JavaFileType.INSTANCE, """
+      class A {
+          void foo() {
+              for(;;) {
+              
+              }.<caret>
+          }
+      }""".trimIndent())
+
+    val elements = myFixture.completeBasic()
+    assertTrue(elements.any { element -> element.lookupString.contains("Extract method", ignoreCase = true) })
+  }
+
+  fun testExtractMethodControlFlowIfStatement() {
+    Registry.get("ide.completion.command.force.enabled").setValue(true, getTestRootDisposable())
+    myFixture.configureByText(JavaFileType.INSTANCE, """
+      class A {
+          void foo() {
+              if(true) {.<caret>
+              
+              }
+          }
+    """)
+  }
+
+  fun testExtractMethodControlFlowIfElseStatement() {
+    Registry.get("ide.completion.command.force.enabled").setValue(true, getTestRootDisposable())
+    myFixture.configureByText(JavaFileType.INSTANCE, """
+      class A {
+          void foo() {
+              if(true) {
+              
+              } else {.<caret>
+              
+              }
+          }
+    """)
+  }
+
+  fun testExtractMethodControlFlowIfElseIfStatement() {
+    Registry.get("ide.completion.command.force.enabled").setValue(true, getTestRootDisposable())
+    myFixture.configureByText(JavaFileType.INSTANCE, """
+      class A {
+          void foo() {
+              if(true) {
+              
+              } else if(true) {
+              
+              } else if (true) {
+              } else {.<caret>
+              }
+          }
+    """)
+  }
+
   fun testMoveMethod() {
     Registry.get("ide.completion.command.force.enabled").setValue(true, getTestRootDisposable())
     myFixture.configureByText(JavaFileType.INSTANCE, """

platform/lang-impl/src/com/intellij/codeInsight/completion/command/commands/AbstractExtractConstantFieldParameterVariableMethodCompletionCommand.kt

@@ -101,9 +101,13 @@ abstract class AbstractExtractMethodCompletionCommandProvider(
     synonyms = synonyms,
   ) {
   override fun isApplicable(offset: Int, psiFile: PsiFile, editor: Editor?): Boolean {
-    if (findControlFlowStatement(offset, psiFile) != null) return true
-
-    return findOutermostExpression(offset, psiFile, editor) != null && super.isApplicable(offset, psiFile, editor)
+    val controlFlowStatement = findControlFlowStatement(offset, psiFile)
+    if (controlFlowStatement != null) {
+      editor?.selectionModel?.setSelection(controlFlowStatement.textRange.startOffset, controlFlowStatement.textRange.endOffset)
+      return super.isApplicable(offset, psiFile, editor)
+    } else {
+      return findOutermostExpression(offset, psiFile, editor) != null && super.isApplicable(offset, psiFile, editor)
+    }
   }
 
   override fun createCommand(context: CommandCompletionProviderContext): ActionCompletionCommand {

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/commands/KotlinExtractConstantFieldParameterVariableMethodCompletionCommand.kt

@@ -53,8 +53,16 @@ internal class KotlinExtractMethodCompletionCommandProvider : AbstractExtractMet
         val containerNode = parent.parent
         if (containerNode !is KtContainerNodeForControlStructureBody) return null
 
-        val controlFlow = containerNode.parent
-        if (controlFlow is KtLoopExpression || controlFlow is KtIfExpression) return controlFlow
+        val controlFlowExpression = containerNode.parent
+        if (controlFlowExpression is KtLoopExpression) return controlFlowExpression
+        else if (controlFlowExpression is KtIfExpression) {
+            var expression = controlFlowExpression
+            while (true) {
+                val parent = expression.parent
+                val grandParent = parent?.parent
+                if (parent is KtContainerNodeForControlStructureBody && grandParent is KtIfExpression) expression = grandParent else return expression
+            }
+        }
         return null
     }
 

plugins/kotlin/fir/tests/test/org/jetbrains/kotlin/idea/fir/completion/commands/K2CommandCompletionTest.kt

@@ -672,6 +672,78 @@ class K2CommandCompletionTest : KotlinLightCodeInsightFixtureTestCase() {
         assertTrue(elements.any { element -> element.lookupString.contains("Extract function", ignoreCase = true) })
     }
 
+    fun testExtractMethodInControlFlowForAfterLBrace() {
+        Registry.get("ide.completion.command.force.enabled").setValue(true, getTestRootDisposable())
+        myFixture.configureByText(
+            "x.kt", """
+            fun foo() {
+                for (i in 1..10) {.<caret>
+                    println(i)
+                }
+            }
+        """.trimIndent())
+        val elements = myFixture.completeBasic()
+        assertTrue(elements.any { element -> element.lookupString.contains("Extract function", ignoreCase = true) })
+    }
+
+    fun testExtractMethodInControlFlowForAfterRBrace() {
+        Registry.get("ide.completion.command.force.enabled").setValue(true, getTestRootDisposable())
+        myFixture.configureByText(
+            "x.kt", """
+            fun foo() {
+                for (i in 1..10) {
+                }.<caret>
+            }
+            """
+        )
+        val elements = myFixture.completeBasic()
+        assertTrue(elements.any { element -> element.lookupString.contains("Extract function", ignoreCase = true) })
+    }
+
+    fun testExtractMethodInControlFlowIf() {
+        Registry.get("ide.completion.command.force.enabled").setValue(true, getTestRootDisposable())
+        myFixture.configureByText(
+            "x.kt", """
+            fun foo() {
+                if(true) {.<caret>
+                }
+            }
+            """
+        )
+        val elements = myFixture.completeBasic()
+        assertTrue(elements.any { element -> element.lookupString.contains("Extract function", ignoreCase = true) })
+    }
+
+    fun testExtractMethodInControlFlowIfElse() {
+        Registry.get("ide.completion.command.force.enabled").setValue(true, getTestRootDisposable())
+        myFixture.configureByText(
+            "x.kt", """
+            fun foo() {
+                if(true) {
+                } else {.<caret>
+                }
+            }
+            """
+        )
+        val elements = myFixture.completeBasic()
+        assertTrue(elements.any { element -> element.lookupString.contains("Extract function", ignoreCase = true) })
+    }
+
+    fun testExtractMethodInControlFlowIfElseIf() {
+        Registry.get("ide.completion.command.force.enabled").setValue(true, getTestRootDisposable())
+        myFixture.configureByText(
+            "x.kt", """
+            fun foo() {
+                if(true) {
+                } else if (true) {
+                } else if (true) {.<caret>
+                } else {
+                }
+            }
+            """
+        )
+    }
+
     fun testInlineMethod() {
         Registry.get("ide.completion.command.force.enabled").setValue(true, getTestRootDisposable())
         myFixture.configureByText(