From ed84affe001ae15d912fbcc078bc9f66c6441d32 Mon Sep 17 00:00:00 2001
From: Tagir Valeev <tagir.valeev@jetbrains.com>
Date: Thu, 28 Nov 2024 16:14:38 +0100
Subject: [PATCH] [kotlin-dfa] Slight adjustments and more complex tests for
 IDEA-363059

GitOrigin-RevId: a874b1ff7471f6bd7f8880769ebb02f81267b755
---
 .../debugger/dfaassist/K2DfaAssistProvider.kt | 11 ++--
 .../idea/k2/debugger/test/K2DfaAssistTest.kt  | 58 +++++++++++++++++++
 2 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/plugins/kotlin/jvm-debugger/evaluation/k2/src/org/jetbrains/kotlin/idea/k2/debugger/dfaassist/K2DfaAssistProvider.kt b/plugins/kotlin/jvm-debugger/evaluation/k2/src/org/jetbrains/kotlin/idea/k2/debugger/dfaassist/K2DfaAssistProvider.kt
index 9e3e083ac608..a18420972768 100644
--- a/plugins/kotlin/jvm-debugger/evaluation/k2/src/org/jetbrains/kotlin/idea/k2/debugger/dfaassist/K2DfaAssistProvider.kt
+++ b/plugins/kotlin/jvm-debugger/evaluation/k2/src/org/jetbrains/kotlin/idea/k2/debugger/dfaassist/K2DfaAssistProvider.kt
@@ -97,9 +97,12 @@ class K2DfaAssistProvider : DfaAssistProvider {
         while(true) {
             val function = current.parentOfType<KtFunction>()
             if (function != null) {
-                val lambda = function.parent as? KtLambdaExpression
-                val arg = lambda?.parent as? KtLambdaArgument
-                val call = arg?.parent as? KtCallExpression
+                val realFun = function.parent as? KtLambdaExpression ?: function
+                val arg = realFun?.parent as? KtValueArgument
+                val call = when(arg) {
+                    is KtLambdaArgument -> arg.parent
+                    else -> arg?.parent?.parent
+                } as? KtCallExpression
                 if (call != null) {
                     val inline = analyze(call) {
                         val functionCall = call.resolveToCall()?.singleFunctionCallOrNull()
@@ -122,7 +125,6 @@ class K2DfaAssistProvider : DfaAssistProvider {
     ): Value? {
         val qualifier = dfaVar.qualifier
         val descriptor = dfaVar.descriptor
-        val scope = anchor.getScope()
         val variables = (proxy as StackFrameProxyImpl).visibleVariables()
         val inlineDepth = getInlineDepth(variables)
         val inlineSuffix = KotlinDebuggerConstants.INLINE_FUN_VAR_SUFFIX.repeat(inlineDepth)
@@ -201,6 +203,7 @@ class K2DfaAssistProvider : DfaAssistProvider {
                         var value: Value? = null
                         if (variable == null) {
                             val psi = symbol.psi
+                            val scope = anchor.getScope()
                             if (psi != null && scope != null && psi.containingFile == scope.containingFile && !scope.isAncestor(psi)) {
                                 // Captured variable
                                 val capturedName = AsmUtil.CAPTURED_PREFIX + name
diff --git a/plugins/kotlin/jvm-debugger/test/k2/test/org/jetbrains/kotlin/idea/k2/debugger/test/K2DfaAssistTest.kt b/plugins/kotlin/jvm-debugger/test/k2/test/org/jetbrains/kotlin/idea/k2/debugger/test/K2DfaAssistTest.kt
index 859155fa6d9b..8b9875c59b93 100644
--- a/plugins/kotlin/jvm-debugger/test/k2/test/org/jetbrains/kotlin/idea/k2/debugger/test/K2DfaAssistTest.kt
+++ b/plugins/kotlin/jvm-debugger/test/k2/test/org/jetbrains/kotlin/idea/k2/debugger/test/K2DfaAssistTest.kt
@@ -783,6 +783,64 @@ class K2DfaAssistTest : DfaAssistTest(), ExpectedPluginModeProvider {
         }
     }
 
+    fun testFunCapture() {
+        val text = """
+            fun main() {
+                test(1)
+            }
+            
+            fun test(i: Int) {
+                block2 {
+                    block(fun() {
+                        <caret>if (i > 0/*TRUE*/) {
+                            print("foo")
+                        }
+                    })
+                }
+            }
+            
+            fun block2(block: () -> Unit) {
+                block()
+            }
+            
+            inline fun block(block: () -> Unit) {
+                block()
+            }
+        """.trimIndent()
+        doTest(text) { vm, frame ->
+            frame.addVariable("\$i", MockIntegerValue(vm,  1))
+        }
+    }
+
+    fun testLambdaArgCapture() {
+        val text = """
+            fun main() {
+                test(1)
+            }
+            
+            fun test(i: Int) {
+                block2 {
+                    block({
+                        <caret>if (i > 0/*TRUE*/) {
+                            print("foo")
+                        }
+                    })
+                }
+            }
+            
+            fun block2(block: () -> Unit) {
+                block()
+            }
+            
+            inline fun block(block: () -> Unit) {
+                block()
+            }
+        """.trimIndent()
+        doTest(text) { vm, frame ->
+            frame.addVariable("\$i", MockIntegerValue(vm,  1))
+        }
+    }
+
     private fun doTest(text: String, mockValues: BiConsumer<MockVirtualMachine, MockStackFrame>) {
         doTest(text, mockValues, "Test.kt")
     }