[FIR-PLUGIN] Add status transformer for default visibility

plugins/fir/fir-plugin-prototype/plugin-annotations/src/org/jetbrains/kotlin/fir/plugin/annotations.kt

@@ -17,3 +17,9 @@ annotation class D
 annotation class E
 annotation class F
 annotation class G
+
+annotation class AllPublic(val visibility: Visibility)
+
+enum class Visibility {
+    Public, Internal, Private, Protected
+}

plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/AllOpenVisibilityTransformer.kt

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2010-2020 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.fir.plugin
+
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.descriptors.Visibility
+import org.jetbrains.kotlin.fir.FirEffectiveVisibility
+import org.jetbrains.kotlin.fir.FirEffectiveVisibilityImpl
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirAnnotatedDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationStatus
+import org.jetbrains.kotlin.fir.declarations.FirPluginKey
+import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
+import org.jetbrains.kotlin.fir.expressions.arguments
+import org.jetbrains.kotlin.fir.extensions.FirStatusTransformerExtension
+import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
+import org.jetbrains.kotlin.fir.extensions.predicate.hasOrUnder
+import org.jetbrains.kotlin.fir.extensions.transform
+import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
+import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
+import org.jetbrains.kotlin.fir.types.ConeClassLikeType
+import org.jetbrains.kotlin.fir.types.coneTypeSafe
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+
+class AllOpenVisibilityTransformer(session: FirSession) : FirStatusTransformerExtension(session) {
+    companion object {
+        private val AllPublicClassId = ClassId(FqName("org.jetbrains.kotlin.fir.plugin"), Name.identifier("AllPublic"))
+        private val VisibilityClassId = ClassId(FqName("org.jetbrains.kotlin.fir.plugin"), Name.identifier("Visibility"))
+
+        private val PublicName = Name.identifier("Public")
+        private val InternalName = Name.identifier("Internal")
+        private val PrivateName = Name.identifier("Private")
+        private val ProtectedName = Name.identifier("Protected")
+    }
+
+    override fun transformStatus(
+        declaration: FirDeclaration,
+        owners: List<FirAnnotatedDeclaration>,
+        status: FirDeclarationStatus
+    ): FirDeclarationStatus {
+        val visibility = findVisibility(declaration, owners) ?: return status
+        if (visibility == status.visibility) return status
+        return status.transform(visibility = visibility)
+    }
+
+    private fun findVisibility(declaration: FirDeclaration, owners: List<FirAnnotatedDeclaration>): Visibility? {
+        (declaration as? FirAnnotatedDeclaration)?.visibilityFromAnnotation()?.let { return it }
+        for (owner in owners) {
+            owner.visibilityFromAnnotation()?.let { return it }
+        }
+        return null
+    }
+
+    private fun FirAnnotatedDeclaration.visibilityFromAnnotation(): Visibility? {
+        val annotation = annotations.firstOrNull {
+            it.annotationTypeRef.coneTypeSafe<ConeClassLikeType>()?.lookupTag?.classId == AllPublicClassId
+        } ?: return null
+        val argument = annotation.arguments.firstOrNull() as? FirQualifiedAccessExpression ?: return null
+        val symbol = (argument.calleeReference as? FirResolvedNamedReference)?.resolvedSymbol as? FirVariableSymbol<*> ?: return null
+        val name = symbol.callableId.takeIf { it.classId == VisibilityClassId }?.callableName ?: return null
+        return when (name) {
+            PublicName -> Visibilities.PUBLIC
+            InternalName -> Visibilities.INTERNAL
+            PrivateName -> Visibilities.PRIVATE
+            ProtectedName -> Visibilities.PROTECTED
+            else -> null
+        }
+    }
+
+    override val predicate: DeclarationPredicate = hasOrUnder(AllPublicClassId.asSingleFqName())
+
+    override val key: FirPluginKey
+        get() = AllOpenPluginKey
+}

plugins/fir/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/FirAllOpenComponentRegistrar.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.plugin.generators.*
 class FirAllOpenComponentRegistrar : FirExtensionRegistrar() {
     override fun ExtensionRegistrarContext.configurePlugin() {
         +::AllOpenStatusTransformer
+        +::AllOpenVisibilityTransformer
         +::AllOpenSupertypeGenerator
 
         // Declaration generators

plugins/fir/fir-plugin-prototype/testData/status/visibilityTransformation.kt

@@ -0,0 +1,30 @@
+import org.jetbrains.kotlin.fir.plugin.AllPublic
+import org.jetbrains.kotlin.fir.plugin.Visibility
+
+@AllPublic(Visibility.Protected)
+class A {
+    val x: String
+
+    fun foo() {}
+
+    class Nested {
+        fun bar() {
+
+        }
+    }
+}
+
+@AllPublic(Visibility.Private)
+class B {
+    val x: String
+
+    fun foo() {
+
+    }
+
+    class Nested {
+        fun bar() {
+
+        }
+    }
+}

plugins/fir/fir-plugin-prototype/testData/status/visibilityTransformation.txt

@@ -0,0 +1,45 @@
+FILE: visibilityTransformation.kt
+    @R|org/jetbrains/kotlin/fir/plugin/AllPublic|(Q|org/jetbrains/kotlin/fir/plugin/Visibility|.R|org/jetbrains/kotlin/fir/plugin/Visibility.Protected|) protected final class A : R|kotlin/Any| {
+        protected constructor(): R|A| {
+            super<R|kotlin/Any|>()
+        }
+
+        protected final val x: R|kotlin/String|
+            public get(): R|kotlin/String|
+
+        protected final fun foo(): R|kotlin/Unit| {
+        }
+
+        protected final class Nested : R|kotlin/Any| {
+            protected constructor(): R|A.Nested| {
+                super<R|kotlin/Any|>()
+            }
+
+            protected final fun bar(): R|kotlin/Unit| {
+            }
+
+        }
+
+    }
+    @R|org/jetbrains/kotlin/fir/plugin/AllPublic|(Q|org/jetbrains/kotlin/fir/plugin/Visibility|.R|org/jetbrains/kotlin/fir/plugin/Visibility.Private|) private final class B : R|kotlin/Any| {
+        private constructor(): R|B| {
+            super<R|kotlin/Any|>()
+        }
+
+        private final val x: R|kotlin/String|
+            public get(): R|kotlin/String|
+
+        private final fun foo(): R|kotlin/Unit| {
+        }
+
+        private final class Nested : R|kotlin/Any| {
+            private constructor(): R|B.Nested| {
+                super<R|kotlin/Any|>()
+            }
+
+            private final fun bar(): R|kotlin/Unit| {
+            }
+
+        }
+
+    }

plugins/fir/fir-plugin-prototype/tests/org/jetbrains/kotlin/fir/plugin/FirAllOpenDiagnosticTestGenerated.java

@@ -100,6 +100,11 @@ public class FirAllOpenDiagnosticTestGenerated extends AbstractFirAllOpenDiagnos
         public void testSimpleAnnotation() throws Exception {
             runTest("plugins/fir/fir-plugin-prototype/testData/status/simpleAnnotation.kt");
         }
+
+        @TestMetadata("visibilityTransformation.kt")
+        public void testVisibilityTransformation() throws Exception {
+            runTest("plugins/fir/fir-plugin-prototype/testData/status/visibilityTransformation.kt");
+        }
     }
 
     @TestMetadata("plugins/fir/fir-plugin-prototype/testData/supertypes")