[ClsStubs] Add context receivers on function types

analysis/decompiled/decompiler-to-file-stubs/testData/clsFileStubBuilder/ContextReceiversOnFunctionType/ContextReceiversOnFunctionType.kt

@@ -0,0 +1,12 @@
+// JVM_FILE_NAME: ContextReceiversOnFunctionTypeKt
+// !LANGUAGE: +ContextReceivers
+
+fun f(g: context(A, B) Int.(Int) -> Int) {}
+
+class A {
+    val valueA: Int = 10
+}
+
+class B {
+    val valueB: Int = 11
+}

analysis/decompiled/decompiler-to-file-stubs/testData/clsFileStubBuilder/ContextReceiversOnFunctionType/ContextReceiversOnFunctionType.txt

@@ -0,0 +1,41 @@
+PsiJetFileStubImpl[package=]
+  PACKAGE_DIRECTIVE
+  IMPORT_LIST
+  FUN[fqName=f, hasBlockBody=true, hasBody=true, hasTypeParameterListBeforeFunctionName=false, isExtension=false, isTopLevel=true, mayHaveContract=false, name=f]
+    MODIFIER_LIST[public]
+    VALUE_PARAMETER_LIST
+      VALUE_PARAMETER[fqName=null, hasDefaultValue=false, hasValOrVar=false, isMutable=false, name=g]
+        TYPE_REFERENCE
+          FUNCTION_TYPE
+            CONTEXT_RECEIVER_LIST
+              CONTEXT_RECEIVER[label=null]
+                TYPE_REFERENCE
+                  USER_TYPE
+                    REFERENCE_EXPRESSION[referencedName=A]
+              CONTEXT_RECEIVER[label=null]
+                TYPE_REFERENCE
+                  USER_TYPE
+                    REFERENCE_EXPRESSION[referencedName=B]
+            FUNCTION_TYPE_RECEIVER
+              TYPE_REFERENCE
+                USER_TYPE
+                  USER_TYPE
+                    REFERENCE_EXPRESSION[referencedName=kotlin]
+                  REFERENCE_EXPRESSION[referencedName=Int]
+            VALUE_PARAMETER_LIST
+              VALUE_PARAMETER[fqName=null, hasDefaultValue=false, hasValOrVar=false, isMutable=false, name=null]
+                TYPE_REFERENCE
+                  USER_TYPE
+                    USER_TYPE
+                      REFERENCE_EXPRESSION[referencedName=kotlin]
+                    REFERENCE_EXPRESSION[referencedName=Int]
+            TYPE_REFERENCE
+              USER_TYPE
+                USER_TYPE
+                  REFERENCE_EXPRESSION[referencedName=kotlin]
+                REFERENCE_EXPRESSION[referencedName=Int]
+    TYPE_REFERENCE
+      USER_TYPE
+        USER_TYPE
+          REFERENCE_EXPRESSION[referencedName=kotlin]
+        REFERENCE_EXPRESSION[referencedName=Unit]

analysis/decompiled/decompiler-to-file-stubs/tests/org/jetbrains/kotlin/analysis/decompiler/stub/files/ClsStubBuilderTestGenerated.java

@@ -108,6 +108,12 @@ public class ClsStubBuilderTestGenerated extends AbstractClsStubBuilderTest {
         runTest("analysis/decompiled/decompiler-to-file-stubs/testData/clsFileStubBuilder/ContextReceiversOnClass/");
     }
 
+    @Test
+    @TestMetadata("ContextReceiversOnFunctionType")
+    public void testContextReceiversOnFunctionType() throws Exception {
+        runTest("analysis/decompiled/decompiler-to-file-stubs/testData/clsFileStubBuilder/ContextReceiversOnFunctionType/");
+    }
+
     @Test
     @TestMetadata("ContextReceiversOnTopLevelCallables")
     public void testContextReceiversOnTopLevelCallables() throws Exception {

analysis/decompiled/decompiler-to-psi/tests/org/jetbrains/kotlin/analysis/decompiler/psi/ByDecompiledPsiStubBuilderTestGenerated.java

@@ -108,6 +108,12 @@ public class ByDecompiledPsiStubBuilderTestGenerated extends AbstractByDecompile
         runTest("analysis/decompiled/decompiler-to-file-stubs/testData/clsFileStubBuilder/ContextReceiversOnClass/");
     }
 
+    @Test
+    @TestMetadata("ContextReceiversOnFunctionType")
+    public void testContextReceiversOnFunctionType() throws Exception {
+        runTest("analysis/decompiled/decompiler-to-file-stubs/testData/clsFileStubBuilder/ContextReceiversOnFunctionType/");
+    }
+
     @Test
     @TestMetadata("ContextReceiversOnTopLevelCallables")
     public void testContextReceiversOnTopLevelCallables() throws Exception {

analysis/decompiled/decompiler-to-stubs/src/org/jetbrains/kotlin/analysis/decompiler/stub/TypeClsStubBuilder.kt

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.metadata.ProtoBuf.Type
 import org.jetbrains.kotlin.metadata.ProtoBuf.Type.Argument.Projection
 import org.jetbrains.kotlin.metadata.ProtoBuf.TypeParameter.Variance
 import org.jetbrains.kotlin.metadata.deserialization.*
+import org.jetbrains.kotlin.metadata.jvm.JvmProtoBuf
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -110,19 +111,29 @@ class TypeClsStubBuilder(private val c: ClsStubBuilderContext) {
                 it.classId.asSingleFqName() == StandardNames.FqNames.extensionFunctionType
             }
 
+            val (contextReceiverAnnotations, otherAnnotations) = notExtensionAnnotations.partition {
+                it.classId.asSingleFqName() == StandardNames.FqNames.contextFunctionTypeParams
+            }
+
             val isExtension = extensionAnnotations.isNotEmpty()
             val isSuspend = Flags.SUSPEND_TYPE.get(type.flags)
 
             val nullableWrapper = if (isSuspend) {
                 val wrapper = nullableTypeParent(parent, type)
-                createTypeAnnotationStubs(wrapper, type, notExtensionAnnotations)
+                createTypeAnnotationStubs(wrapper, type, otherAnnotations)
                 wrapper
             } else {
-                createTypeAnnotationStubs(parent, type, notExtensionAnnotations)
+                createTypeAnnotationStubs(parent, type, otherAnnotations)
                 nullableTypeParent(parent, type)
             }
 
-            createFunctionTypeStub(nullableWrapper, type, isExtension, isSuspend)
+            val numContextReceivers = if (contextReceiverAnnotations.isEmpty()) {
+                0
+            } else {
+                val argument = type.getExtension(JvmProtoBuf.typeAnnotation).find { c.nameResolver.getClassId(it.id).asSingleFqName() == StandardNames.FqNames.contextFunctionTypeParams }!!.getArgument(0)
+                argument.value.intValue.toInt()
+            }
+            createFunctionTypeStub(nullableWrapper, type, isExtension, isSuspend, numContextReceivers)
 
             return
         }
@@ -225,20 +236,33 @@ class TypeClsStubBuilder(private val c: ClsStubBuilderContext) {
         parent: StubElement<out PsiElement>,
         type: Type,
         isExtensionFunctionType: Boolean,
-        isSuspend: Boolean
+        isSuspend: Boolean,
+        numContextReceivers: Int,
     ) {
         val typeArgumentList = type.argumentList
         val functionType = KotlinPlaceHolderStubImpl<KtFunctionType>(parent, KtStubElementTypes.FUNCTION_TYPE)
+        var processedTypes = 0
+
+        if (numContextReceivers != 0) {
+            ContextReceiversListStubBuilder(c).createContextReceiverStubs(
+                functionType,
+                typeArgumentList.subList(
+                    processedTypes,
+                    processedTypes + numContextReceivers
+                ).map { it.type(c.typeTable)!! })
+            processedTypes += numContextReceivers
+        }
+
         if (isExtensionFunctionType) {
             val functionTypeReceiverStub =
                 KotlinPlaceHolderStubImpl<KtFunctionTypeReceiver>(functionType, KtStubElementTypes.FUNCTION_TYPE_RECEIVER)
-            val receiverTypeProto = typeArgumentList.first().type(c.typeTable)!!
+            val receiverTypeProto = typeArgumentList[processedTypes].type(c.typeTable)!!
             createTypeReferenceStub(functionTypeReceiverStub, receiverTypeProto)
+            processedTypes++
         }
 
         val parameterList = KotlinPlaceHolderStubImpl<KtParameterList>(functionType, KtStubElementTypes.VALUE_PARAMETER_LIST)
-        val typeArgumentsWithoutReceiverAndReturnType =
-            typeArgumentList.subList(if (isExtensionFunctionType) 1 else 0, typeArgumentList.size - 1)
+        val typeArgumentsWithoutReceiverAndReturnType = typeArgumentList.subList(processedTypes, typeArgumentList.size - 1)
         var suspendParameterType: Type? = null
 
         for ((index, argument) in typeArgumentsWithoutReceiverAndReturnType.withIndex()) {

compiler/psi/src/org/jetbrains/kotlin/psi/stubs/KotlinStubVersions.kt

@@ -23,12 +23,12 @@ object KotlinStubVersions {
     // Though only kotlin declarations (no code in the bodies) are stubbed, please do increase this version
     //      if you are not 100% sure it can be avoided.
     // Increasing this version will lead to reindexing of all kotlin source files on the first IDE startup with the new version.
-    const val SOURCE_STUB_VERSION = 148
+    const val SOURCE_STUB_VERSION = 149
 
     // Binary stub version should be increased if stub format (org.jetbrains.kotlin.psi.stubs.impl) is changed
     //      or changes are made to the core stub building code (org.jetbrains.kotlin.idea.decompiler.stubBuilder).
     // Increasing this version will lead to reindexing of all binary files that are potentially kotlin binaries (including all class files).
-    private const val BINARY_STUB_VERSION = 84
+    private const val BINARY_STUB_VERSION = 85
 
     // Classfile stub version should be increased if changes are made to classfile stub building subsystem (org.jetbrains.kotlin.idea.decompiler.classFile)
     // Increasing this version will lead to reindexing of all classfiles.