[LL FIR] FirElementBuilder: accept KtAnnotated instead of KtDeclaration as annotation owner

^KT-65780

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

@@ -14,14 +14,15 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.FileStruct
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.FirElementsRecorder
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.declarationCanBeLazilyResolved
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.findSourceNonLocalFirDeclaration
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.requireTypeIntersectionWith
 import org.jetbrains.kotlin.analysis.utils.printer.parentOfType
 import org.jetbrains.kotlin.analysis.utils.printer.parentsOfType
 import org.jetbrains.kotlin.fir.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.correspondingProperty
 import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirClass
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.declarations.FirReceiverParameter
@@ -111,11 +112,11 @@ internal class FirElementBuilder(
     private inline fun <T : KtElement> getFirForNonBodyElement(
         element: KtElement,
         anchorElementProvider: (KtElement) -> T?,
-        declarationProvider: (T) -> KtDeclaration?,
+        annotatedElementProvider: (T) -> KtAnnotated?,
-        resolveAndFindFirForAnchor: (FirDeclaration, T) -> FirElement?,
+        resolveAndFindFirForAnchor: (FirAnnotationContainer, T) -> FirElement?,
     ): FirElement? {
         val anchorElement = anchorElementProvider(element) ?: return null
-        val declaration = declarationProvider(anchorElement) ?: return null
+        val declaration = annotatedElementProvider(anchorElement) ?: return null
         val nonLocalDeclaration = declaration.getNonLocalContainingOrThisDeclaration()
         if (declaration != nonLocalDeclaration) return null
 
@@ -131,23 +132,27 @@ internal class FirElementBuilder(
         return findElementInside(firElement = anchorFir, element = element, stopAt = anchorElement)
     }
 
-    private fun KtAnnotationEntry.owner(): KtDeclaration? {
+    private fun KtAnnotationEntry.owner(): KtAnnotated? {
-        val parent = parent
+        val modifierList = when (val parent = parent) {
-        val modifierList = parent as? KtModifierList ?: (parent as? KtAnnotation)?.parent as? KtModifierList ?: return null
+            is KtModifierList -> parent
-        return modifierList.owner as? KtDeclaration
+            is KtAnnotation -> parent.parent as? KtModifierList
+            else -> null
+        }
+
+        return modifierList?.owner as? KtDeclaration
     }
 
     private fun getFirForElementInsideAnnotations(element: KtElement): FirElement? = getFirForNonBodyElement(
         element = element,
         anchorElementProvider = { it.parentOfType<KtAnnotationEntry>(withSelf = true) },
-        declarationProvider = { it.owner() },
+        annotatedElementProvider = { it.owner() },
         resolveAndFindFirForAnchor = { declaration, anchor -> declaration.resolveAndFindAnnotation(anchor, goDeep = true) },
     )
 
     private fun getFirForElementInsideTypes(element: KtElement): FirElement? = getFirForNonBodyElement(
         element = element,
         anchorElementProvider = { it.parentsOfType<KtTypeReference>(withSelf = true).lastOrNull() },
-        declarationProvider = {
+        annotatedElementProvider = {
             when (val parent = it.parent) {
                 is KtDeclaration -> parent
                 is KtSuperTypeListEntry, is KtConstructorCalleeExpression, is KtTypeConstraint -> parent.parentOfType<KtDeclaration>()
@@ -179,7 +184,9 @@ internal class FirElementBuilder(
         return firElement
     }
 
-    private fun FirDeclaration.resolveAndFindTypeRefAnchor(typeReference: KtTypeReference): FirElement? {
+    private fun FirAnnotationContainer.resolveAndFindTypeRefAnchor(typeReference: KtTypeReference): FirElement? {
+        requireTypeIntersectionWith<FirElementWithResolveState>()
+
         lazyResolveToPhase(FirResolvePhase.ANNOTATION_ARGUMENTS)
 
         if (this is FirCallableDeclaration) {
@@ -218,7 +225,12 @@ internal class FirElementBuilder(
         return null
     }
 
-    private fun FirDeclaration.resolveAndFindAnnotation(annotationEntry: KtAnnotationEntry, goDeep: Boolean = false): FirAnnotation? {
+    private fun FirAnnotationContainer.resolveAndFindAnnotation(
+        annotationEntry: KtAnnotationEntry,
+        goDeep: Boolean = false,
+    ): FirAnnotation? {
+        requireTypeIntersectionWith<FirElementWithResolveState>()
+
         lazyResolveToPhase(FirResolvePhase.ANNOTATION_ARGUMENTS)
         findAnnotation(annotationEntry)?.let { return it }
 

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/exceptionUtils.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2024 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.
  */
 
@@ -14,6 +14,9 @@ import org.jetbrains.kotlin.fir.utils.exceptions.withConeTypeEntry
 import org.jetbrains.kotlin.fir.utils.exceptions.withFirEntry
 import org.jetbrains.kotlin.utils.exceptions.ExceptionAttachmentBuilder
 import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
+import org.jetbrains.kotlin.utils.exceptions.requireWithAttachment
+import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.contract
 
 
 fun errorWithFirSpecificEntries(
@@ -39,3 +42,17 @@ fun errorWithFirSpecificEntries(
         additionalInfos()
     }
 }
+
+@OptIn(ExperimentalContracts::class)
+inline fun <reified R> Any.requireTypeIntersectionWith() {
+    contract { returns() implies (this@requireTypeIntersectionWith is R) }
+
+    requireWithAttachment(
+        this is R,
+        { "${this::class.simpleName} must be ${R::class.simpleName}" },
+    ) {
+        if (this@requireTypeIntersectionWith is FirElement) {
+            withFirEntry("container", this@requireTypeIntersectionWith)
+        }
+    }
+}