From da96a1ee3356f68dd356b076e0717dc3f11ab5a0 Mon Sep 17 00:00:00 2001
From: Andrey Sokolov <andrey.sokolov@jetbrains.com>
Date: Thu, 22 Sep 2022 18:37:24 +0200
Subject: [PATCH] Similar usages: Declaration statement in for statement should
 consider for statement as context

GitOrigin-RevId: dea5b3a87a8cc1de9888ebf3e155b4cce1690d53
---
 .../JavaUsageSimilarityFeaturesProvider.java  | 34 +++++++++++--------
 .../DeclarationInForStatement.java            |  7 ++++
 .../JavaUsagesBySimilarityTest.java           |  7 ++++
 3 files changed, 33 insertions(+), 15 deletions(-)
 create mode 100644 java/java-tests/testData/findSimilar/DeclarationInForStatement.java

diff --git a/java/java-impl/src/com/intellij/find/findUsages/similarity/JavaUsageSimilarityFeaturesProvider.java b/java/java-impl/src/com/intellij/find/findUsages/similarity/JavaUsageSimilarityFeaturesProvider.java
index 7b7e41138571..5b652f4046b3 100644
--- a/java/java-impl/src/com/intellij/find/findUsages/similarity/JavaUsageSimilarityFeaturesProvider.java
+++ b/java/java-impl/src/com/intellij/find/findUsages/similarity/JavaUsageSimilarityFeaturesProvider.java
@@ -37,21 +37,25 @@ public class JavaUsageSimilarityFeaturesProvider implements UsageSimilarityFeatu
   }
 
   public @Nullable PsiElement getContainingStatement(@NotNull PsiElement element) {
-    return PsiTreeUtil.getParentOfType(element, PsiDeclarationStatement.class,
-                                       PsiExpressionStatement.class,
-                                       PsiIfStatement.class,
-                                       PsiWhileStatement.class,
-                                       PsiTryStatement.class,
-                                       PsiThrowStatement.class,
-                                       PsiSwitchStatement.class,
-                                       PsiReturnStatement.class,
-                                       PsiLoopStatement.class,
-                                       PsiImportStatement.class,
-                                       PsiForStatement.class,
-                                       PsiForeachStatement.class,
-                                       PsiConditionalLoopStatement.class,
-                                       PsiBlockStatement.class,
-                                       PsiMethod.class);
+    PsiElement containingStatement = PsiTreeUtil.getParentOfType(element, PsiDeclarationStatement.class,
+                                                                 PsiExpressionStatement.class,
+                                                                 PsiIfStatement.class,
+                                                                 PsiWhileStatement.class,
+                                                                 PsiTryStatement.class,
+                                                                 PsiThrowStatement.class,
+                                                                 PsiSwitchStatement.class,
+                                                                 PsiReturnStatement.class,
+                                                                 PsiLoopStatement.class,
+                                                                 PsiImportStatement.class,
+                                                                 PsiForStatement.class,
+                                                                 PsiForeachStatement.class,
+                                                                 PsiConditionalLoopStatement.class,
+                                                                 PsiBlockStatement.class,
+                                                                 PsiMethod.class);
+    if (containingStatement instanceof PsiDeclarationStatement && containingStatement.getParent() instanceof PsiForStatement) {
+      return containingStatement.getParent();
+    }
+    return containingStatement;
   }
 
   private static @NotNull Bag getParentStatementFeatures(@NotNull PsiElement context) {
diff --git a/java/java-tests/testData/findSimilar/DeclarationInForStatement.java b/java/java-tests/testData/findSimilar/DeclarationInForStatement.java
new file mode 100644
index 000000000000..de7e2c3df09f
--- /dev/null
+++ b/java/java-tests/testData/findSimilar/DeclarationInForStatement.java
@@ -0,0 +1,7 @@
+class A {
+  public static void main(String[] args) {
+
+    List<String> modules = new ArrayList<String>();
+    for (int i = modules.s<caret>ize(); i < modules.size(); i++)
+  }
+}
\ No newline at end of file
diff --git a/java/java-tests/testSrc/com/intellij/find/groupSimilar/JavaUsagesBySimilarityTest.java b/java/java-tests/testSrc/com/intellij/find/groupSimilar/JavaUsagesBySimilarityTest.java
index 6f6e79817ab0..91d98396f227 100644
--- a/java/java-tests/testSrc/com/intellij/find/groupSimilar/JavaUsagesBySimilarityTest.java
+++ b/java/java-tests/testSrc/com/intellij/find/groupSimilar/JavaUsagesBySimilarityTest.java
@@ -89,6 +89,13 @@ public class JavaUsagesBySimilarityTest extends JavaCodeInsightFixtureTestCase {
     }
   }
 
+  public void testDeclarationInForStatement() {
+    myFixture.configureByFile("DeclarationInForStatement.java");
+    PsiElement elementAtCaret = myFixture.getReferenceAtCaretPosition().getElement();
+    final Bag features = new JavaUsageSimilarityFeaturesProvider().getFeatures(elementAtCaret);
+    assertEquals(1, features.get("FOR"));
+  }
+
   public void testBag() {
     final Bag bag = new Bag("a", "b");
     assertEquals("a : 1\n" +