[LL FIR] introduce NonLocalAnnotationVisitor

- Unified and optimized visitor instead of transformer with many
  redundant operations.
- `LLFirAnnotationArgumentsTargetResolver` now has the correct state
  keeper for nested declarations – e.g., property accessors or value
  parameters.
- `FirAllLazyAnnotationCalculatorVisitor` in tests now almost the same
  as in the production as now it uses the same visitor without recursive
  visiting

^KT-63606 Fixed
^KT-64000 Fixed
^KT-64579 Fixed

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

@@ -9,6 +9,7 @@ import com.intellij.psi.PsiElement
 import kotlinx.collections.immutable.PersistentList
 import kotlinx.collections.immutable.persistentListOf
 import kotlinx.collections.immutable.toPersistentList
+import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignation
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.withFirDesignationEntry
@@ -30,7 +31,6 @@ import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
 import org.jetbrains.kotlin.fir.scopes.kotlinScopeProvider
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
-import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildTypeProjectionWithVariance
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
@@ -52,8 +52,9 @@ internal object FirLazyBodiesCalculator {
         )
     }
 
+    @TestOnly
     fun calculateAllLazyExpressionsInFile(firFile: FirFile) {
-        firFile.transformSingle(FirAllLazyAnnotationCalculatorTransformer, firFile.moduleData.session)
+        firFile.accept(RecursiveLazyAnnotationCalculatorVisitor, firFile.moduleData.session)
         firFile.transformSingle(FirAllLazyBodiesCalculatorTransformer, persistentListOf())
     }
 
@@ -62,7 +63,7 @@ internal object FirLazyBodiesCalculator {
     }
 
     fun calculateAnnotations(firElement: FirElement, session: FirSession) {
-        firElement.transformSingle(FirTargetLazyAnnotationCalculatorTransformer, session)
+        firElement.accept(LazyAnnotationCalculatorVisitor, session)
     }
 
     fun calculateLazyArgumentsForAnnotation(annotationCall: FirAnnotationCall, session: FirSession): FirArgumentList {
@@ -624,108 +625,31 @@ private fun calculateLazyBodyForCodeFragment(designation: FirDesignation) {
     codeFragment.replaceBlock(newCodeFragment.block)
 }
 
-private object FirAllLazyAnnotationCalculatorTransformer : FirLazyAnnotationTransformer() {
-    override fun <E : FirElement> transformElement(element: E, data: FirSession): E {
-        element.transformChildren(this, data)
-        return element
+/**
+ * This object is supposed to be used only for tests.
+ *
+ * @see LazyAnnotationCalculatorVisitor
+ */
+private object RecursiveLazyAnnotationCalculatorVisitor : RecursiveNonLocalAnnotationVisitor<FirSession>() {
+    override fun processAnnotation(annotation: FirAnnotation, data: FirSession) {
+        calculateAnnotationCallIfNeeded(annotation, data)
     }
 }
 
-private object FirTargetLazyAnnotationCalculatorTransformer : FirLazyAnnotationTransformer() {
-    override fun <E : FirElement> transformElement(element: E, data: FirSession): E {
-        element.transformChildren(this, data)
-        return element
-    }
-
-    override fun transformRegularClass(regularClass: FirRegularClass, data: FirSession): FirStatement {
-        regularClass.transformAnnotations(this, data)
-        regularClass.transformTypeParameters(this, data)
-        regularClass.transformSuperTypeRefs(this, data)
-        regularClass.contextReceivers.forEach {
-            it.transformSingle(this, data)
-        }
-
-        return regularClass
-    }
-
-    override fun transformScript(script: FirScript, data: FirSession): FirScript {
-        script.transformAnnotations(this, data)
-        return script
-    }
-
-    override fun transformBlock(block: FirBlock, data: FirSession): FirStatement {
-        // We shouldn't process blocks because there are no lazy annotations
-        return block
-    }
-
-    override fun transformFile(file: FirFile, data: FirSession): FirFile {
-        file.transformAnnotationsContainer(this, data)
-        return file
+/**
+ * Calculates all non-local lazy annotations on a provided declaration.
+ */
+private object LazyAnnotationCalculatorVisitor : NonLocalAnnotationVisitor<FirSession>() {
+    override fun processAnnotation(annotation: FirAnnotation, data: FirSession) {
+        calculateAnnotationCallIfNeeded(annotation, data)
     }
 }
 
-private abstract class FirLazyAnnotationTransformer : FirTransformer<FirSession>() {
-    override fun <E : FirElement> transformElement(element: E, data: FirSession): E {
-        element.transformChildren(this, data)
-        return element
-    }
+private fun calculateAnnotationCallIfNeeded(annotation: FirAnnotation, session: FirSession) {
+    if (annotation !is FirAnnotationCall || !FirLazyBodiesCalculator.needCalculatingAnnotationCall(annotation)) return
 
-    override fun transformErrorTypeRef(errorTypeRef: FirErrorTypeRef, data: FirSession): FirTypeRef {
-        visitTypeAnnotations(errorTypeRef, data)
-        return super.transformErrorTypeRef(errorTypeRef, data)
-    }
-
-    override fun transformResolvedTypeRef(resolvedTypeRef: FirResolvedTypeRef, data: FirSession): FirTypeRef {
-        visitTypeAnnotations(resolvedTypeRef, data)
-        return super.transformResolvedTypeRef(resolvedTypeRef, data)
-    }
-
-    private fun visitTypeAnnotations(resolvedTypeRef: FirResolvedTypeRef, data: FirSession) {
-        resolvedTypeRef.coneType.forEachType { coneType ->
-            for (typeArgumentAnnotation in coneType.customAnnotations) {
-                typeArgumentAnnotation.accept(this, data)
-            }
-        }
-    }
-
-    override fun transformAnnotationCall(annotationCall: FirAnnotationCall, data: FirSession): FirStatement {
-        if (FirLazyBodiesCalculator.needCalculatingAnnotationCall(annotationCall)) {
-            val newArgumentList = FirLazyBodiesCalculator.calculateLazyArgumentsForAnnotation(annotationCall, data)
-            annotationCall.replaceArgumentList(newArgumentList)
-        }
-
-        super.transformAnnotationCall(annotationCall, data)
-        return annotationCall
-    }
-
-    override fun transformErrorAnnotationCall(errorAnnotationCall: FirErrorAnnotationCall, data: FirSession): FirStatement {
-        transformAnnotationCall(errorAnnotationCall, data)
-        return errorAnnotationCall
-    }
-
-    override fun transformExpression(expression: FirExpression, data: FirSession): FirStatement {
-        if (expression is FirLazyExpression) {
-            return expression
-        }
-
-        return super.transformExpression(expression, data)
-    }
-
-    override fun transformBlock(block: FirBlock, data: FirSession): FirStatement {
-        if (block is FirLazyBlock) {
-            return block
-        }
-
-        return super.transformBlock(block, data)
-    }
-
-    override fun transformDelegatedConstructorCall(delegatedConstructorCall: FirDelegatedConstructorCall, data: FirSession): FirStatement {
-        if (delegatedConstructorCall is FirLazyDelegatedConstructorCall) {
-            return delegatedConstructorCall
-        }
-
-        return super.transformDelegatedConstructorCall(delegatedConstructorCall, data)
-    }
+    val newArgumentList = FirLazyBodiesCalculator.calculateLazyArgumentsForAnnotation(annotation, session)
+    annotation.replaceArgumentList(newArgumentList)
 }
 
 private object FirAllLazyBodiesCalculatorTransformer : FirLazyBodiesCalculatorTransformer() {

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

@@ -0,0 +1,215 @@
+/*
+ * 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.
+ */
+
+package org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve
+
+import org.jetbrains.kotlin.fir.FirAnnotationContainer
+import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirElementWithResolveState
+import org.jetbrains.kotlin.fir.FirFileAnnotationsContainer
+import org.jetbrains.kotlin.fir.declarations.FirAnonymousInitializer
+import org.jetbrains.kotlin.fir.declarations.FirBackingField
+import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirConstructor
+import org.jetbrains.kotlin.fir.declarations.FirContextReceiver
+import org.jetbrains.kotlin.fir.declarations.FirDanglingModifierList
+import org.jetbrains.kotlin.fir.declarations.FirEnumEntry
+import org.jetbrains.kotlin.fir.declarations.FirErrorPrimaryConstructor
+import org.jetbrains.kotlin.fir.declarations.FirErrorProperty
+import org.jetbrains.kotlin.fir.declarations.FirFunction
+import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.declarations.FirPropertyAccessor
+import org.jetbrains.kotlin.fir.declarations.FirReceiverParameter
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.FirScript
+import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
+import org.jetbrains.kotlin.fir.declarations.FirTypeAlias
+import org.jetbrains.kotlin.fir.declarations.FirTypeParameter
+import org.jetbrains.kotlin.fir.declarations.FirTypeParameterRefsOwner
+import org.jetbrains.kotlin.fir.declarations.FirValueParameter
+import org.jetbrains.kotlin.fir.declarations.FirVariable
+import org.jetbrains.kotlin.fir.expressions.FirAnnotation
+import org.jetbrains.kotlin.fir.expressions.FirAnnotationArgumentMapping
+import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
+import org.jetbrains.kotlin.fir.expressions.FirArgumentList
+import org.jetbrains.kotlin.fir.expressions.FirErrorAnnotationCall
+import org.jetbrains.kotlin.fir.types.FirErrorTypeRef
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.fir.types.customAnnotations
+import org.jetbrains.kotlin.fir.types.forEachType
+import org.jetbrains.kotlin.fir.visitors.FirVisitor
+
+/**
+ * Visitor for annotations outside bodies.
+ * Such annotations are processed in [ANNOTATION_ARGUMENTS][org.jetbrains.kotlin.fir.declarations.FirResolvePhase.ANNOTATION_ARGUMENTS]
+ * phase.
+ *
+ * This visitor is not recursive and processes only the target declaration (without unrelated nested declarations).
+ * See [RecursiveNonLocalAnnotationVisitor] for the recursive visitor.
+ *
+ * @see processAnnotation
+ * @see RecursiveNonLocalAnnotationVisitor
+ */
+internal abstract class NonLocalAnnotationVisitor<T> : FirVisitor<Unit, T>() {
+    abstract fun processAnnotation(annotation: FirAnnotation, data: T)
+
+    /**
+     * Skip all [FirElementWithResolveState] without explicit override
+     */
+    override fun visitElement(element: FirElement, data: T) {
+        if (element is FirElementWithResolveState) return
+
+        element.acceptChildren(this, data)
+    }
+
+    /**
+     * Skip argument list as the compiler do not support annotations inside annotation arguments
+     */
+    override fun visitArgumentList(argumentList: FirArgumentList, data: T) {}
+    override fun visitAnnotationArgumentMapping(annotationArgumentMapping: FirAnnotationArgumentMapping, data: T) {}
+
+    override fun visitResolvedTypeRef(resolvedTypeRef: FirResolvedTypeRef, data: T) {
+        visitTypeAnnotations(resolvedTypeRef, data)
+        resolvedTypeRef.acceptChildren(this, data)
+    }
+
+    override fun visitErrorTypeRef(errorTypeRef: FirErrorTypeRef, data: T) {
+        visitResolvedTypeRef(errorTypeRef, data)
+    }
+
+    private fun visitTypeAnnotations(resolvedTypeRef: FirResolvedTypeRef, data: T) {
+        resolvedTypeRef.coneType.forEachType { coneType ->
+            for (typeArgumentAnnotation in coneType.customAnnotations) {
+                typeArgumentAnnotation.accept(this, data)
+            }
+        }
+    }
+
+    override fun visitAnnotation(annotation: FirAnnotation, data: T) {
+        processAnnotation(annotation, data)
+        annotation.acceptChildren(this, data)
+    }
+
+    override fun visitAnnotationCall(annotationCall: FirAnnotationCall, data: T) {
+        visitAnnotation(annotationCall, data)
+    }
+
+    override fun visitErrorAnnotationCall(errorAnnotationCall: FirErrorAnnotationCall, data: T) {
+        visitAnnotation(errorAnnotationCall, data)
+    }
+
+    override fun visitFileAnnotationsContainer(fileAnnotationsContainer: FirFileAnnotationsContainer, data: T) {
+        visitAnnotationContainer(fileAnnotationsContainer, data)
+    }
+
+    override fun visitDanglingModifierList(danglingModifierList: FirDanglingModifierList, data: T) {
+        visitAnnotationContainer(danglingModifierList, data)
+    }
+
+    override fun visitScript(script: FirScript, data: T) {
+        visitAnnotationContainer(script, data)
+    }
+
+    override fun visitAnonymousInitializer(anonymousInitializer: FirAnonymousInitializer, data: T) {
+        visitAnnotationContainer(anonymousInitializer, data)
+    }
+
+    override fun visitMemberDeclaration(memberDeclaration: FirMemberDeclaration, data: T) {
+        visitTypeParameterRefsOwner(memberDeclaration, data)
+        visitAnnotationContainer(memberDeclaration, data)
+    }
+
+    override fun visitTypeAlias(typeAlias: FirTypeAlias, data: T) {
+        visitMemberDeclaration(typeAlias, data)
+        typeAlias.expandedTypeRef.accept(this, data)
+    }
+
+    override fun visitRegularClass(regularClass: FirRegularClass, data: T) {
+        visitMemberDeclaration(regularClass, data)
+        visitContextReceivers(regularClass.contextReceivers, data)
+        regularClass.superTypeRefs.forEach { it.accept(this, data) }
+    }
+
+    override fun visitCallableDeclaration(callableDeclaration: FirCallableDeclaration, data: T) {
+        visitMemberDeclaration(callableDeclaration, data)
+        visitContextReceivers(callableDeclaration.contextReceivers, data)
+        callableDeclaration.receiverParameter?.accept(this, data)
+        callableDeclaration.returnTypeRef.accept(this, data)
+    }
+
+    fun visitContextReceivers(contextReceivers: List<FirContextReceiver>, data: T) {
+        contextReceivers.forEach { it.accept(this, data) }
+    }
+
+    override fun visitAnnotationContainer(annotationContainer: FirAnnotationContainer, data: T) {
+        annotationContainer.annotations.forEach { it.accept(this, data) }
+    }
+
+    override fun visitTypeParameterRefsOwner(typeParameterRefsOwner: FirTypeParameterRefsOwner, data: T) {
+        typeParameterRefsOwner.typeParameters.forEach { it.accept(this, data) }
+    }
+
+    override fun visitReceiverParameter(receiverParameter: FirReceiverParameter, data: T) {
+        receiverParameter.acceptChildren(this, data)
+    }
+
+    override fun visitContextReceiver(contextReceiver: FirContextReceiver, data: T) {
+        contextReceiver.acceptChildren(this, data)
+    }
+
+    override fun visitTypeParameter(typeParameter: FirTypeParameter, data: T) {
+        typeParameter.acceptChildren(this, data)
+    }
+
+    override fun visitFunction(function: FirFunction, data: T) {
+        visitCallableDeclaration(function, data)
+        function.valueParameters.forEach { it.accept(this, data) }
+    }
+
+    override fun visitSimpleFunction(simpleFunction: FirSimpleFunction, data: T) {
+        visitFunction(simpleFunction, data)
+    }
+
+    override fun visitConstructor(constructor: FirConstructor, data: T) {
+        visitFunction(constructor, data)
+    }
+
+    override fun visitErrorPrimaryConstructor(errorPrimaryConstructor: FirErrorPrimaryConstructor, data: T) {
+        visitFunction(errorPrimaryConstructor, data)
+    }
+
+    override fun visitPropertyAccessor(propertyAccessor: FirPropertyAccessor, data: T) {
+        visitFunction(propertyAccessor, data)
+    }
+
+    override fun visitVariable(variable: FirVariable, data: T) {
+        visitCallableDeclaration(variable, data)
+        variable.getter?.accept(this, data)
+        variable.setter?.accept(this, data)
+        variable.backingField?.accept(this, data)
+    }
+
+    override fun visitEnumEntry(enumEntry: FirEnumEntry, data: T) {
+        visitVariable(enumEntry, data)
+    }
+
+    override fun visitProperty(property: FirProperty, data: T) {
+        visitVariable(property, data)
+    }
+
+    override fun visitErrorProperty(errorProperty: FirErrorProperty, data: T) {
+        visitVariable(errorProperty, data)
+    }
+
+    override fun visitValueParameter(valueParameter: FirValueParameter, data: T) {
+        visitVariable(valueParameter, data)
+    }
+
+    override fun visitBackingField(backingField: FirBackingField, data: T) {
+        visitVariable(backingField, data)
+    }
+}

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

@@ -0,0 +1,38 @@
+/*
+ * 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.
+ */
+
+package org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve
+
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.forEachDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.FirScript
+
+/**
+ * In opposite to [NonLocalAnnotationVisitor] processes not only the target declaration,
+ * but also its nested declarations.
+ *
+ * @see NonLocalAnnotationVisitor
+ */
+internal abstract class RecursiveNonLocalAnnotationVisitor<T> : NonLocalAnnotationVisitor<T>() {
+    override fun visitFile(file: FirFile, data: T) {
+        super.visitFile(file, data)
+
+        file.annotationsContainer?.accept(this, data)
+        file.forEachDeclaration { it.accept(this, data) }
+    }
+
+    override fun visitScript(script: FirScript, data: T) {
+        super.visitScript(script, data)
+
+        script.forEachDeclaration { it.accept(this, data) }
+    }
+
+    override fun visitRegularClass(regularClass: FirRegularClass, data: T) {
+        super.visitRegularClass(regularClass, data)
+
+        regularClass.forEachDeclaration { it.accept(this, data) }
+    }
+}

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

@@ -8,8 +8,8 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.transformers
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveTarget
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.throwUnexpectedFirElementError
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.FirLazyBodiesCalculator
+import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.NonLocalAnnotationVisitor
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirSession
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.AnnotationVisitorVoid
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkAnnotationsAreResolved
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
@@ -24,7 +24,6 @@ import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
 import org.jetbrains.kotlin.fir.types.FirTypeProjection
-import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 
 internal object LLFirAnnotationArgumentsLazyResolver : LLFirLazyResolver(FirResolvePhase.ANNOTATION_ARGUMENTS) {
@@ -134,10 +133,10 @@ private class LLFirAnnotationArgumentsTargetResolver(resolveTarget: LLFirResolve
     }
 
     private class ForeignAnnotationsContext(val collection: MutableCollection<FirBasedSymbol<*>>, val currentSymbol: FirCallableSymbol<*>)
-    private object ForeignAnnotationsCollector : AnnotationVisitorVoid<ForeignAnnotationsContext>() {
-        override fun visitAnnotation(annotation: FirAnnotation, data: ForeignAnnotationsContext) {}
-        override fun visitAnnotationCall(annotationCall: FirAnnotationCall, data: ForeignAnnotationsContext) {
-            val symbolToPostpone = annotationCall.containingDeclarationSymbol.symbolToPostponeIfCanBeResolvedOnDemand() ?: return
+    private object ForeignAnnotationsCollector : NonLocalAnnotationVisitor<ForeignAnnotationsContext>() {
+        override fun processAnnotation(annotation: FirAnnotation, data: ForeignAnnotationsContext) {
+            if (annotation !is FirAnnotationCall) return
+            val symbolToPostpone = annotation.containingDeclarationSymbol.symbolToPostponeIfCanBeResolvedOnDemand() ?: return
             if (symbolToPostpone != data.currentSymbol) {
                 data.collection += symbolToPostpone
             }
@@ -242,19 +241,13 @@ internal object AnnotationArgumentsStateKeepers {
     }
 
     val DECLARATION: StateKeeper<FirElementWithResolveState, FirSession> = stateKeeper { target, session ->
-        val visitor = object : FirVisitorVoid() {
-            override fun visitElement(element: FirElement) {
-                when (element) {
-                    is FirDeclaration -> if (element !== target) return // Avoid nested declarations
-                    is FirAnnotation -> entity(element, ANNOTATION, session)
-                    is FirStatement -> return
-                }
-
-                element.acceptChildren(this)
+        val visitor = object : NonLocalAnnotationVisitor<Unit>() {
+            override fun processAnnotation(annotation: FirAnnotation, data: Unit) {
+                entity(annotation, ANNOTATION, session)
             }
         }
 
-        target.accept(visitor)
+        target.accept(visitor, Unit)
     }
 }
 

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

@@ -7,6 +7,7 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.util
 
+import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.NonLocalAnnotationVisitor
 import org.jetbrains.kotlin.fir.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
@@ -28,12 +29,8 @@ import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.FirTypeProjectionWithVariance
 import org.jetbrains.kotlin.fir.types.FirTypeRef
-import org.jetbrains.kotlin.fir.types.coneType
-import org.jetbrains.kotlin.fir.types.customAnnotations
-import org.jetbrains.kotlin.fir.types.forEachType
 import org.jetbrains.kotlin.fir.utils.exceptions.withFirEntry
 import org.jetbrains.kotlin.fir.utils.exceptions.withFirSymbolEntry
-import org.jetbrains.kotlin.fir.visitors.FirVisitor
 import org.jetbrains.kotlin.utils.exceptions.ExceptionAttachmentBuilder
 import org.jetbrains.kotlin.utils.exceptions.checkWithAttachment
 import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
@@ -248,36 +245,10 @@ internal fun FirAbstractBodyResolveTransformerDispatcher.checkAnnotationCallIsRe
     checkAnnotationIsResolved(annotationCall, annotationContainer)
 }
 
-private object AnnotationChecker : AnnotationVisitorVoid<FirAnnotationContainer>() {
-    override fun visitAnnotation(annotation: FirAnnotation, data: FirAnnotationContainer) {
+private object AnnotationChecker : NonLocalAnnotationVisitor<FirAnnotationContainer>() {
+    override fun processAnnotation(annotation: FirAnnotation, data: FirAnnotationContainer) {
         checkAnnotationIsResolved(annotation, data)
     }
-
-    override fun visitAnnotationCall(annotationCall: FirAnnotationCall, data: FirAnnotationContainer) {
-        checkAnnotationIsResolved(annotationCall, data)
-    }
-}
-
-/**
- * This visitor is responsible for visiting all annotations recursively
- */
-internal abstract class AnnotationVisitorVoid<T> : FirVisitor<Unit, T>() {
-    abstract override fun visitAnnotation(annotation: FirAnnotation, data: T)
-    abstract override fun visitAnnotationCall(annotationCall: FirAnnotationCall, data: T)
-
-    override fun visitResolvedTypeRef(resolvedTypeRef: FirResolvedTypeRef, data: T) {
-        resolvedTypeRef.acceptChildren(this, data)
-
-        resolvedTypeRef.coneType.forEachType {
-            it.type.attributes.customAnnotations.forEach { typeArgumentAnnotation ->
-                typeArgumentAnnotation.accept(this, data)
-            }
-        }
-    }
-
-    override fun visitElement(element: FirElement, data: T) {
-        element.acceptChildren(this, data)
-    }
 }
 
 internal fun checkAnnotationsAreResolved(annotationContainer: FirAnnotationContainer) {