From 369c86ebb05d0c9f385fa9fb01aa4bfc139c2d7e Mon Sep 17 00:00:00 2001
From: Anastasiya Shadrina <anastasia.shadrina@jetbrains.com>
Date: Sun, 12 Sep 2021 23:49:50 +0700
Subject: [PATCH] [FE] Fix contextual functional types rendering

---
 .../fromKEEP/autoCloseScope.txt                |  2 +-
 .../contextReceivers/fromKEEP/clickHandler.txt |  4 ++--
 .../contextReceivers/functionalType.txt        |  8 ++++----
 .../kotlin/renderer/DescriptorRendererImpl.kt  | 18 +++++++++++++++---
 4 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/autoCloseScope.txt b/compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/autoCloseScope.txt
index 1b08c399033..dc0fc2f8c01 100644
--- a/compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/autoCloseScope.txt
+++ b/compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/autoCloseScope.txt
@@ -1,7 +1,7 @@
 package
 
 public fun test(): kotlin.Unit
-public fun withAutoClose(/*0*/ block: @kotlin.ContextFunctionTypeParams(count = 1) () -> kotlin.Unit): kotlin.Unit
+public fun withAutoClose(/*0*/ block: context(AutoCloseScope) () -> kotlin.Unit): kotlin.Unit
 public fun File.open(): InputStream
 
 public interface AutoCloseScope {
diff --git a/compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/clickHandler.txt b/compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/clickHandler.txt
index e968e693630..acdc000f929 100644
--- a/compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/clickHandler.txt
+++ b/compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/clickHandler.txt
@@ -1,6 +1,6 @@
 package
 
-public fun handleClick(/*0*/ clickHandler: ClickHandler /* = @kotlin.ContextFunctionTypeParams(count = 1) (ClickEvent) -> kotlin.Unit */): kotlin.Unit
+public fun handleClick(/*0*/ clickHandler: ClickHandler /* = context(Button) (ClickEvent) -> kotlin.Unit */): kotlin.Unit
 
 public final class Button {
     public constructor Button()
@@ -15,4 +15,4 @@ public final class ClickEvent {
     public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
     public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
 }
-public typealias ClickHandler = @kotlin.ContextFunctionTypeParams(count = 1) (ClickEvent) -> kotlin.Unit
+public typealias ClickHandler = context(Button) (ClickEvent) -> kotlin.Unit
\ No newline at end of file
diff --git a/compiler/testData/diagnostics/tests/extensions/contextReceivers/functionalType.txt b/compiler/testData/diagnostics/tests/extensions/contextReceivers/functionalType.txt
index 1c94626caf6..3721cf0b860 100644
--- a/compiler/testData/diagnostics/tests/extensions/contextReceivers/functionalType.txt
+++ b/compiler/testData/diagnostics/tests/extensions/contextReceivers/functionalType.txt
@@ -1,10 +1,10 @@
 package
 
-public fun f2(/*0*/ g: @kotlin.ContextFunctionTypeParams(count = 1) (Param) -> kotlin.Unit): kotlin.Unit
-public fun f4(/*0*/ g: @kotlin.ContextFunctionTypeParams(count = 1) () -> kotlin.Unit): kotlin.Unit
+public fun f2(/*0*/ g: context(C) (Param) -> kotlin.Unit): kotlin.Unit
+public fun f4(/*0*/ g: context(C) () -> kotlin.Unit): kotlin.Unit
 public fun test(): kotlin.Unit
-public fun R.f1(/*0*/ g: @kotlin.ContextFunctionTypeParams(count = 1) (R.(Param) -> kotlin.Unit)): kotlin.Unit
-public fun R.f3(/*0*/ g: @kotlin.ContextFunctionTypeParams(count = 1) (R.() -> kotlin.Unit)): kotlin.Unit
+public fun R.f1(/*0*/ g: context(C) R.(Param) -> kotlin.Unit): kotlin.Unit
+public fun R.f3(/*0*/ g: context(C) R.() -> kotlin.Unit): kotlin.Unit
 
 public final class C {
     public constructor C()
diff --git a/core/descriptors/src/org/jetbrains/kotlin/renderer/DescriptorRendererImpl.kt b/core/descriptors/src/org/jetbrains/kotlin/renderer/DescriptorRendererImpl.kt
index b2d28efc5b4..89938abb52f 100644
--- a/core/descriptors/src/org/jetbrains/kotlin/renderer/DescriptorRendererImpl.kt
+++ b/core/descriptors/src/org/jetbrains/kotlin/renderer/DescriptorRendererImpl.kt
@@ -36,7 +36,7 @@ internal class DescriptorRendererImpl(
 
     private val functionTypeAnnotationsRenderer: DescriptorRendererImpl by lazy {
         withOptions {
-            excludedTypeAnnotationClasses += listOf(StandardNames.FqNames.extensionFunctionType)
+            excludedTypeAnnotationClasses += listOf(StandardNames.FqNames.extensionFunctionType, StandardNames.FqNames.contextFunctionTypeParams)
         } as DescriptorRendererImpl
     }
 
@@ -317,15 +317,27 @@ internal class DescriptorRendererImpl(
 
     private fun StringBuilder.renderFunctionType(type: KotlinType) {
         val lengthBefore = length
-        // we need special renderer to skip @ExtensionFunctionType
+        // we need special renderer to skip @ExtensionFunctionType and @ContextFunctionTypeParams
         with(functionTypeAnnotationsRenderer) {
             renderAnnotations(type)
         }
         val hasAnnotations = length != lengthBefore
 
+        val receiverType = type.getReceiverTypeFromFunctionType()
+        val contextReceiversTypes = type.getContextReceiverTypesFromFunctionType()
+        if (contextReceiversTypes.isNotEmpty()) {
+            append("context(")
+            val withoutLast = contextReceiversTypes.subList(0, contextReceiversTypes.lastIndex)
+            for (contextReceiverType in withoutLast) {
+                renderNormalizedType(contextReceiverType)
+                append(", ")
+            }
+            renderNormalizedType(contextReceiversTypes.last())
+            append(") ")
+        }
+
         val isSuspend = type.isSuspendFunctionType
         val isNullable = type.isMarkedNullable
-        val receiverType = type.getReceiverTypeFromFunctionType()
 
         val needParenthesis = isNullable || (hasAnnotations && receiverType != null)
         if (needParenthesis) {