From bdc4c6c8c1684d25520db719c6bb2f42f82b76c8 Mon Sep 17 00:00:00 2001
From: Tagir Valeev <tagir.valeev@jetbrains.com>
Date: Fri, 15 Sep 2023 18:01:05 +0200
Subject: [PATCH] [java-intentions] IDEA-331850 Support 'Unroll Loop' intention
 when there are declarations in a loop

GitOrigin-RevId: 5fa9904997f3eade008dae8278ffc44dd0e9172d
---
 .../intention/impl/UnrollLoopAction.java      | 41 +++++++++++++++----
 .../unrollLoop/afterDeclareConflict.java      | 10 +++++
 .../quickFix/unrollLoop/afterDeclareOne.java  |  7 ++++
 .../quickFix/unrollLoop/afterDeclareTwo.java  | 13 ++++++
 .../unrollLoop/beforeDeclareConflict.java     | 10 +++++
 .../quickFix/unrollLoop/beforeDeclareOne.java |  9 ++++
 .../quickFix/unrollLoop/beforeDeclareTwo.java |  9 ++++
 7 files changed, 90 insertions(+), 9 deletions(-)
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/afterDeclareConflict.java
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/afterDeclareOne.java
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/afterDeclareTwo.java
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/beforeDeclareConflict.java
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/beforeDeclareOne.java
 create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/beforeDeclareTwo.java

diff --git a/java/java-impl/src/com/intellij/codeInsight/intention/impl/UnrollLoopAction.java b/java/java-impl/src/com/intellij/codeInsight/intention/impl/UnrollLoopAction.java
index 4718bee6baaa..25ccae85baaf 100644
--- a/java/java-impl/src/com/intellij/codeInsight/intention/impl/UnrollLoopAction.java
+++ b/java/java-impl/src/com/intellij/codeInsight/intention/impl/UnrollLoopAction.java
@@ -62,7 +62,6 @@ public class UnrollLoopAction extends PsiUpdateModCommandAction<PsiLoopStatement
     }
     PsiStatement[] statements = ControlFlowUtils.unwrapBlock(body);
     if (statements.length == 0) return null;
-    if (ContainerUtil.exists(statements, PsiDeclarationStatement.class::isInstance)) return null;
     if (VariableAccessUtils.variableIsAssigned(iterationParameter, body)) return null;
     for (PsiStatement statement : statements) {
       if (isLoopBreak(statement)) continue;
@@ -193,7 +192,7 @@ public class UnrollLoopAction extends PsiUpdateModCommandAction<PsiLoopStatement
 
   @Override
   protected void invoke(@NotNull ActionContext context, @NotNull PsiLoopStatement loop, @NotNull ModPsiUpdater updater) {
-    if (!(loop.getParent() instanceof PsiCodeBlock)) return;
+    if (!(loop.getParent() instanceof PsiCodeBlock parentBlock)) return;
     List<PsiExpression> expressions = extractExpressions(loop);
     if (expressions.isEmpty()) return;
     Project project = context.project();
@@ -215,14 +214,19 @@ public class UnrollLoopAction extends PsiUpdateModCommandAction<PsiLoopStatement
       PsiElement[] children;
       if (body instanceof PsiBlockStatement) {
         PsiCodeBlock block = ((PsiBlockStatement)Objects.requireNonNull(loop.getBody())).getCodeBlock();
-        PsiElement firstBodyElement = block.getFirstBodyElement();
-        PsiElement lastBodyElement = block.getLastBodyElement();
-        if (firstBodyElement != null && lastBodyElement != null) {
-          ct.markRangeUnchanged(firstBodyElement, lastBodyElement);
+        if (canUnwrapBlock(block, parentBlock, expressions)) {
+          PsiElement firstBodyElement = block.getFirstBodyElement();
+          PsiElement lastBodyElement = block.getLastBodyElement();
+          if (firstBodyElement != null && lastBodyElement != null) {
+            ct.markRangeUnchanged(firstBodyElement, lastBodyElement);
+          }
+          children = ((PsiBlockStatement)body).getCodeBlock().getChildren();
+          // Skip {braces}
+          children = Arrays.copyOfRange(children, 1, children.length-1);
+        } else {
+          ct.markUnchanged(loop.getBody());
+          children = new PsiElement[]{body};
         }
-        children = ((PsiBlockStatement)body).getCodeBlock().getChildren();
-        // Skip {braces}
-        children = Arrays.copyOfRange(children, 1, children.length-1);
       } else {
         ct.markUnchanged(loop.getBody());
         children = new PsiElement[]{body};
@@ -247,6 +251,25 @@ public class UnrollLoopAction extends PsiUpdateModCommandAction<PsiLoopStatement
     ct.deleteAndRestoreComments(loop);
   }
 
+  private static boolean canUnwrapBlock(@NotNull PsiCodeBlock block, PsiCodeBlock parentBlock, List<PsiExpression> expressions) {
+    for (PsiStatement statement : block.getStatements()) {
+      if (statement instanceof PsiDeclarationStatement declaration) {
+        if (expressions.size() > 1) return false;
+        for (PsiElement element : declaration.getDeclaredElements()) {
+          if (element instanceof PsiVariable variable) {
+            String name = variable.getName();
+            if (name != null) {
+              if (PsiResolveHelper.getInstance(block.getProject()).resolveReferencedVariable(name, parentBlock) != null) {
+                return false;
+              }
+            }
+          }
+        }
+      }
+    }
+    return true;
+  }
+
   private static boolean isLoopBreak(PsiStatement statement) {
     if (!(statement instanceof PsiIfStatement ifStatement)) return false;
     if (ifStatement.getElseBranch() != null || ifStatement.getCondition() == null) return false;
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/afterDeclareConflict.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/afterDeclareConflict.java
new file mode 100644
index 000000000000..459c6839e181
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/afterDeclareConflict.java
@@ -0,0 +1,10 @@
+// "Unroll loop" "true-preview"
+class X {
+  void test() {
+      {
+          int y = 1 + 1;
+          System.out.println(y);
+      }
+      int y = 2;
+  }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/afterDeclareOne.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/afterDeclareOne.java
new file mode 100644
index 000000000000..f98afa1af23e
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/afterDeclareOne.java
@@ -0,0 +1,7 @@
+// "Unroll loop" "true-preview"
+class X {
+  void test() {
+      int y = 1 + 1;
+      System.out.println(y);
+  }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/afterDeclareTwo.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/afterDeclareTwo.java
new file mode 100644
index 000000000000..14eda8faec44
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/afterDeclareTwo.java
@@ -0,0 +1,13 @@
+// "Unroll loop" "true-preview"
+class X {
+  void test() {
+      {
+          int y = 1 + 1;
+          System.out.println(y);
+      }
+      {
+          int y = 2 + 1;
+          System.out.println(y);
+      }
+  }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/beforeDeclareConflict.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/beforeDeclareConflict.java
new file mode 100644
index 000000000000..1fe09b338469
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/beforeDeclareConflict.java
@@ -0,0 +1,10 @@
+// "Unroll loop" "true-preview"
+class X {
+  void test() {
+    <caret>for (int x : new int[]{1}) {
+      int y = x + 1;
+      System.out.println(y);
+    }
+    int y = 2;
+  }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/beforeDeclareOne.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/beforeDeclareOne.java
new file mode 100644
index 000000000000..f020e4b8f21e
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/beforeDeclareOne.java
@@ -0,0 +1,9 @@
+// "Unroll loop" "true-preview"
+class X {
+  void test() {
+    <caret>for (int x : new int[]{1}) {
+      int y = x + 1;
+      System.out.println(y);
+    }
+  }
+}
\ No newline at end of file
diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/beforeDeclareTwo.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/beforeDeclareTwo.java
new file mode 100644
index 000000000000..69611aff7371
--- /dev/null
+++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/unrollLoop/beforeDeclareTwo.java
@@ -0,0 +1,9 @@
+// "Unroll loop" "true-preview"
+class X {
+  void test() {
+    <caret>for (int x : new int[]{1, 2}) {
+      int y = x + 1;
+      System.out.println(y);
+    }
+  }
+}
\ No newline at end of file