From c7e3da870919337a606c7dc492c694e6f8e439c3 Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Fri, 13 May 2022 16:20:42 +0300
Subject: [PATCH] [FIR] Extend DeclarationPredicate DSL

---
 .../kotlin/fir/extensions/FirPredicateBasedProvider.kt |  9 ++++++++-
 .../fir/extensions/predicate/DeclarationPredicate.kt   | 10 ++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirPredicateBasedProvider.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirPredicateBasedProvider.kt
index dd8d42e684b..871268ca937 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirPredicateBasedProvider.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirPredicateBasedProvider.kt
@@ -28,12 +28,19 @@ abstract class FirPredicateBasedProvider : FirSessionComponent {
     abstract fun fileHasPluginAnnotations(file: FirFile): Boolean
     abstract fun matches(predicate: DeclarationPredicate, declaration: FirDeclaration): Boolean
 
-    open fun registerAnnotatedDeclaration(declaration: FirDeclaration, owners: PersistentList<FirDeclaration>) {
+    fun matches(predicate: DeclarationPredicate, declaration: FirBasedSymbol<*>): Boolean {
+        return matches(predicate, declaration.fir)
     }
 
     fun matches(predicates: List<DeclarationPredicate>, declaration: FirDeclaration): Boolean {
         return predicates.any { matches(it, declaration) }
     }
+
+    fun matches(predicates: List<DeclarationPredicate>, declaration: FirBasedSymbol<*>): Boolean {
+        return matches(predicates, declaration.fir)
+    }
+
+    open fun registerAnnotatedDeclaration(declaration: FirDeclaration, owners: PersistentList<FirDeclaration>) {}
 }
 
 @NoMutableState
diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/extensions/predicate/DeclarationPredicate.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/extensions/predicate/DeclarationPredicate.kt
index f4ac76d1109..fae1020a3f7 100644
--- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/extensions/predicate/DeclarationPredicate.kt
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/extensions/predicate/DeclarationPredicate.kt
@@ -115,6 +115,7 @@ class UnderMetaAnnotated(metaAnnotations: Set<AnnotationFqn>) : MetaAnnotated(me
 infix fun DeclarationPredicate.or(other: DeclarationPredicate): DeclarationPredicate = DeclarationPredicate.Or(this, other)
 infix fun DeclarationPredicate.and(other: DeclarationPredicate): DeclarationPredicate = DeclarationPredicate.And(this, other)
 
+// ------------------- varargs -------------------
 fun under(vararg annotations: AnnotationFqn): DeclarationPredicate = UnderAnnotatedWith(annotations.toSet())
 fun has(vararg annotations: AnnotationFqn): DeclarationPredicate = AnnotatedWith(annotations.toSet())
 fun metaUnder(vararg metaAnnotations: AnnotationFqn): DeclarationPredicate = UnderMetaAnnotated(metaAnnotations.toSet())
@@ -122,3 +123,12 @@ fun metaHas(vararg metaAnnotations: AnnotationFqn): DeclarationPredicate = Annot
 
 fun hasOrUnder(vararg annotations: AnnotationFqn): DeclarationPredicate = has(*annotations) or under(*annotations)
 fun metaHasOrUnder(vararg metaAnnotations: AnnotationFqn): DeclarationPredicate = metaHas(*metaAnnotations) or metaUnder(*metaAnnotations)
+
+// ------------------- colelctions -------------------
+fun under(annotations: Collection<AnnotationFqn>): DeclarationPredicate = UnderAnnotatedWith(annotations.toSet())
+fun has(annotations: Collection<AnnotationFqn>): DeclarationPredicate = AnnotatedWith(annotations.toSet())
+fun metaUnder(metaAnnotations: Collection<AnnotationFqn>): DeclarationPredicate = UnderMetaAnnotated(metaAnnotations.toSet())
+fun metaHas(metaAnnotations: Collection<AnnotationFqn>): DeclarationPredicate = AnnotatedWithMeta(metaAnnotations.toSet())
+
+fun hasOrUnder(annotations: Collection<AnnotationFqn>): DeclarationPredicate = has(annotations) or under(annotations)
+fun metaHasOrUnder(metaAnnotations: Collection<AnnotationFqn>): DeclarationPredicate = metaHas(metaAnnotations) or metaUnder(metaAnnotations)