From ea1ff1d62ca96ea4788923d317d2ca67ba6096d8 Mon Sep 17 00:00:00 2001
From: Frederik Haselmeier <frederik.haselmeier@jetbrains.com>
Date: Tue, 26 Mar 2024 12:33:44 +0100
Subject: [PATCH] [Kotlin] Fixed for loop declarations leaking to outer scope
 when using PSI-only based completion

KTIJ-19995

GitOrigin-RevId: e01effa55285da3258b029ce4be963a7a040db2e
---
 .../idea/completion/implCommon/PsiTreeCompletion.kt    |  4 ++++
 .../testData/dumb/ForLoopDeclarationInBody.kt          |  8 ++++++++
 .../testData/dumb/ForLoopDeclarationShouldNotLeak.kt   |  9 +++++++++
 .../completion/test/DumbCompletionTestGenerated.java   | 10 ++++++++++
 .../fir/completion/FirDumbCompletionTestGenerated.java | 10 ++++++++++
 5 files changed, 41 insertions(+)
 create mode 100644 plugins/kotlin/completion/testData/dumb/ForLoopDeclarationInBody.kt
 create mode 100644 plugins/kotlin/completion/testData/dumb/ForLoopDeclarationShouldNotLeak.kt

diff --git a/plugins/kotlin/completion/impl-shared/src/org/jetbrains/kotlin/idea/completion/implCommon/PsiTreeCompletion.kt b/plugins/kotlin/completion/impl-shared/src/org/jetbrains/kotlin/idea/completion/implCommon/PsiTreeCompletion.kt
index 6b522a2b5566..03acc5f8fa42 100644
--- a/plugins/kotlin/completion/impl-shared/src/org/jetbrains/kotlin/idea/completion/implCommon/PsiTreeCompletion.kt
+++ b/plugins/kotlin/completion/impl-shared/src/org/jetbrains/kotlin/idea/completion/implCommon/PsiTreeCompletion.kt
@@ -238,6 +238,10 @@ class PsiTreeCompletion {
             return super.visitClassBody(classBody, data.registerSubScope(classBody.textRange, scopeType))
         }
 
+        override fun visitForExpression(expression: KtForExpression, data: LookupScope): Void? {
+            return super.visitForExpression(expression, data.registerSubScope(expression.textRange, ScopeType.LOCAL))
+        }
+
         override fun visitNamedFunction(function: KtNamedFunction, data: LookupScope): Void? {
             // We handle this declaration separately because it needs to open a new scope for the
             // function parameters, so they do not leak to the outside scope.
diff --git a/plugins/kotlin/completion/testData/dumb/ForLoopDeclarationInBody.kt b/plugins/kotlin/completion/testData/dumb/ForLoopDeclarationInBody.kt
new file mode 100644
index 000000000000..92c012aaf30e
--- /dev/null
+++ b/plugins/kotlin/completion/testData/dumb/ForLoopDeclarationInBody.kt
@@ -0,0 +1,8 @@
+fun test() {
+    for (prefixTest in test) {
+        val a = prefix<caret>
+    }
+}
+
+// EXIST: prefixTest
+// NOTHING_ELSE
\ No newline at end of file
diff --git a/plugins/kotlin/completion/testData/dumb/ForLoopDeclarationShouldNotLeak.kt b/plugins/kotlin/completion/testData/dumb/ForLoopDeclarationShouldNotLeak.kt
new file mode 100644
index 000000000000..9928b837b3e5
--- /dev/null
+++ b/plugins/kotlin/completion/testData/dumb/ForLoopDeclarationShouldNotLeak.kt
@@ -0,0 +1,9 @@
+fun test() {
+    for (prefixTest in test) {
+
+    }
+    val a = prefix<caret>
+}
+
+// ABSENT: prefixTest
+// NOTHING_ELSE
\ No newline at end of file
diff --git a/plugins/kotlin/completion/tests-k1/test/org/jetbrains/kotlin/idea/completion/test/DumbCompletionTestGenerated.java b/plugins/kotlin/completion/tests-k1/test/org/jetbrains/kotlin/idea/completion/test/DumbCompletionTestGenerated.java
index 904c1a464de1..a09a28e3edb7 100644
--- a/plugins/kotlin/completion/tests-k1/test/org/jetbrains/kotlin/idea/completion/test/DumbCompletionTestGenerated.java
+++ b/plugins/kotlin/completion/tests-k1/test/org/jetbrains/kotlin/idea/completion/test/DumbCompletionTestGenerated.java
@@ -118,6 +118,16 @@ public class DumbCompletionTestGenerated extends AbstractDumbCompletionTest {
         runTest("../testData/dumb/ExtensionProperty.kt");
     }
 
+    @TestMetadata("ForLoopDeclarationInBody.kt")
+    public void testForLoopDeclarationInBody() throws Exception {
+        runTest("../testData/dumb/ForLoopDeclarationInBody.kt");
+    }
+
+    @TestMetadata("ForLoopDeclarationShouldNotLeak.kt")
+    public void testForLoopDeclarationShouldNotLeak() throws Exception {
+        runTest("../testData/dumb/ForLoopDeclarationShouldNotLeak.kt");
+    }
+
     @TestMetadata("Function.kt")
     public void testFunction() throws Exception {
         runTest("../testData/dumb/Function.kt");
diff --git a/plugins/kotlin/fir/tests/test/org/jetbrains/kotlin/idea/fir/completion/FirDumbCompletionTestGenerated.java b/plugins/kotlin/fir/tests/test/org/jetbrains/kotlin/idea/fir/completion/FirDumbCompletionTestGenerated.java
index 2dc5cbd17d46..d576eedef774 100644
--- a/plugins/kotlin/fir/tests/test/org/jetbrains/kotlin/idea/fir/completion/FirDumbCompletionTestGenerated.java
+++ b/plugins/kotlin/fir/tests/test/org/jetbrains/kotlin/idea/fir/completion/FirDumbCompletionTestGenerated.java
@@ -118,6 +118,16 @@ public class FirDumbCompletionTestGenerated extends AbstractFirDumbCompletionTes
         runTest("../../completion/testData/dumb/ExtensionProperty.kt");
     }
 
+    @TestMetadata("ForLoopDeclarationInBody.kt")
+    public void testForLoopDeclarationInBody() throws Exception {
+        runTest("../../completion/testData/dumb/ForLoopDeclarationInBody.kt");
+    }
+
+    @TestMetadata("ForLoopDeclarationShouldNotLeak.kt")
+    public void testForLoopDeclarationShouldNotLeak() throws Exception {
+        runTest("../../completion/testData/dumb/ForLoopDeclarationShouldNotLeak.kt");
+    }
+
     @TestMetadata("Function.kt")
     public void testFunction() throws Exception {
         runTest("../../completion/testData/dumb/Function.kt");