From ce8489a8a5cebf9a24024698cc3fb174afa29a32 Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Thu, 22 Dec 2022 13:22:37 +0200
Subject: [PATCH] [FIR] Add ability to create accessor to declaration attribute
 on symbol

---
 .../jetbrains/kotlin/fir/plugin/TypeUtils.kt  | 37 +++++++++++++++++++
 .../resolve/providers/FirSymbolProvider.kt    |  4 ++
 .../declarations/FirDeclarationAttributes.kt  | 16 ++++++++
 3 files changed, 57 insertions(+)
 create mode 100644 compiler/fir/plugin-utils/src/org/jetbrains/kotlin/fir/plugin/TypeUtils.kt

diff --git a/compiler/fir/plugin-utils/src/org/jetbrains/kotlin/fir/plugin/TypeUtils.kt b/compiler/fir/plugin-utils/src/org/jetbrains/kotlin/fir/plugin/TypeUtils.kt
new file mode 100644
index 00000000000..dbc9cc3fc3a
--- /dev/null
+++ b/compiler/fir/plugin-utils/src/org/jetbrains/kotlin/fir/plugin/TypeUtils.kt
@@ -0,0 +1,37 @@
+/*
+ * Copyright 2010-2022 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.fir.FirSession
+import org.jetbrains.kotlin.fir.resolve.constructType
+import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
+import org.jetbrains.kotlin.fir.symbols.constructStarProjectedType
+import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.fir.types.ConeClassLikeType
+import org.jetbrains.kotlin.fir.types.ConeTypeProjection
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
+import org.jetbrains.kotlin.name.ClassId
+
+public fun ClassId.createConeType(
+    session: FirSession,
+    typeArguments: Array<ConeTypeProjection> = emptyArray(),
+    nullable: Boolean = false
+): ConeClassLikeType {
+    val symbol = session.symbolProvider.getClassLikeSymbolByClassId(this) as? FirClassSymbol<*>
+
+    return when {
+        symbol != null -> when {
+            typeArguments.isEmpty() -> symbol.constructStarProjectedType(isNullable = nullable)
+            else -> symbol.constructType(typeArguments, nullable)
+        }
+        else -> ConeClassLikeTypeImpl(
+            ConeClassLikeLookupTagImpl(this),
+            typeArguments,
+            nullable
+        )
+    }
+}
diff --git a/compiler/fir/providers/src/org/jetbrains/kotlin/fir/resolve/providers/FirSymbolProvider.kt b/compiler/fir/providers/src/org/jetbrains/kotlin/fir/resolve/providers/FirSymbolProvider.kt
index 325fe17a3f0..f37eecb1ead 100644
--- a/compiler/fir/providers/src/org/jetbrains/kotlin/fir/resolve/providers/FirSymbolProvider.kt
+++ b/compiler/fir/providers/src/org/jetbrains/kotlin/fir/resolve/providers/FirSymbolProvider.kt
@@ -85,5 +85,9 @@ inline fun <reified T : FirBasedSymbol<*>> FirSymbolProvider.getSymbolByTypeRef(
     return getSymbolByLookupTag(lookupTag) as? T
 }
 
+fun FirSymbolProvider.getRegularClassSymbolByClassId(classId: ClassId): FirRegularClassSymbol? {
+    return getClassLikeSymbolByClassId(classId) as? FirRegularClassSymbol
+}
+
 val FirSession.symbolProvider: FirSymbolProvider by FirSession.sessionComponentAccessor()
 val FirSession.dependenciesSymbolProvider: FirSymbolProvider by FirSession.sessionComponentAccessor<FirDependenciesSymbolProvider>()
diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationAttributes.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationAttributes.kt
index a6d6f1819bc..ba431a84a73 100644
--- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationAttributes.kt
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationAttributes.kt
@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.declarations
 
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.util.ConeTypeRegistry
 import org.jetbrains.kotlin.util.ArrayMap
 import org.jetbrains.kotlin.util.AttributeArrayOwner
@@ -54,6 +55,11 @@ object FirDeclarationDataRegistry : ConeTypeRegistry<FirDeclarationDataKey, Any>
         return DeclarationDataAccessor(generateAnyNullableAccessor(kClass), kClass)
     }
 
+    fun <K : FirDeclarationDataKey> symbolAccessor(key: K): SymbolDataAccessor {
+        val kClass = key::class
+        return SymbolDataAccessor(generateAnyNullableAccessor(kClass), kClass)
+    }
+
     fun <K : FirDeclarationDataKey, V : Any> attributesAccessor(key: K): ReadWriteProperty<FirDeclarationAttributes, V?> {
         val kClass = key::class
         return AttributeDataAccessor(generateNullableAccessor(kClass), kClass)
@@ -73,6 +79,16 @@ object FirDeclarationDataRegistry : ConeTypeRegistry<FirDeclarationDataKey, Any>
         }
     }
 
+    class SymbolDataAccessor(
+        private val dataAccessor: NullableArrayMapAccessor<FirDeclarationDataKey, Any, *>,
+        val key: KClass<out FirDeclarationDataKey>
+    ) {
+        operator fun <V> getValue(thisRef: FirBasedSymbol<*>, property: KProperty<*>): V? {
+            @Suppress("UNCHECKED_CAST")
+            return dataAccessor.getValue(thisRef.fir.attributes, property) as? V
+        }
+    }
+
     private class AttributeDataAccessor<V : Any>(
         val dataAccessor: NullableArrayMapAccessor<FirDeclarationDataKey, Any, V>,
         val key: KClass<out FirDeclarationDataKey>