[LL API] Use a common type of 'KtDeclaration' and 'KtCodeFragment'

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

@@ -52,7 +52,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
         else -> true
     }
 
-    private fun PsiElement.getNonLocalContainingOrThisDeclarationCodeFragmentAware(predicate: (KtDeclaration) -> Boolean): KtElement? {
+    private fun PsiElement.getNonLocalContainingOrThisDeclarationCodeFragmentAware(predicate: (KtDeclaration) -> Boolean): KtAnnotated? {
         return getNonLocalContainingOrThisDeclaration(predicate) ?: containingFile as? KtCodeFragment
     }
 
@@ -342,7 +342,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
     }
 
     private fun requiredResolvePhase(
-        container: KtElement,
+        container: KtAnnotated,
         elementToReplace: PsiElement,
     ): FirResolvePhase {
         assert(container is KtDeclaration || container is KtCodeFragment)

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/inBlockModification.kt

@@ -22,9 +22,9 @@ import org.jetbrains.kotlin.fir.expressions.FirLazyExpression
 import org.jetbrains.kotlin.fir.expressions.builder.buildLazyBlock
 import org.jetbrains.kotlin.fir.expressions.builder.buildLazyExpression
 import org.jetbrains.kotlin.fir.expressions.impl.FirContractCallBlock
+import org.jetbrains.kotlin.psi.KtAnnotated
 import org.jetbrains.kotlin.psi.KtCodeFragment
 import org.jetbrains.kotlin.psi.KtDeclaration
-import org.jetbrains.kotlin.psi.KtElement
 
 @Deprecated(
     "Temporarily left for binary compatibility. Use invalidateAfterInBlockModification(KtElement) instead.",
@@ -34,7 +34,7 @@ import org.jetbrains.kotlin.psi.KtElement
 @LLFirInternals
 @Suppress("unused")
 fun invalidateAfterInBlockModification(declaration: KtDeclaration): Boolean {
-    return invalidateAfterInBlockModification(declaration as KtElement)
+    return invalidateAfterInBlockModification(declaration as KtAnnotated)
 }
 
 /**
@@ -42,7 +42,7 @@ fun invalidateAfterInBlockModification(declaration: KtDeclaration): Boolean {
  * @return **false** if it is not in-block modification
  */
 @LLFirInternals
-fun invalidateAfterInBlockModification(declaration: KtElement): Boolean {
+fun invalidateAfterInBlockModification(declaration: KtAnnotated): Boolean {
     ApplicationManager.getApplication().assertIsWriteThread()
 
     val project = declaration.project

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/FirLazyBodiesCalculator.kt

@@ -87,7 +87,7 @@ private inline fun <reified T : FirDeclaration> revive(
         session = session,
         scopeProvider = session.kotlinScopeProvider,
         designation = designation,
-        rootNonLocalDeclaration = psiFactory(designation) as KtElement,
+        rootNonLocalDeclaration = psiFactory(designation) as KtAnnotated,
     ) as T
 }
 

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/RawFirNonLocalDeclarationBuilder.kt

@@ -68,7 +68,7 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
             session: FirSession,
             scopeProvider: FirScopeProvider,
             designation: FirDesignation,
-            rootNonLocalDeclaration: KtElement,
+            rootNonLocalDeclaration: KtAnnotated,
         ): FirDeclaration {
             val functionsToRebind = when (val originalDeclaration = designation.target) {
                 is FirFunction -> setOf(originalDeclaration)