[kotlin] K2: add inlay hints for lambda context parameters

The new hint is enabled and disabled together with the hints for this and implicit it. KTIJ-34412 GitOrigin-RevId: ab9865b2860209e81dc96f3cceeb0df6e1ff5e0f
2026-03-22 15:19:59 +07:00 · 2025-06-04 19:23:16 +03:00
parent d7c09230a6
commit 60ac1284df
12 changed files with 292 additions and 108 deletions
--- a/plugins/kotlin/code-insight/kotlin.code-insight.k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/hints/KtLambdasHintsProvider.kt
+++ b/plugins/kotlin/code-insight/kotlin.code-insight.k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/hints/KtLambdasHintsProvider.kt
@@ -6,9 +6,12 @@ import com.intellij.codeInsight.hints.filtering.Matcher
 import com.intellij.codeInsight.hints.filtering.MatcherConstructor
 import com.intellij.psi.PsiElement
 import com.intellij.psi.createSmartPointer
+import org.jetbrains.kotlin.analysis.api.KaExperimentalApi
+import org.jetbrains.kotlin.analysis.api.KaSession
 import org.jetbrains.kotlin.analysis.api.analyze
 import org.jetbrains.kotlin.analysis.api.resolution.successfulFunctionCallOrNull
 import org.jetbrains.kotlin.analysis.api.resolution.symbol
+import org.jetbrains.kotlin.analysis.api.symbols.KaAnonymousFunctionSymbol
 import org.jetbrains.kotlin.idea.codeInsight.hints.SHOW_IMPLICIT_RECEIVERS_AND_PARAMS
 import org.jetbrains.kotlin.idea.codeInsight.hints.SHOW_RETURN_EXPRESSIONS
 import org.jetbrains.kotlin.idea.codeInsight.hints.isFollowedByNewLine
@@ -16,6 +19,7 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
+import org.jetbrains.kotlin.utils.addToStdlib.ifNotEmpty
 import kotlin.contracts.ExperimentalContracts
 import kotlin.contracts.contract

@@ -104,30 +108,76 @@ class KtLambdasHintsProvider : AbstractKtInlayHintsProvider() {
        sink.whenOptionEnabled(SHOW_IMPLICIT_RECEIVERS_AND_PARAMS.name) {
            analyze(functionLiteral) {
                val anonymousFunctionSymbol = functionLiteral.symbol
-                anonymousFunctionSymbol.receiverParameter?.let { receiverSymbol ->
-                    val skipped = functionLiteral.getParentOfType<KtCallExpression>(false, KtBlockExpression::class.java)?.let { callExpression ->
-                        val functionCall = callExpression.resolveToCall()?.successfulFunctionCallOrNull() ?: return@let true
-                        val functionSymbol = functionCall.symbol
-                        functionSymbol.isExcludeListed(excludeListMatchers)
-                    }
+                printContextParameters(lambdaExpression, anonymousFunctionSymbol, sink)
+                printReceiverParameter(lambdaExpression, anonymousFunctionSymbol, sink)
+                printImplicitIt(lambdaExpression, anonymousFunctionSymbol, sink)
+            }
+        }
+    }

-                    if (skipped != true) {
-                        sink.addPresentation(InlineInlayPosition(lbrace.textRange.endOffset, true), hintFormat = HintFormat.default) {
-                            text("this: ")
-                            printKtType(receiverSymbol.returnType)
-                        }
-                    }
+    @OptIn(KaExperimentalApi::class)
+    private fun KaSession.printContextParameters(
+        lambdaExpression: KtLambdaExpression,
+        anonymousFunctionSymbol: KaAnonymousFunctionSymbol,
+        sink: InlayTreeSink,
+    ) {
+        anonymousFunctionSymbol.contextParameters.ifNotEmpty {
+            val contextParameters = this
+            sink.addPresentation(
+                position = InlineInlayPosition(lambdaExpression.leftCurlyBrace.textRange.endOffset, true),
+                hintFormat = HintFormat.default,
+            ) {
+                text("context(")
+                for (contextParameter in contextParameters.dropLast(1)) {
+                    printKtType(contextParameter.returnType)
+                    text(", ")
+                }
+                printKtType(contextParameters.last().returnType)
+                text(")")
+            }
+        }
+    }
+
+    private fun KaSession.printReceiverParameter(
+        lambdaExpression: KtLambdaExpression,
+        anonymousFunctionSymbol: KaAnonymousFunctionSymbol,
+        sink: InlayTreeSink,
+    ) {
+        anonymousFunctionSymbol.receiverParameter?.let { receiverSymbol ->
+            val skipped = lambdaExpression.functionLiteral.getParentOfType<KtCallExpression>(false, KtBlockExpression::class.java)
+                ?.let { callExpression ->
+                    val functionCall = callExpression.resolveToCall()?.successfulFunctionCallOrNull() ?: return@let true
+                    val functionSymbol = functionCall.symbol
+                    functionSymbol.isExcludeListed(excludeListMatchers)
                }

-                anonymousFunctionSymbol.valueParameters.singleOrNull()?.let { singleParameterSymbol ->
-                    val type = singleParameterSymbol.takeIf { it.isImplicitLambdaParameter }
-                        ?.returnType?.takeUnless { it.isUnitType } ?: return@let
-                    sink.addPresentation(InlineInlayPosition(lbrace.textRange.endOffset, true), hintFormat = HintFormat.default) {
-                        text("it: ")
-                        printKtType(type)
-                    }
+            if (skipped != true) {
+                sink.addPresentation(
+                    position = InlineInlayPosition(lambdaExpression.leftCurlyBrace.textRange.endOffset, true),
+                    hintFormat = HintFormat.default,
+                ) {
+                    text("this: ")
+                    printKtType(receiverSymbol.returnType)
                }
            }
        }
    }
+
+    private fun KaSession.printImplicitIt(
+        lambdaExpression: KtLambdaExpression,
+        anonymousFunctionSymbol: KaAnonymousFunctionSymbol,
+        sink: InlayTreeSink,
+    ) {
+        anonymousFunctionSymbol.valueParameters.singleOrNull()?.let { singleParameterSymbol ->
+            val type = singleParameterSymbol.takeIf { it.isImplicitLambdaParameter }
+                ?.returnType?.takeUnless { it.isUnitType } ?: return@let
+            sink.addPresentation(
+                position = InlineInlayPosition(lambdaExpression.leftCurlyBrace.textRange.endOffset, true),
+                hintFormat = HintFormat.default
+            ) {
+                text("it: ")
+                printKtType(type)
+            }
+        }
+    }
 }
--- a/plugins/kotlin/code-insight/kotlin.code-insight.k2/test/org/jetbrains/kotlin/idea/k2/hints/KtLambdasHintsProviderGenerated.java
+++ b/plugins/kotlin/code-insight/kotlin.code-insight.k2/test/org/jetbrains/kotlin/idea/k2/hints/KtLambdasHintsProviderGenerated.java
@@ -19,114 +19,172 @@ import org.junit.runner.RunWith;
@TestDataPath("$CONTENT_ROOT")
@RunWith(JUnit3RunnerWithInners.class)
@TestMetadata("../../idea/tests/testData/codeInsight/hints/lambda")
-public class KtLambdasHintsProviderGenerated extends AbstractKtLambdasHintsProvider {
-    @java.lang.Override
-    @org.jetbrains.annotations.NotNull
-    public final KotlinPluginMode getPluginMode() {
-        return KotlinPluginMode.K2;
+public abstract class KtLambdasHintsProviderGenerated extends AbstractKtLambdasHintsProvider {
+    @RunWith(JUnit3RunnerWithInners.class)
+    @TestMetadata("../../idea/tests/testData/codeInsight/hints/lambda/context")
+    public static class Context extends AbstractKtLambdasHintsProvider {
+        @java.lang.Override
+        @org.jetbrains.annotations.NotNull
+        public final KotlinPluginMode getPluginMode() {
+            return KotlinPluginMode.K2;
+        }
+
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+        }
+
+        @TestMetadata("ContextAndExplicitParameter.kt")
+        public void testContextAndExplicitParameter() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/context/ContextAndExplicitParameter.kt");
+        }
+
+        @TestMetadata("ContextAndIt.kt")
+        public void testContextAndIt() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/context/ContextAndIt.kt");
+        }
+
+        @TestMetadata("ContextAndThis.kt")
+        public void testContextAndThis() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/context/ContextAndThis.kt");
+        }
+
+        @TestMetadata("ContextNoNewLine.kt")
+        public void testContextNoNewLine() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/context/ContextNoNewLine.kt");
+        }
+
+        @TestMetadata("ContextSingle.kt")
+        public void testContextSingle() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/context/ContextSingle.kt");
+        }
+
+        @TestMetadata("ContextThisAndExplicitParameter.kt")
+        public void testContextThisAndExplicitParameter() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/context/ContextThisAndExplicitParameter.kt");
+        }
+
+        @TestMetadata("ContextThisAndIt.kt")
+        public void testContextThisAndIt() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/context/ContextThisAndIt.kt");
+        }
+
+        @TestMetadata("ContextTriple.kt")
+        public void testContextTriple() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/context/ContextTriple.kt");
+        }
    }

-    private void runTest(String testDataFilePath) throws Exception {
-        KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
-    }
+    @RunWith(JUnit3RunnerWithInners.class)
+    @TestMetadata("../../idea/tests/testData/codeInsight/hints/lambda")
+    public static class Uncategorized extends AbstractKtLambdasHintsProvider {
+        @java.lang.Override
+        @org.jetbrains.annotations.NotNull
+        public final KotlinPluginMode getPluginMode() {
+            return KotlinPluginMode.K2;
+        }

-    @TestMetadata("AnnotatedStatement.kt")
-    public void testAnnotatedStatement() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/AnnotatedStatement.kt");
-    }
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+        }

-    @TestMetadata("DisabledHints.kt")
-    public void testDisabledHints() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/DisabledHints.kt");
-    }
+        @TestMetadata("AnnotatedStatement.kt")
+        public void testAnnotatedStatement() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/AnnotatedStatement.kt");
+        }

-    @TestMetadata("Elvis.kt")
-    public void testElvis() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/Elvis.kt");
-    }
+        @TestMetadata("DisabledHints.kt")
+        public void testDisabledHints() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/DisabledHints.kt");
+        }

-    @TestMetadata("If.kt")
-    public void testIf() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/If.kt");
-    }
+        @TestMetadata("Elvis.kt")
+        public void testElvis() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/Elvis.kt");
+        }

-    @TestMetadata("IfBranchValue.kt")
-    public void testIfBranchValue() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/IfBranchValue.kt");
-    }
+        @TestMetadata("If.kt")
+        public void testIf() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/If.kt");
+        }

-    @TestMetadata("ImplicitIt.kt")
-    public void testImplicitIt() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/ImplicitIt.kt");
-    }
+        @TestMetadata("IfBranchValue.kt")
+        public void testIfBranchValue() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/IfBranchValue.kt");
+        }

-    @TestMetadata("ImplicitSingleLine.kt")
-    public void testImplicitSingleLine() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/ImplicitSingleLine.kt");
-    }
+        @TestMetadata("ImplicitIt.kt")
+        public void testImplicitIt() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/ImplicitIt.kt");
+        }

-    @TestMetadata("ImplicitThis.kt")
-    public void testImplicitThis() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/ImplicitThis.kt");
-    }
+        @TestMetadata("ImplicitSingleLine.kt")
+        public void testImplicitSingleLine() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/ImplicitSingleLine.kt");
+        }

-    @TestMetadata("ImplicitThisAndIt.kt")
-    public void testImplicitThisAndIt() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/ImplicitThisAndIt.kt");
-    }
+        @TestMetadata("ImplicitThis.kt")
+        public void testImplicitThis() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/ImplicitThis.kt");
+        }

-    @TestMetadata("ImplicitThisAndNamedSingleParameter.kt")
-    public void testImplicitThisAndNamedSingleParameter() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/ImplicitThisAndNamedSingleParameter.kt");
-    }
+        @TestMetadata("ImplicitThisAndIt.kt")
+        public void testImplicitThisAndIt() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/ImplicitThisAndIt.kt");
+        }

-    @TestMetadata("Label.kt")
-    public void testLabel() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/Label.kt");
-    }
+        @TestMetadata("ImplicitThisAndNamedSingleParameter.kt")
+        public void testImplicitThisAndNamedSingleParameter() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/ImplicitThisAndNamedSingleParameter.kt");
+        }

-    @TestMetadata("LabeledStatement.kt")
-    public void testLabeledStatement() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/LabeledStatement.kt");
-    }
+        @TestMetadata("Label.kt")
+        public void testLabel() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/Label.kt");
+        }

-    @TestMetadata("Nested.kt")
-    public void testNested() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/Nested.kt");
-    }
+        @TestMetadata("LabeledStatement.kt")
+        public void testLabeledStatement() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/LabeledStatement.kt");
+        }

-    @TestMetadata("NoHintForSingleExpression.kt")
-    public void testNoHintForSingleExpression() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/NoHintForSingleExpression.kt");
-    }
+        @TestMetadata("Nested.kt")
+        public void testNested() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/Nested.kt");
+        }

-    @TestMetadata("OneLineIf.kt")
-    public void testOneLineIf() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/OneLineIf.kt");
-    }
+        @TestMetadata("NoHintForSingleExpression.kt")
+        public void testNoHintForSingleExpression() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/NoHintForSingleExpression.kt");
+        }

-    @TestMetadata("PostfixPrefixExpr.kt")
-    public void testPostfixPrefixExpr() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/PostfixPrefixExpr.kt");
-    }
+        @TestMetadata("OneLineIf.kt")
+        public void testOneLineIf() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/OneLineIf.kt");
+        }

-    @TestMetadata("Qualified.kt")
-    public void testQualified() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/Qualified.kt");
-    }
+        @TestMetadata("PostfixPrefixExpr.kt")
+        public void testPostfixPrefixExpr() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/PostfixPrefixExpr.kt");
+        }

-    @TestMetadata("ReturnFunType.kt")
-    public void testReturnFunType() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/ReturnFunType.kt");
-    }
+        @TestMetadata("Qualified.kt")
+        public void testQualified() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/Qualified.kt");
+        }

-    @TestMetadata("SimpleCase.kt")
-    public void testSimpleCase() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/SimpleCase.kt");
-    }
+        @TestMetadata("ReturnFunType.kt")
+        public void testReturnFunType() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/ReturnFunType.kt");
+        }

-    @TestMetadata("When.kt")
-    public void testWhen() throws Exception {
-        runTest("../../idea/tests/testData/codeInsight/hints/lambda/When.kt");
+        @TestMetadata("SimpleCase.kt")
+        public void testSimpleCase() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/SimpleCase.kt");
+        }
+
+        @TestMetadata("When.kt")
+        public void testWhen() throws Exception {
+            runTest("../../idea/tests/testData/codeInsight/hints/lambda/When.kt");
+        }
    }
 }
--- a/plugins/kotlin/idea/tests/test/org/jetbrains/kotlin/idea/codeInsight/hints/AbstractKotlinInlayHintsProviderTest.kt
+++ b/plugins/kotlin/idea/tests/test/org/jetbrains/kotlin/idea/codeInsight/hints/AbstractKotlinInlayHintsProviderTest.kt
@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.idea.test.ExpectedPluginModeProvider
 import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescriptor
 import org.jetbrains.kotlin.idea.test.runAll
 import org.jetbrains.kotlin.idea.test.setUpWithKotlinPlugin
+import org.jetbrains.kotlin.idea.test.withCustomCompilerOptions
 import java.io.File

 abstract class AbstractKotlinInlayHintsProviderTest : DeclarativeInlayHintsProviderTestCase(),
@@ -92,7 +93,9 @@ abstract class AbstractKotlinInlayHintsProviderTest : DeclarativeInlayHintsProvi
        val options: Map<String, Boolean> = calculateOptions(fileContents)

        try {
-            doTestProvider("${file.name.substringBefore(".")}.kt", fileContents, inlayHintsProvider, options, file)
+            withCustomCompilerOptions(fileContents, project, module) {
+                doTestProvider("${file.name.substringBefore(".")}.kt", fileContents, inlayHintsProvider, options, file)
+            }
        } catch (e: FileComparisonFailedError) {
            throw FileComparisonFailedError(
                e.message,
--- a/plugins/kotlin/idea/tests/testData/codeInsight/hints/lambda/context/ContextAndExplicitParameter.kt
+++ b/plugins/kotlin/idea/tests/testData/codeInsight/hints/lambda/context/ContextAndExplicitParameter.kt
@@ -0,0 +1,9 @@
+// MODE: receivers_params
+// COMPILER_ARGUMENTS: -Xcontext-parameters
+
+val v = myWithContext("foo", 1) {/*<# context(|[kotlin.String:kotlin.fqn.class]String|) #>*/ param ->
+}
+
+fun <C, I> myWithContext(c: C, i: I, block: context(C)(I) -> Unit) {
+    with(c) { block(i) }
+}
--- a/plugins/kotlin/idea/tests/testData/codeInsight/hints/lambda/context/ContextAndIt.kt
+++ b/plugins/kotlin/idea/tests/testData/codeInsight/hints/lambda/context/ContextAndIt.kt
@@ -0,0 +1,9 @@
+// MODE: receivers_params
+// COMPILER_ARGUMENTS: -Xcontext-parameters
+
+val v = myWithContext("foo", 1) {/*<# context(|[kotlin.String:kotlin.fqn.class]String|) #>*//*<# it: |[kotlin.Int:kotlin.fqn.class]Int #>*/
+}
+
+fun <C, I> myWithContext(c: C, i: I, block: context(C)(I) -> Unit) {
+    with(c) { block(i) }
+}
--- a/plugins/kotlin/idea/tests/testData/codeInsight/hints/lambda/context/ContextAndThis.kt
+++ b/plugins/kotlin/idea/tests/testData/codeInsight/hints/lambda/context/ContextAndThis.kt
@@ -0,0 +1,9 @@
+// MODE: receivers_params
+// COMPILER_ARGUMENTS: -Xcontext-parameters
+
+val v = myWithContext("foo", 1) {/*<# context(|[kotlin.String:kotlin.fqn.class]String|) #>*//*<# this: |[kotlin.Int:kotlin.fqn.class]Int #>*/
+}
+
+fun <C, R> myWithContext(c: C, r: R, block: context(C) R.() -> Unit) {
+    with(c) { r.block() }
+}
--- a/plugins/kotlin/idea/tests/testData/codeInsight/hints/lambda/context/ContextNoNewLine.kt
+++ b/plugins/kotlin/idea/tests/testData/codeInsight/hints/lambda/context/ContextNoNewLine.kt
@@ -0,0 +1,8 @@
+// MODE: receivers_params
+// COMPILER_ARGUMENTS: -Xcontext-parameters
+
+val v = myWithContext("foo") { }
+
+fun <C> myWithContext(context: C, block: context(C)() -> Unit) {
+    with(context) { block() }
+}
--- a/plugins/kotlin/idea/tests/testData/codeInsight/hints/lambda/context/ContextSingle.kt
+++ b/plugins/kotlin/idea/tests/testData/codeInsight/hints/lambda/context/ContextSingle.kt
@@ -0,0 +1,9 @@
+// MODE: receivers_params
+// COMPILER_ARGUMENTS: -Xcontext-parameters
+
+val v = myWithContext("foo") {/*<# context(|[kotlin.String:kotlin.fqn.class]String|) #>*/
+}
+
+fun <C> myWithContext(context: C, block: context(C)() -> Unit) {
+    with(context) { block() }
+}
--- a/plugins/kotlin/idea/tests/testData/codeInsight/hints/lambda/context/ContextThisAndExplicitParameter.kt
+++ b/plugins/kotlin/idea/tests/testData/codeInsight/hints/lambda/context/ContextThisAndExplicitParameter.kt
@@ -0,0 +1,9 @@
+// MODE: receivers_params
+// COMPILER_ARGUMENTS: -Xcontext-parameters
+
+val v = myWithContext("foo", 1, 2.0) {/*<# context(|[kotlin.String:kotlin.fqn.class]String|) #>*//*<# this: |[kotlin.Int:kotlin.fqn.class]Int #>*/ param ->
+}
+
+fun <C, R, I> myWithContext(c: C, r: R, i: I, block: context(C) R.(I) -> Unit) {
+    with(c) { r.block(i) }
+}
--- a/plugins/kotlin/idea/tests/testData/codeInsight/hints/lambda/context/ContextThisAndIt.kt
+++ b/plugins/kotlin/idea/tests/testData/codeInsight/hints/lambda/context/ContextThisAndIt.kt
@@ -0,0 +1,9 @@
+// MODE: receivers_params
+// COMPILER_ARGUMENTS: -Xcontext-parameters
+
+val v = myWithContext("foo", 1, 2.0) {/*<# context(|[kotlin.String:kotlin.fqn.class]String|) #>*//*<# this: |[kotlin.Int:kotlin.fqn.class]Int #>*//*<# it: |[kotlin.Double:kotlin.fqn.class]Double #>*/
+}
+
+fun <C, R, I> myWithContext(c: C, r: R, i: I, block: context(C) R.(I) -> Unit) {
+    with(c) { r.block(i) }
+}
--- a/plugins/kotlin/idea/tests/testData/codeInsight/hints/lambda/context/ContextTriple.kt
+++ b/plugins/kotlin/idea/tests/testData/codeInsight/hints/lambda/context/ContextTriple.kt
@@ -0,0 +1,9 @@
+// MODE: receivers_params
+// COMPILER_ARGUMENTS: -Xcontext-parameters
+
+val v = myWithContext("foo", 1, 2.0) {/*<# context(|[kotlin.String:kotlin.fqn.class]String|, |[kotlin.Int:kotlin.fqn.class]Int|, |[kotlin.Double:kotlin.fqn.class]Double|) #>*/
+}
+
+fun <T1, T2, T3> myWithContext(t1: T1, t2: T2, t3: T3, block: context(T1, T2, T3)() -> Unit) {
+    with(t1, t2, t3) { block() }
+}
--- a/plugins/kotlin/util/test-generator-fe10/test/org/jetbrains/kotlin/fe10/testGenerator/Fe10GenerateTests.kt
+++ b/plugins/kotlin/util/test-generator-fe10/test/org/jetbrains/kotlin/fe10/testGenerator/Fe10GenerateTests.kt
@@ -1100,7 +1100,9 @@ private fun assembleWorkspace(): TWorkspace = workspace(KotlinPluginMode.K1) {
        }

        testClass<AbstractKotlinLambdasHintsProvider> {
-            model("codeInsight/hints/lambda")
+            model("codeInsight/hints/lambda", excludedDirectories = listOf(
+                "context", // K2
+            ))
        }
        testClass<AbstractKotlinValuesHintsProviderTest> {
            model("codeInsight/hints/values")