From c5c7b01cd0cd6d429f2d29d98d7fc86313e21ae6 Mon Sep 17 00:00:00 2001
From: Roman Golyshev <roman.golyshev@jetbrains.com>
Date: Wed, 9 Apr 2025 12:24:23 +0200
Subject: [PATCH] KTIJ-33579 [kotlin] Use `ImportContextWithFixedReceiverType`
 in `ComponentFunctionImportQuickFixFactory` to properly handle destructuring
 in lambdas

^KTIJ-33579 Fixed

GitOrigin-RevId: b3f921bba5ef097f1998692aed8fac2e99c303b5
---
 .../ComponentFunctionImportQuickFixFactory.kt    | 16 +++++++++++++++-
 ...tionInLambdaParameterExtensionComponent1.test |  2 --
 ...tionInLambdaParameterExtensionComponent2.test |  2 --
 3 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/plugins/kotlin/code-insight/fixes-k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/fixes/imprt/factories/ComponentFunctionImportQuickFixFactory.kt b/plugins/kotlin/code-insight/fixes-k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/fixes/imprt/factories/ComponentFunctionImportQuickFixFactory.kt
index 99d36b8df74e..b019f956f54f 100644
--- a/plugins/kotlin/code-insight/fixes-k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/fixes/imprt/factories/ComponentFunctionImportQuickFixFactory.kt
+++ b/plugins/kotlin/code-insight/fixes-k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/fixes/imprt/factories/ComponentFunctionImportQuickFixFactory.kt
@@ -8,6 +8,7 @@ import org.jetbrains.kotlin.idea.base.analysis.api.utils.KtSymbolFromIndexProvid
 import org.jetbrains.kotlin.idea.k2.codeinsight.fixes.imprt.*
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtParameter
 
 internal object ComponentFunctionImportQuickFixFactory : AbstractImportQuickFixFactory() {
     override fun KaSession.detectPositionContext(diagnostic: KaDiagnosticWithPsi<*>): ImportContext? =
@@ -15,7 +16,20 @@ internal object ComponentFunctionImportQuickFixFactory : AbstractImportQuickFixF
             is KaFirDiagnostic.ComponentFunctionMissing,
             is KaFirDiagnostic.ComponentFunctionAmbiguity -> {
                 val destructuredExpression = diagnostic.psi as? KtExpression ?: return null
-                DefaultImportContext(destructuredExpression, ImportPositionTypeAndReceiver.OperatorCall(destructuredExpression))
+
+                val destructuredType = when (destructuredExpression) {
+                    // destructuring in lambda parameter position (e.g. `foo { (a, b) -> ... }`)
+                    is KtParameter -> destructuredExpression.returnType
+
+                    // regular assignment destructuring (e.g. `val (a, b) = ...`)
+                    else -> destructuredExpression.expressionType
+                } ?: return null
+
+                ImportContextWithFixedReceiverType(
+                    destructuredExpression,
+                    ImportPositionType.OperatorCall,
+                    explicitReceiverType = destructuredType,
+                )
             }
 
             else -> null
diff --git a/plugins/kotlin/idea/tests/testData/quickfix/autoImports/multiDeclarationInLambdaParameterExtensionComponent1.test b/plugins/kotlin/idea/tests/testData/quickfix/autoImports/multiDeclarationInLambdaParameterExtensionComponent1.test
index cf7647feb926..41b346347a58 100644
--- a/plugins/kotlin/idea/tests/testData/quickfix/autoImports/multiDeclarationInLambdaParameterExtensionComponent1.test
+++ b/plugins/kotlin/idea/tests/testData/quickfix/autoImports/multiDeclarationInLambdaParameterExtensionComponent1.test
@@ -1,7 +1,6 @@
 // FILE: first.before.kt
 // "Import operator 'Some.component1'" "true"
 // IGNORE_K1
-// IGNORE_K2
 
 package testing
 
@@ -26,7 +25,6 @@ operator fun Some.component2() = 3
 // FILE: first.after.kt
 // "Import operator 'Some.component1'" "true"
 // IGNORE_K1
-// IGNORE_K2
 
 package testing
 
diff --git a/plugins/kotlin/idea/tests/testData/quickfix/autoImports/multiDeclarationInLambdaParameterExtensionComponent2.test b/plugins/kotlin/idea/tests/testData/quickfix/autoImports/multiDeclarationInLambdaParameterExtensionComponent2.test
index 2fcf9a83e328..c26b30baf0dd 100644
--- a/plugins/kotlin/idea/tests/testData/quickfix/autoImports/multiDeclarationInLambdaParameterExtensionComponent2.test
+++ b/plugins/kotlin/idea/tests/testData/quickfix/autoImports/multiDeclarationInLambdaParameterExtensionComponent2.test
@@ -1,7 +1,6 @@
 // FILE: first.before.kt
 // "Import operator 'Some.component2'" "true"
 // IGNORE_K1
-// IGNORE_K2
 
 package testing
 
@@ -26,7 +25,6 @@ operator fun Some.component2() = 3
 // FILE: first.after.kt
 // "Import operator 'Some.component2'" "true"
 // IGNORE_K1
-// IGNORE_K2
 
 package testing