[stubs] fix cls stubs for declarations with context receivers

^KTIJ-21243 fixed

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

@@ -9,10 +9,7 @@ import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.metadata.ProtoBuf
 import org.jetbrains.kotlin.metadata.ProtoBuf.MemberKind
 import org.jetbrains.kotlin.metadata.ProtoBuf.Modality
-import org.jetbrains.kotlin.metadata.deserialization.Flags
-import org.jetbrains.kotlin.metadata.deserialization.hasReceiver
-import org.jetbrains.kotlin.metadata.deserialization.receiverType
-import org.jetbrains.kotlin.metadata.deserialization.returnType
+import org.jetbrains.kotlin.metadata.deserialization.*
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
 import org.jetbrains.kotlin.psi.stubs.impl.KotlinFunctionStubImpl
@@ -89,10 +86,12 @@ abstract class CallableClsStubBuilder(
 ) {
     protected val c = outerContext.child(typeParameters)
     protected val typeStubBuilder = TypeClsStubBuilder(c)
+    private val contextReceiversListStubBuilder = ContextReceiversListStubBuilder(c)
     protected val isTopLevel: Boolean get() = protoContainer is ProtoContainer.Package
     protected val callableStub: StubElement<out PsiElement> by lazy(LazyThreadSafetyMode.NONE) { doCreateCallableStub(parent) }
 
     fun build() {
+        contextReceiversListStubBuilder.createContextReceiverStubs(callableStub, contextReceiverTypes)
         createModifierListStub()
         val typeConstraintListData = typeStubBuilder.createTypeParameterListStub(callableStub, typeParameters)
         createReceiverTypeReferenceStub()
@@ -105,6 +104,7 @@ abstract class CallableClsStubBuilder(
     abstract val receiverAnnotations: List<ClassIdWithTarget>
 
     abstract val returnType: ProtoBuf.Type?
+    abstract val contextReceiverTypes: List<ProtoBuf.Type>
 
     private fun createReceiverTypeReferenceStub() {
         receiverType?.let {
@@ -144,6 +144,9 @@ private class FunctionClsStubBuilder(
     override val returnType: ProtoBuf.Type
         get() = functionProto.returnType(c.typeTable)
 
+    override val contextReceiverTypes: List<ProtoBuf.Type>
+        get() = functionProto.contextReceiverTypes(c.typeTable)
+
     override fun createValueParameterList() {
         typeStubBuilder.createValueParameterListStub(callableStub, functionProto, functionProto.valueParameterList, protoContainer)
     }
@@ -200,6 +203,9 @@ private class PropertyClsStubBuilder(
     override val returnType: ProtoBuf.Type
         get() = propertyProto.returnType(c.typeTable)
 
+    override val contextReceiverTypes: List<ProtoBuf.Type>
+        get() = propertyProto.contextReceiverTypes(c.typeTable)
+
     override fun createValueParameterList() {
     }
 
@@ -261,6 +267,9 @@ private class ConstructorClsStubBuilder(
     override val returnType: ProtoBuf.Type?
         get() = null
 
+    override val contextReceiverTypes: List<ProtoBuf.Type>
+        get() = emptyList()
+
     override fun createValueParameterList() {
         typeStubBuilder.createValueParameterListStub(callableStub, constructorProto, constructorProto.valueParameterList, protoContainer)
     }

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

@@ -12,10 +12,7 @@ import org.jetbrains.kotlin.descriptors.SourceElement
 import org.jetbrains.kotlin.lexer.KtModifierKeywordToken
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.metadata.ProtoBuf
-import org.jetbrains.kotlin.metadata.deserialization.Flags
-import org.jetbrains.kotlin.metadata.deserialization.NameResolver
-import org.jetbrains.kotlin.metadata.deserialization.TypeTable
-import org.jetbrains.kotlin.metadata.deserialization.supertypes
+import org.jetbrains.kotlin.metadata.deserialization.*
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.psi.KtClassBody
 import org.jetbrains.kotlin.psi.KtSuperTypeEntry
@@ -57,6 +54,7 @@ private class ClassClsStubBuilder(
     private val c = outerContext.child(
         classProto.typeParameterList, classId.shortClassName, nameResolver, thisAsProtoContainer.typeTable, thisAsProtoContainer
     )
+    private val contextReceiversListStubBuilder = ContextReceiversListStubBuilder(c)
     private val typeStubBuilder = TypeClsStubBuilder(c)
     private val supertypeIds = run {
         val supertypeIds = classProto.supertypes(c.typeTable).map { c.nameResolver.getClassId(it.className) }
@@ -85,6 +83,7 @@ private class ClassClsStubBuilder(
 
     private fun createClassOrObjectStubAndModifierListStub(): StubElement<out PsiElement> {
         val classOrObjectStub = doCreateClassOrObjectStub()
+        contextReceiversListStubBuilder.createContextReceiverStubs(classOrObjectStub, classProto.contextReceiverTypes(c.typeTable))
         val modifierList = createModifierListForClass(classOrObjectStub)
         if (Flags.HAS_ANNOTATIONS.get(classProto.flags)) {
             createAnnotationStubs(c.components.annotationLoader.loadClassAnnotations(thisAsProtoContainer), modifierList)

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

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.analysis.decompiler.stub
+
+import com.intellij.psi.stubs.StubElement
+import org.jetbrains.kotlin.metadata.ProtoBuf
+import org.jetbrains.kotlin.psi.KtContextReceiver
+import org.jetbrains.kotlin.psi.KtContextReceiverList
+import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
+import org.jetbrains.kotlin.psi.stubs.impl.KotlinPlaceHolderStubImpl
+
+internal class ContextReceiversListStubBuilder(c: ClsStubBuilderContext) {
+    private val typeStubBuilder = TypeClsStubBuilder(c)
+
+    fun createContextReceiverStubs(parent: StubElement<*>, contextReceiverTypes: List<ProtoBuf.Type>) {
+        if (contextReceiverTypes.isEmpty()) return
+        val contextReceiverListStub =
+            KotlinPlaceHolderStubImpl<KtContextReceiverList>(parent, KtStubElementTypes.CONTEXT_RECEIVER_LIST)
+        for (contextReceiverType in contextReceiverTypes) {
+            val contextReceiverStub =
+                KotlinPlaceHolderStubImpl<KtContextReceiver>(contextReceiverListStub, KtStubElementTypes.CONTEXT_RECEIVER)
+            typeStubBuilder.createTypeReferenceStub(contextReceiverStub, contextReceiverType)
+        }
+    }
+}