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
2023-06-28 16:42:41 +02:00
parent d7d2c1f260
commit 88450c69b3
27 changed files with 423 additions and 4 deletions
@@ -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");
+            }
+        }
    }
 }
@@ -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");
+            }
+        }
    }
 }
@@ -0,0 +1,8 @@
+// FILE: main.kt
+
+abstract class Outer {
+    annotation class Ann
+}
+
+<expr>@Outer.Ann</expr>
+class Foo : Outer()
@@ -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
@@ -0,0 +1,7 @@
+// FILE: main.kt
+
+abstract class Outer {
+    annotation class Ann
+}
+
+class Foo <expr>@Outer.Ann</expr> constructor() : Outer()
@@ -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
@@ -0,0 +1,7 @@
+// FILE: main.kt
+
+abstract class Outer {
+    annotation class Ann
+}
+
+class Foo(<expr>@Outer.Ann</expr> val p: Any) : Outer()
@@ -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
@@ -0,0 +1,8 @@
+// FILE: main.kt
+
+abstract class Outer {
+    abstract class Nested : Outer()
+}
+
+context(<expr>Outer.Nested</expr>)
+class Foo : Outer.Nested()
@@ -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
@@ -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()
@@ -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()
@@ -0,0 +1,7 @@
+// FILE: main.kt
+
+abstract class Outer {
+    abstract class Nested : Outer()
+}
+
+<expr>class Foo : Outer.Nested { constructor(): super() }</expr>
@@ -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
@@ -0,0 +1,7 @@
+// FILE: main.kt
+
+abstract class Outer {
+    abstract class Nested : Outer()
+}
+
+<expr>class Foo : Outer.Nested()</expr>
@@ -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
@@ -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>)
@@ -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()
@@ -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>
@@ -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()
@@ -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() }
@@ -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
@@ -0,0 +1,7 @@
+// FILE: main.kt
+
+abstract class Outer {
+    abstract class Nested : Outer()
+}
+
+class Foo<T : <expr>Outer.Nested</expr>> : Outer.Nested()
@@ -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
@@ -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 }
+}
@@ -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 {
@@ -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)
 }