From 08b63534311164465f2733af677e505e4690f3bb Mon Sep 17 00:00:00 2001
From: Roman Golyshev <roman.golyshev@jetbrains.com>
Date: Thu, 25 Nov 2021 18:26:03 +0300
Subject: [PATCH] FIR: Add IDE implementation for
 `FirRegisteredPluginAnnotations`

It works on `KotlinAnnotationsResolver`, which means that it uses
PSI indices to search for the annotations
---
 .../LLFirIdeRegisteredPluginAnnotations.kt    | 66 +++++++++++++++++++
 1 file changed, 66 insertions(+)
 create mode 100644 analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirIdeRegisteredPluginAnnotations.kt

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirIdeRegisteredPluginAnnotations.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirIdeRegisteredPluginAnnotations.kt
new file mode 100644
index 00000000000..5b19cf4d4cd
--- /dev/null
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirIdeRegisteredPluginAnnotations.kt
@@ -0,0 +1,66 @@
+/*
+ * Copyright 2010-2021 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.providers
+
+import org.jetbrains.kotlin.analysis.providers.KotlinAnnotationsResolver
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.caches.*
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.extensions.AbstractFirRegisteredPluginAnnotations
+import org.jetbrains.kotlin.fir.extensions.AnnotationFqn
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.KtClass
+
+internal class LLFirIdeRegisteredPluginAnnotations(
+    session: FirSession,
+    private val annotationsResolver: KotlinAnnotationsResolver
+) : AbstractFirRegisteredPluginAnnotations(session) {
+
+    private val annotationsFromPlugins: MutableSet<AnnotationFqn> = mutableSetOf()
+
+    override val annotations: Set<AnnotationFqn>
+        get() = allAnnotationsCache.getValue()
+
+    private val allAnnotationsCache: FirLazyValue<Set<AnnotationFqn>, Nothing?> = session.firCachesFactory.createLazyValue {
+        // at this point, both metaAnnotations and annotationsFromPlugins should be collected
+        val result = metaAnnotations.flatMapTo(mutableSetOf()) { getAnnotationsWithMetaAnnotation(it) }
+
+        if (result.isEmpty()) {
+            annotationsFromPlugins
+        } else {
+            result += annotationsFromPlugins
+            result
+        }
+    }
+
+    // MetaAnnotation -> Annotations
+    private val annotationsWithMetaAnnotationCache: FirCache<AnnotationFqn, Set<AnnotationFqn>, Nothing?> =
+        session.firCachesFactory.createCache { metaAnnotation -> collectAnnotationsWithMetaAnnotation(metaAnnotation) }
+
+    override fun getAnnotationsWithMetaAnnotation(metaAnnotation: AnnotationFqn): Collection<AnnotationFqn> {
+        return annotationsWithMetaAnnotationCache.getValue(metaAnnotation)
+    }
+
+    private fun collectAnnotationsWithMetaAnnotation(metaAnnotation: AnnotationFqn): Set<FqName> {
+        val annotatedDeclarations = annotationsResolver.declarationsByAnnotation(ClassId.topLevel(metaAnnotation))
+
+        return annotatedDeclarations
+            .asSequence()
+            .filterIsInstance<KtClass>()
+            .filter { it.isAnnotation() && it.isTopLevel() }
+            .mapNotNull { it.fqName }
+            .toSet()
+    }
+
+    override fun saveAnnotationsFromPlugin(annotations: Collection<AnnotationFqn>) {
+        annotationsFromPlugins += annotations
+    }
+
+    override fun registerUserDefinedAnnotation(metaAnnotation: AnnotationFqn, annotationClasses: Collection<FirRegularClass>) {
+        error("This method should never be called in IDE mode")
+    }
+}