From 25ab863af7679427b0e15e3fea30513291e3531a Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Wed, 18 Mar 2020 15:17:57 +0300
Subject: [PATCH] Add FlyweightCapableTreeStructure to FirLightSourceElement &
 use it

---
 .../fir/lightTree/converter/BaseConverter.kt   |  2 +-
 .../jetbrains/kotlin/fir/FirModifierList.kt    | 18 +++++++++++++++---
 .../jetbrains/kotlin/fir/FirSourceElement.kt   | 18 +++++++++++-------
 3 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/compiler/fir/lightTree/src/org/jetbrains/kotlin/fir/lightTree/converter/BaseConverter.kt b/compiler/fir/lightTree/src/org/jetbrains/kotlin/fir/lightTree/converter/BaseConverter.kt
index 86b089e9036..170d1b9fa87 100644
--- a/compiler/fir/lightTree/src/org/jetbrains/kotlin/fir/lightTree/converter/BaseConverter.kt
+++ b/compiler/fir/lightTree/src/org/jetbrains/kotlin/fir/lightTree/converter/BaseConverter.kt
@@ -32,7 +32,7 @@ open class BaseConverter(
     protected val implicitType = buildImplicitTypeRef()
 
     override fun LighterASTNode.toFirSourceElement(): FirSourceElement {
-        return FirLightSourceElement(this)
+        return FirLightSourceElement(this, tree)
     }
 
     override val LighterASTNode.elementType: IElementType
diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirModifierList.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirModifierList.kt
index e00dbbcc08e..46268dc48da 100644
--- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirModifierList.kt
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirModifierList.kt
@@ -7,7 +7,10 @@ package org.jetbrains.kotlin.fir
 
 import com.intellij.lang.ASTNode
 import com.intellij.lang.LighterASTNode
+import com.intellij.openapi.util.Ref
+import com.intellij.psi.PsiElement
 import com.intellij.psi.tree.TokenSet
+import com.intellij.util.diff.FlyweightCapableTreeStructure
 import org.jetbrains.kotlin.lexer.KtModifierKeywordToken
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtModifierList
@@ -26,13 +29,22 @@ class FirPsiModifierList(val modifierList: KtModifierList) : FirModifierList() {
 
 }
 
-class FirLightModifierList(val modifierList: LighterASTNode) : FirModifierList() {
+class FirLightModifierList(val modifierList: LighterASTNode, val tree: FlyweightCapableTreeStructure<LighterASTNode>) : FirModifierList() {
     override val modifiers: List<FirLightModifier>
-        get() = TODO()
+        get() {
+            val kidsRef = Ref<Array<LighterASTNode?>>()
+            tree.getChildren(modifierList, kidsRef)
+            val modifierNodes = kidsRef.get()
+            return modifierNodes.filterNotNull()
+                .filter { it.tokenType is KtModifierKeywordToken }
+                .map { FirLightModifier(it, it.tokenType as KtModifierKeywordToken) }
+        }
 }
 
 sealed class FirModifier<Node : Any>(val node: Node, val token: KtModifierKeywordToken)
 
 class FirPsiModifier(node: ASTNode, token: KtModifierKeywordToken) : FirModifier<ASTNode>(node, token)
 
-class FirLightModifier(node: LighterASTNode, token: KtModifierKeywordToken) : FirModifier<LighterASTNode>(node, token)
\ No newline at end of file
+class FirLightModifier(node: LighterASTNode, token: KtModifierKeywordToken) : FirModifier<LighterASTNode>(node, token)
+
+val FirModifier<*>.psi: PsiElement? get() = (this as? FirPsiModifier)?.node?.psi
\ No newline at end of file
diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSourceElement.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSourceElement.kt
index afdd133433c..234d4479a74 100644
--- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSourceElement.kt
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSourceElement.kt
@@ -6,14 +6,16 @@
 package org.jetbrains.kotlin.fir
 
 import com.intellij.lang.LighterASTNode
+import com.intellij.openapi.util.Ref
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.lexer.KtModifierKeywordToken
+import com.intellij.util.diff.FlyweightCapableTreeStructure
+import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.psi.KtModifierListOwner
 
 sealed class FirSourceElement
 
 class FirPsiSourceElement(val psi: PsiElement) : FirSourceElement()
-class FirLightSourceElement(val element: LighterASTNode) : FirSourceElement()
+class FirLightSourceElement(val element: LighterASTNode, val tree: FlyweightCapableTreeStructure<LighterASTNode>) : FirSourceElement()
 
 val FirSourceElement?.psi: PsiElement? get() = (this as? FirPsiSourceElement)?.psi
 
@@ -22,15 +24,17 @@ val FirElement.psi: PsiElement? get() = (source as? FirPsiSourceElement)?.psi
 @Suppress("NOTHING_TO_INLINE")
 inline fun PsiElement.toFirSourceElement(): FirPsiSourceElement = FirPsiSourceElement(this)
 
-@Suppress("NOTHING_TO_INLINE")
-inline fun LighterASTNode.toFirSourceElement(): FirLightSourceElement = FirLightSourceElement(this)
-
 val FirSourceElement?.lightNode: LighterASTNode? get() = (this as? FirLightSourceElement)?.element
 
 fun FirSourceElement?.getModifierList(): FirModifierList? {
     return when (this) {
-        is FirPsiSourceElement -> (psi as? KtModifierListOwner)?.modifierList?.let { FirPsiModifierList(it) }
-        is FirLightSourceElement -> FirLightModifierList(element)
         null -> null
+        is FirPsiSourceElement -> (psi as? KtModifierListOwner)?.modifierList?.let { FirPsiModifierList(it) }
+        is FirLightSourceElement -> {
+            val kidsRef = Ref<Array<LighterASTNode?>>()
+            tree.getChildren(element, kidsRef)
+            val modifierListNode = kidsRef.get().find { it?.tokenType == KtNodeTypes.MODIFIER_LIST } ?: return null
+            FirLightModifierList(modifierListNode, tree)
+        }
     }
 }