KTIJ-26024 [Analysis API] Provide correct context for class header from FirTowerDataContextAllElementsCollector

To correctly provide the context for class header resolution, we save
the additional copy of the context right before `withRegularClass`
wrapper call in `LLFirBodyLazyResolver`. Otherwise we would have to
clear the existing context by hand, which is too cumbersome
(if at all possible)

^KTIJ-26024 Fixed
^KTIJ-24832 Fixed

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/references/FirIdeNormalAnalysisSourceModuleReferenceShortenerTestGenerated.java

@@ -450,5 +450,81 @@ public class FirIdeNormalAnalysisSourceModuleReferenceShortenerTestGenerated ext
         public void testNestedClassFromSupertypes6() throws Exception {
             runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/nestedClassFromSupertypes6.kt");
         }
+
+        @Nested
+        @TestMetadata("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions")
+        @TestDataPath("$PROJECT_ROOT")
+        public class ClassHeaderPositions {
+            @Test
+            public void testAllFilesPresentInClassHeaderPositions() throws Exception {
+                KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions"), Pattern.compile("^(.+)\\.kt$"), null, true);
+            }
+
+            @Test
+            @TestMetadata("annotationOnClass.kt")
+            public void testAnnotationOnClass() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/annotationOnClass.kt");
+            }
+
+            @Test
+            @TestMetadata("annotationOnConstructor.kt")
+            public void testAnnotationOnConstructor() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/annotationOnConstructor.kt");
+            }
+
+            @Test
+            @TestMetadata("annotationOnParameter.kt")
+            public void testAnnotationOnParameter() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/annotationOnParameter.kt");
+            }
+
+            @Test
+            @TestMetadata("contextReceiver.kt")
+            public void testContextReceiver() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/contextReceiver.kt");
+            }
+
+            @Test
+            @TestMetadata("primaryConstructorParameter.kt")
+            public void testPrimaryConstructorParameter() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/primaryConstructorParameter.kt");
+            }
+
+            @Test
+            @TestMetadata("superType.kt")
+            public void testSuperType() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superType.kt");
+            }
+
+            @Test
+            @TestMetadata("superTypeConstructor.kt")
+            public void testSuperTypeConstructor() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superTypeConstructor.kt");
+            }
+
+            @Test
+            @TestMetadata("superTypeConstructorArgument.kt")
+            public void testSuperTypeConstructorArgument() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superTypeConstructorArgument.kt");
+            }
+
+            @Test
+            @TestMetadata("superTypeDelegation.kt")
+            public void testSuperTypeDelegation() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superTypeDelegation.kt");
+            }
+
+            @Test
+            @TestMetadata("superTypeTypeArguments.kt")
+            public void testSuperTypeTypeArguments() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superTypeTypeArguments.kt");
+            }
+
+            @Test
+            @TestMetadata("typeBound.kt")
+            public void testTypeBound() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/typeBound.kt");
+            }
+        }
     }
 }

analysis/analysis-api-standalone/tests-gen/org/jetbrains/kotlin/analysis/api/standalone/fir/test/cases/generated/cases/references/FirStandaloneNormalAnalysisSourceModuleReferenceShortenerTestGenerated.java

@@ -450,5 +450,81 @@ public class FirStandaloneNormalAnalysisSourceModuleReferenceShortenerTestGenera
         public void testNestedClassFromSupertypes6() throws Exception {
             runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/nestedClassFromSupertypes6.kt");
         }
+
+        @Nested
+        @TestMetadata("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions")
+        @TestDataPath("$PROJECT_ROOT")
+        public class ClassHeaderPositions {
+            @Test
+            public void testAllFilesPresentInClassHeaderPositions() throws Exception {
+                KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions"), Pattern.compile("^(.+)\\.kt$"), null, true);
+            }
+
+            @Test
+            @TestMetadata("annotationOnClass.kt")
+            public void testAnnotationOnClass() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/annotationOnClass.kt");
+            }
+
+            @Test
+            @TestMetadata("annotationOnConstructor.kt")
+            public void testAnnotationOnConstructor() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/annotationOnConstructor.kt");
+            }
+
+            @Test
+            @TestMetadata("annotationOnParameter.kt")
+            public void testAnnotationOnParameter() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/annotationOnParameter.kt");
+            }
+
+            @Test
+            @TestMetadata("contextReceiver.kt")
+            public void testContextReceiver() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/contextReceiver.kt");
+            }
+
+            @Test
+            @TestMetadata("primaryConstructorParameter.kt")
+            public void testPrimaryConstructorParameter() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/primaryConstructorParameter.kt");
+            }
+
+            @Test
+            @TestMetadata("superType.kt")
+            public void testSuperType() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superType.kt");
+            }
+
+            @Test
+            @TestMetadata("superTypeConstructor.kt")
+            public void testSuperTypeConstructor() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superTypeConstructor.kt");
+            }
+
+            @Test
+            @TestMetadata("superTypeConstructorArgument.kt")
+            public void testSuperTypeConstructorArgument() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superTypeConstructorArgument.kt");
+            }
+
+            @Test
+            @TestMetadata("superTypeDelegation.kt")
+            public void testSuperTypeDelegation() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superTypeDelegation.kt");
+            }
+
+            @Test
+            @TestMetadata("superTypeTypeArguments.kt")
+            public void testSuperTypeTypeArguments() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superTypeTypeArguments.kt");
+            }
+
+            @Test
+            @TestMetadata("typeBound.kt")
+            public void testTypeBound() throws Exception {
+                runTest("analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/typeBound.kt");
+            }
+        }
     }
 }

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/annotationOnClass.kt

@@ -0,0 +1,8 @@
+// FILE: main.kt
+
+abstract class Outer {
+    annotation class Ann
+}
+
+<expr>@Outer.Ann</expr>
+class Foo : Outer()

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/annotationOnClass.txt

@@ -0,0 +1,7 @@
+Before shortening: @Outer.Ann
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+with SHORTEN_AND_IMPORT:
+[type] Outer.Ann
+with SHORTEN_AND_STAR_IMPORT:
+[type] Outer.Ann

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/annotationOnConstructor.kt

@@ -0,0 +1,7 @@
+// FILE: main.kt
+
+abstract class Outer {
+    annotation class Ann
+}
+
+class Foo <expr>@Outer.Ann</expr> constructor() : Outer()

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/annotationOnConstructor.txt

@@ -0,0 +1,8 @@
+Before shortening: @Outer.Ann
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+[type] Outer.Ann
+with SHORTEN_AND_IMPORT:
+[type] Outer.Ann
+with SHORTEN_AND_STAR_IMPORT:
+[type] Outer.Ann

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/annotationOnParameter.kt

@@ -0,0 +1,7 @@
+// FILE: main.kt
+
+abstract class Outer {
+    annotation class Ann
+}
+
+class Foo(<expr>@Outer.Ann</expr> val p: Any) : Outer()

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/annotationOnParameter.txt

@@ -0,0 +1,8 @@
+Before shortening: @Outer.Ann
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+[type] Outer.Ann
+with SHORTEN_AND_IMPORT:
+[type] Outer.Ann
+with SHORTEN_AND_STAR_IMPORT:
+[type] Outer.Ann

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/contextReceiver.kt

@@ -0,0 +1,8 @@
+// FILE: main.kt
+
+abstract class Outer {
+    abstract class Nested : Outer()
+}
+
+context(<expr>Outer.Nested</expr>)
+class Foo : Outer.Nested()

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/contextReceiver.txt

@@ -0,0 +1,7 @@
+Before shortening: Outer.Nested
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+with SHORTEN_AND_IMPORT:
+[type] Outer.Nested
+with SHORTEN_AND_STAR_IMPORT:
+[type] Outer.Nested

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/primaryConstructorParameter.kt

@@ -0,0 +1,7 @@
+// FILE: main.kt
+
+abstract class Outer {
+    abstract class Nested : Outer()
+}
+
+class Foo(<expr>val param: Outer.Nested = Outer.Nested()</expr>) : Outer.Nested()

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/primaryConstructorParameter.txt

@@ -0,0 +1,11 @@
+Before shortening: val param: Outer.Nested = Outer.Nested()
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+[type] Outer.Nested
+[qualifier] Outer.Nested()
+with SHORTEN_AND_IMPORT:
+[type] Outer.Nested
+[qualifier] Outer.Nested()
+with SHORTEN_AND_STAR_IMPORT:
+[type] Outer.Nested
+[qualifier] Outer.Nested()

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superType.kt

@@ -0,0 +1,7 @@
+// FILE: main.kt
+
+abstract class Outer {
+    abstract class Nested : Outer()
+}
+
+<expr>class Foo : Outer.Nested { constructor(): super() }</expr>

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superType.txt

@@ -0,0 +1,7 @@
+Before shortening: class Foo : Outer.Nested { constructor(): super() }
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+with SHORTEN_AND_IMPORT:
+[type] Outer.Nested
+with SHORTEN_AND_STAR_IMPORT:
+[type] Outer.Nested

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superTypeConstructor.kt

@@ -0,0 +1,7 @@
+// FILE: main.kt
+
+abstract class Outer {
+    abstract class Nested : Outer()
+}
+
+<expr>class Foo : Outer.Nested()</expr>

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superTypeConstructor.txt

@@ -0,0 +1,7 @@
+Before shortening: class Foo : Outer.Nested()
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+with SHORTEN_AND_IMPORT:
+[type] Outer.Nested
+with SHORTEN_AND_STAR_IMPORT:
+[type] Outer.Nested

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superTypeConstructorArgument.kt

@@ -0,0 +1,9 @@
+// FILE: main.kt
+
+abstract class Outer {
+    abstract class Nested(a: Any) : Outer()
+
+    class OtherNested
+}
+
+class Foo : Outer.Nested(<expr>Outer.OtherNested()</expr>)

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superTypeConstructorArgument.txt

@@ -0,0 +1,8 @@
+Before shortening: Outer.OtherNested()
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+[qualifier] Outer.OtherNested()
+with SHORTEN_AND_IMPORT:
+[qualifier] Outer.OtherNested()
+with SHORTEN_AND_STAR_IMPORT:
+[qualifier] Outer.OtherNested()

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superTypeDelegation.kt

@@ -0,0 +1,11 @@
+// FILE: main.kt
+
+abstract class Outer {
+    abstract class Nested : Outer(), Marker
+
+    interface class Marker
+
+    class MarkerImpl() : Marker
+}
+
+class Foo : Outer.Nested(), <expr>Outer.Marker by Outer.MarkerImpl()</expr>

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superTypeDelegation.txt

@@ -0,0 +1,10 @@
+Before shortening: Outer.Marker by Outer.MarkerImpl()
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+[qualifier] Outer.MarkerImpl()
+with SHORTEN_AND_IMPORT:
+[type] Outer.Marker
+[qualifier] Outer.MarkerImpl()
+with SHORTEN_AND_STAR_IMPORT:
+[type] Outer.Marker
+[qualifier] Outer.MarkerImpl()

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superTypeTypeArguments.kt

@@ -0,0 +1,9 @@
+// FILE: main.kt
+
+abstract class Outer {
+    abstract class Nested<T> : Outer()
+
+    class OtherNested
+}
+
+class Foo : <expr>Outer.Nested<Outer.OtherNested></expr> { constructor(): super() }

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/superTypeTypeArguments.txt

@@ -0,0 +1,9 @@
+Before shortening: Outer.Nested<Outer.OtherNested>
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+with SHORTEN_AND_IMPORT:
+[type] Outer.Nested<Outer.OtherNested>
+[type] Outer.OtherNested
+with SHORTEN_AND_STAR_IMPORT:
+[type] Outer.Nested<Outer.OtherNested>
+[type] Outer.OtherNested

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/typeBound.kt

@@ -0,0 +1,7 @@
+// FILE: main.kt
+
+abstract class Outer {
+    abstract class Nested : Outer()
+}
+
+class Foo<T : <expr>Outer.Nested</expr>> : Outer.Nested()

analysis/analysis-api/testData/components/referenceShortener/referenceShortener/nestedClasses/classHeaderPositions/typeBound.txt

@@ -0,0 +1,7 @@
+Before shortening: Outer.Nested
+with DO_NOT_SHORTEN:
+with SHORTEN_IF_ALREADY_IMPORTED:
+with SHORTEN_AND_IMPORT:
+[type] Outer.Nested
+with SHORTEN_AND_STAR_IMPORT:
+[type] Outer.Nested

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirTowerDataContextCollector.kt

@@ -9,11 +9,14 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirTowerDataContext
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirTowerDataContextCollector
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
+import org.jetbrains.kotlin.utils.addIfNotNull
 
 interface FirTowerContextProvider {
     fun getClosestAvailableParentContext(ktElement: KtElement): FirTowerDataContext?
@@ -30,6 +33,14 @@ internal class FileTowerProvider(
 internal class FirTowerDataContextAllElementsCollector : FirTowerDataContextCollector, FirTowerContextProvider {
     private val elementsToContext: MutableMap<KtElement, FirTowerDataContext> = hashMapOf()
 
+    /**
+     * Class headers have context which is different from the context inside of the class, so we have
+     * to collect the scopes for it separately.
+     *
+     * See [containsInHeader] for more info.
+     */
+    private val classesToClassHeaderContext: MutableMap<KtClassOrObject, FirTowerDataContext> = hashMapOf()
+
     override fun addFileContext(file: FirFile, context: FirTowerDataContext) {
         val ktFile = file.psi as? KtFile ?: return
         elementsToContext[ktFile] = context
@@ -43,14 +54,26 @@ internal class FirTowerDataContextAllElementsCollector : FirTowerDataContextColl
 
     override fun addDeclarationContext(declaration: FirDeclaration, context: FirTowerDataContext) {
         val psi = declaration.psi as? KtElement ?: return
+
+        // FIR creates a fake field declaration for the delegated super calls,
+        // and it makes handling class headers more difficult, so
+        // we do not collect contexts for such declarations
+        if (psi is KtDelegatedSuperTypeEntry) return
+
         elementsToContext[psi] = context
     }
 
+    override fun addClassHeaderContext(declaration: FirRegularClass, context: FirTowerDataContext) {
+        val psi = declaration.psi as? KtElement ?: return
+        if (psi !is KtClassOrObject) return
+        classesToClassHeaderContext[psi] = context
+    }
+
     override fun getClosestAvailableParentContext(ktElement: KtElement): FirTowerDataContext? {
         var current: PsiElement? = ktElement
         while (current != null) {
             if (current is KtElement) {
-                getContextForElement(current)?.let { return it }
+                getContextForElementOnPosition(current, ktElement)?.let { return it }
             }
             if (current is KtFile) {
                 break
@@ -60,12 +83,26 @@ internal class FirTowerDataContextAllElementsCollector : FirTowerDataContextColl
         return null
     }
 
-    private fun getContextForElement(element: KtElement): FirTowerDataContext? {
+    /**
-        elementsToContext[element]?.let { return it }
+     * N.B. It is expected that [element] contains [position].
+     */
+    private fun getContextForElementOnPosition(element: KtElement, position: KtElement): FirTowerDataContext? {
+        val contextMapping =
+            if (element is KtClassOrObject && element.containsInHeader(position)) {
+                classesToClassHeaderContext
+            } else {
+                elementsToContext
+            }
+
+        return getContextForElement(element, contextMapping)
+    }
+
+    private fun getContextForElement(element: KtElement, contextMapping: Map<out KtElement, FirTowerDataContext>): FirTowerDataContext? {
+        contextMapping[element]?.let { return it }
 
         if (element is KtDeclaration) {
             val originalDeclaration = element.originalDeclaration
-            originalDeclaration?.let { elementsToContext[it] }?.let { return it }
+            originalDeclaration?.let { contextMapping[it] }?.let { return it }
         }
 
         return null
@@ -77,3 +114,28 @@ private tailrec fun PsiElement.closestBlockLevelOrInitializerExpression(): KtExp
         this is KtExpression && (parent is KtBlockExpression || parent is KtDeclarationWithInitializer) -> this
         else -> parent?.closestBlockLevelOrInitializerExpression()
     }
+
+/**
+ * Returns true if [element] is considered to be a part of [this] class header.
+ *
+ * Here by "header" we understand the parts of the class declaration which resolution is not affected
+ * by the class own supertypes.
+ */
+private fun KtClassOrObject.containsInHeader(element: PsiElement): Boolean {
+    if (this == element) return false
+
+    val classHeaderParts = buildSet<KtElement> {
+        addIfNotNull(getContextReceiverList())
+        addIfNotNull(modifierList) // for annotations declared on class
+        addIfNotNull(typeParameterList)
+
+        for (superTypeEntry in superTypeListEntries) {
+            addIfNotNull(superTypeEntry.typeReference)
+        }
+    }
+
+    if (classHeaderParts.isEmpty()) return false
+
+    val elementParentsUntilClass = element.parentsWithSelf.takeWhile { it != this }
+    return elementParentsUntilClass.any { it in classHeaderParts }
+}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.analysis.utils.errors.checkWithAttachmentBuilder
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.FirFileAnnotationsContainer
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.PrivateForInline
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.getExplicitBackingField
 import org.jetbrains.kotlin.fir.expressions.*
@@ -128,6 +129,9 @@ private class LLFirBodyTargetResolver(
     }
 
     private fun resolveMembersForControlFlowGraph(target: FirRegularClass) {
+        withTypeArguments(target) {
+            transformer.firTowerDataContextCollector?.addClassHeaderContext(target, transformer.context.towerDataContext)
+        }
         withRegularClass(target) {
             transformer.firTowerDataContextCollector?.addDeclarationContext(target, transformer.context.towerDataContext)
 
@@ -173,6 +177,13 @@ private class LLFirBodyTargetResolver(
             else -> throwUnexpectedFirElementError(target)
         }
     }
+
+    private inline fun withTypeArguments(regularClass: FirRegularClass, action: () -> Unit) {
+        @OptIn(PrivateForInline::class)
+        transformer.declarationsTransformer.context.withTypeParametersOf(regularClass) {
+            action()
+        }
+    }
 }
 
 internal object BodyStateKeepers {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirTowerDataContextCollector.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirTowerDataContext
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 
@@ -17,4 +18,21 @@ interface FirTowerDataContextCollector {
     fun addFileContext(file: FirFile, context: FirTowerDataContext)
     fun addStatementContext(statement: FirStatement, context: FirTowerDataContext)
     fun addDeclarationContext(declaration: FirDeclaration, context: FirTowerDataContext)
+
+    /**
+     * Here by "header" we understand the parts of the class declaration which resolution is not affected
+     * by the class own supertypes.
+     *
+     * It includes any type reference mentioned before the body of the class, except for:
+     * - primary constructor declaration with annotations
+     * - super type constructor **value arguments** with annotations
+     *
+     * Also, it includes class type parameters, because they are available anywhere inside the class declaration.
+     *
+     * N.B. We collect header contexts only for [FirRegularClass]es, ignoring
+     * [org.jetbrains.kotlin.fir.declarations.FirAnonymousObject]s. That's because anonymous objects are in
+     * fact expressions, and it should be enough for [FirTowerDataContextCollector] to have a context for
+     * the containing declaration to correctly process everything inside anonymous objects.
+     */
+    fun addClassHeaderContext(declaration: FirRegularClass, context: FirTowerDataContext)
 }