diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirIdePredicateBasedProvider.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirIdePredicateBasedProvider.kt index e482b25330c..f16a40f1498 100644 --- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirIdePredicateBasedProvider.kt +++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirIdePredicateBasedProvider.kt @@ -18,6 +18,8 @@ import org.jetbrains.kotlin.fir.caches.FirCache import org.jetbrains.kotlin.fir.caches.createCache import org.jetbrains.kotlin.fir.caches.firCachesFactory import org.jetbrains.kotlin.fir.caches.getValue +import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration +import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration import org.jetbrains.kotlin.fir.declarations.FirDeclaration import org.jetbrains.kotlin.fir.declarations.FirFile import org.jetbrains.kotlin.fir.extensions.AnnotationFqn @@ -29,6 +31,7 @@ import org.jetbrains.kotlin.fir.psi import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef import org.jetbrains.kotlin.fir.types.classId +import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitorVoid import org.jetbrains.kotlin.fir.visitors.FirVisitor import org.jetbrains.kotlin.name.ClassId import org.jetbrains.kotlin.psi.* @@ -130,6 +133,45 @@ internal class LLFirIdePredicateBasedProvider( override fun visitAncestorMetaAnnotatedWith(predicate: AncestorMetaAnnotatedWith, data: FirDeclaration): Boolean { return metaAnnotationsOnOuterDeclarations(data).any { it in predicate.metaAnnotations } } + + override fun visitParentAnnotatedWith(predicate: ParentAnnotatedWith, data: FirDeclaration): Boolean { + val parent = data.directParentDeclaration ?: return false + val parentPredicate = AnnotatedWith(predicate.annotations) + + return parentPredicate.accept(this, parent) + } + + override fun visitHasAnnotatedWith(predicate: HasAnnotatedWith, data: FirDeclaration): Boolean { + val childPredicate = AnnotatedWith(predicate.annotations) + + return data.anyDirectChildDeclarationMatches(childPredicate) + } + + override fun visitParentMetaAnnotatedWith(predicate: ParentMetaAnnotatedWith, data: FirDeclaration): Boolean { + val parent = data.directParentDeclaration ?: return false + val parentPredicate = MetaAnnotatedWith(predicate.annotations) + + return parentPredicate.accept(this, parent) + } + + override fun visitHasMetaAnnotatedWith(predicate: HasMetaAnnotatedWith, data: FirDeclaration): Boolean { + val childPredicate = MetaAnnotatedWith(predicate.annotations) + + return data.anyDirectChildDeclarationMatches(childPredicate) + } + + private val FirDeclaration.directParentDeclaration: FirDeclaration? + get() = getOwnersOfDeclaration(this)?.lastOrNull()?.fir + + private fun FirDeclaration.anyDirectChildDeclarationMatches(childPredicate: DeclarationPredicate): Boolean { + var result = false + + this.forEachDirectChildDeclaration { + result = result || childPredicate.accept(this@Matcher, it) + } + + return result + } } private fun annotationsOnDeclaration(declaration: FirDeclaration): Set { @@ -217,3 +259,26 @@ private inline fun FirFile.forEachElementWithContainers( accept(declarationsCollector, persistentListOf()) } + +/** + * Calls [action] on every direct child declaration of [this] declaration. + */ +private inline fun FirDeclaration.forEachDirectChildDeclaration(crossinline action: (child: FirDeclaration) -> Unit) { + this.acceptChildren(object : FirDefaultVisitorVoid() { + override fun visitElement(element: FirElement) { + // we must visit only direct children + } + + override fun visitFile(file: FirFile) { + action(file) + } + + override fun visitCallableDeclaration(callableDeclaration: FirCallableDeclaration) { + action(callableDeclaration) + } + + override fun visitClassLikeDeclaration(classLikeDeclaration: FirClassLikeDeclaration) { + action(classLikeDeclaration) + } + }) +}