FirSourceElement: introduce common 'treeStructure' property

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/FirSourceChildren.kt

@@ -35,7 +35,7 @@ private fun FirPsiSourceElement<*>.getChild(types: Set<IElementType>, index: Int
 }
 
 private fun FirLightSourceElement.getChild(types: Set<IElementType>, index: Int, depth: Int): FirSourceElement? {
-    val visitor = LighterTreeElementFinderByType(tree, types, index, depth)
+    val visitor = LighterTreeElementFinderByType(treeStructure, types, index, depth)
 
-    return visitor.find(lighterASTNode)?.let { it.toFirLightSourceElement(it.startOffset, it.endOffset, tree) }
+    return visitor.find(lighterASTNode)?.let { it.toFirLightSourceElement(it.startOffset, it.endOffset, treeStructure) }
 }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/FirModifierList.kt

@@ -28,9 +28,9 @@ fun FirSourceElement?.getModifierList(): FirModifierList? {
         is FirPsiSourceElement<*> -> (psi as? KtModifierListOwner)?.modifierList?.let { FirPsiModifierList(it) }
         is FirLightSourceElement -> {
             val kidsRef = Ref<Array<LighterASTNode?>>()
-            tree.getChildren(lighterASTNode, kidsRef)
+            treeStructure.getChildren(lighterASTNode, kidsRef)
             val modifierListNode = kidsRef.get().find { it?.tokenType == KtNodeTypes.MODIFIER_LIST } ?: return null
-            FirLightModifierList(modifierListNode, tree)
+            FirLightModifierList(modifierListNode, treeStructure)
         }
     }
 }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirAnnotationArgumentChecker.kt

@@ -241,9 +241,9 @@ object FirAnnotationArgumentChecker : FirBasicDeclarationChecker() {
             is FirPsiSourceElement<*> ->
                 source.psi.parent.toFirPsiSourceElement()
             is FirLightSourceElement -> {
-                val elementOfParent = source.tree.getParent(source.lighterASTNode) ?: source.lighterASTNode
+                val elementOfParent = source.treeStructure.getParent(source.lighterASTNode) ?: source.lighterASTNode
 
-                elementOfParent.toFirLightSourceElement(elementOfParent.startOffset, elementOfParent.endOffset, source.tree)
+                elementOfParent.toFirLightSourceElement(elementOfParent.startOffset, elementOfParent.endOffset, source.treeStructure)
             }
             else ->
                 source

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirAnnotationClassDeclarationChecker.kt

@@ -53,7 +53,7 @@ object FirAnnotationClassDeclarationChecker : FirBasicDeclarationChecker() {
                             }
                             is FirLightSourceElement -> {
                                 val kidsRef = Ref<Array<LighterASTNode?>>()
-                                parameterSourceElement.tree.getChildren(parameterSourceElement.lighterASTNode, kidsRef)
+                                parameterSourceElement.treeStructure.getChildren(parameterSourceElement.lighterASTNode, kidsRef)
 
                                 if (kidsRef.get().any { it?.tokenType == VAR_KEYWORD })
                                     reporter.report(parameterSourceElement, FirErrors.VAR_ANNOTATION_PARAMETER)

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirConstructorInInterfaceChecker.kt

@@ -40,7 +40,7 @@ object FirConstructorInInterfaceChecker : FirBasicDeclarationChecker() {
         if (localPsi != null && localPsi !is PsiErrorElement) {
             return localPsi.hasPrimaryConstructor()
         } else if (this is FirLightSourceElement) {
-            return localLightNode.hasPrimaryConstructor(tree)
+            return localLightNode.hasPrimaryConstructor(treeStructure)
         }
 
         return false

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirDelegationInInterfaceChecker.kt

@@ -40,7 +40,7 @@ object FirDelegationInInterfaceChecker : FirMemberDeclarationChecker() {
         if (localPsi != null && localPsi !is PsiErrorElement) {
             return localPsi.findSuperTypeDelegation()
         } else if (this is FirLightSourceElement) {
-            return localLightNode.findSuperTypeDelegation(tree)
+            return localLightNode.findSuperTypeDelegation(treeStructure)
         }
 
         return -1

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirExposedVisibilityDeclarationChecker.kt

@@ -234,9 +234,9 @@ object FirExposedVisibilityDeclarationChecker : FirMemberDeclarationChecker() {
         is FirPsiSourceElement<*> -> (this.psi as? PsiNameIdentifierOwner)?.nameIdentifier?.toFirPsiSourceElement()
         is FirLightSourceElement -> {
             val kidsRef = Ref<Array<LighterASTNode?>>()
-            this.tree.getChildren(lighterASTNode, kidsRef)
+            this.treeStructure.getChildren(lighterASTNode, kidsRef)
             val identifier = kidsRef.get().find { it?.tokenType == KtTokens.IDENTIFIER }
-            identifier?.toFirLightSourceElement(this.tree.getStartOffset(identifier), this.tree.getEndOffset(identifier), this.tree)
+            identifier?.toFirLightSourceElement(this.treeStructure.getStartOffset(identifier), this.treeStructure.getEndOffset(identifier), this.treeStructure)
         }
     }
 

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirSupertypeInitializedInInterfaceChecker.kt

@@ -40,7 +40,7 @@ object FirSupertypeInitializedInInterfaceChecker : FirMemberDeclarationChecker()
         if (localPsi != null && localPsi !is PsiErrorElement) {
             return localPsi.findSuperTypeCall()
         } else if (this is FirLightSourceElement) {
-            return localLightNode.findSuperTypeCall(tree)
+            return localLightNode.findSuperTypeCall(treeStructure)
         }
 
         return -1

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirSupertypeInitializedWithoutPrimaryConstructor.kt

@@ -44,7 +44,7 @@ object FirSupertypeInitializedWithoutPrimaryConstructor : FirMemberDeclarationCh
         if (localPsi != null && localPsi !is PsiErrorElement) {
             return localPsi.anySupertypeHasConstructorParentheses()
         } else if (this is FirLightSourceElement) {
-            return localLightNode.anySupertypeHasConstructorParentheses(tree)
+            return localLightNode.anySupertypeHasConstructorParentheses(treeStructure)
         }
 
         return false

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirTypeArgumentsNotAllowedExpressionChecker.kt

@@ -42,7 +42,7 @@ object FirTypeArgumentsNotAllowedExpressionChecker : FirQualifiedAccessChecker()
         if (localPsi != null && localPsi !is PsiErrorElement) {
             return localPsi.hasAnyArguments()
         } else if (this is FirLightSourceElement) {
-            return localLight.hasAnyArguments(this.tree)
+            return localLight.hasAnyArguments(this.treeStructure)
         }
 
         return false

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/extended/CanBeReplacedWithOperatorAssignmentChecker.kt

@@ -71,7 +71,7 @@ object CanBeReplacedWithOperatorAssignmentChecker : FirExpressionChecker<FirVari
         source: FirLightSourceElement,
         prevOperator: LighterASTNode? = null
     ): Boolean {
-        val tree = source.tree
+        val tree = source.treeStructure
         val childrenNullable = Ref<Array<LighterASTNode?>>()
         tree.getChildren(expression, childrenNullable)
         val children = childrenNullable.get().filterNotNull()

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/extended/CanBeValChecker.kt

@@ -102,7 +102,7 @@ object CanBeValChecker : AbstractFirPropertyInitializationChecker() {
         is FirPsiSourceElement<*> -> fir.psi?.children?.size?.minus(1) // -1 cuz we don't need expression node after equals operator
         is FirLightSourceElement -> {
             val source = fir.source as FirLightSourceElement
-            val tree = (fir.source as FirLightSourceElement).tree
+            val tree = (fir.source as FirLightSourceElement).treeStructure
             val children = Ref<Array<LighterASTNode?>>()
             tree.getChildren(source.lighterASTNode, children)
             children.get().filterNotNull().filter { it.tokenType == KtNodeTypes.DESTRUCTURING_DECLARATION_ENTRY }.size

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/extended/RedundantSingleExpressionStringTemplateChecker.kt

@@ -54,10 +54,10 @@ object RedundantSingleExpressionStringTemplateChecker : FirBasicExpressionChecke
     }
 
     private fun LighterASTNode.stringParentChildrenCount(source: FirLightSourceElement): Int? {
-        val parent = source.tree.getParent(this)
+        val parent = source.treeStructure.getParent(this)
         return if (parent?.tokenType == KtNodeTypes.STRING_TEMPLATE) {
             val childrenOfParent = Ref<Array<LighterASTNode>>()
-            source.tree.getChildren(parent!!, childrenOfParent)
+            source.treeStructure.getChildren(parent!!, childrenOfParent)
             childrenOfParent.get().filter { it is PsiBuilder.Marker }.size
         } else {
             parent?.stringParentChildrenCount(source)

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/fir/ValueParameter.kt

@@ -46,7 +46,7 @@ class ValueParameter(
             val parameterSource = firValueParameter.source as? FirLightSourceElement
             val parameterNode = parameterSource?.lighterASTNode
             source = parameterNode?.toFirLightSourceElement(
-                parameterSource.startOffset, parameterSource.endOffset, parameterSource.tree,
+                parameterSource.startOffset, parameterSource.endOffset, parameterSource.treeStructure,
                 FirFakeSourceElementKind.PropertyFromParameter
             )
             this.session = session

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSourceElement.kt

@@ -7,7 +7,9 @@ package org.jetbrains.kotlin.fir
 
 import com.intellij.lang.LighterASTNode
 import com.intellij.lang.TreeBackedLighterAST
+import com.intellij.openapi.util.Ref
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
 import com.intellij.psi.tree.IElementType
 import com.intellij.util.diff.FlyweightCapableTreeStructure
 
@@ -154,6 +156,7 @@ sealed class FirSourceElement {
     abstract val endOffset: Int
     abstract val kind: FirSourceElementKind
     abstract val lighterASTNode: LighterASTNode
+    abstract val treeStructure: FlyweightCapableTreeStructure<LighterASTNode>
 }
 
 // NB: in certain situations, psi.node could be null
@@ -169,6 +172,37 @@ sealed class FirPsiSourceElement<out P : PsiElement>(val psi: P) : FirSourceElem
         get() = psi.textRange.endOffset
 
     override val lighterASTNode by lazy { TreeBackedLighterAST.wrap(psi.node) }
+
+    override val treeStructure: FlyweightCapableTreeStructure<LighterASTNode> by lazy { WrappedTreeStructure(psi.containingFile) }
+
+    private class WrappedTreeStructure(file: PsiFile) : FlyweightCapableTreeStructure<LighterASTNode> {
+        private val lighterAST = TreeBackedLighterAST(file.node)
+
+        private fun LighterASTNode.unwrap() = lighterAST.unwrap(this)
+
+        override fun toString(node: LighterASTNode): CharSequence = node.unwrap().text
+
+        override fun getRoot(): LighterASTNode = lighterAST.root
+
+        override fun getParent(node: LighterASTNode): LighterASTNode? = node.unwrap().psi.parent?.node?.let { TreeBackedLighterAST.wrap(it) }
+
+        override fun getChildren(node: LighterASTNode, nodesRef: Ref<Array<LighterASTNode>>): Int {
+            val children = node.unwrap().psi.children
+            if (children.isEmpty()) {
+                nodesRef.set(LighterASTNode.EMPTY_ARRAY)
+            } else {
+                nodesRef.set(children.map { TreeBackedLighterAST.wrap(it.node) }.toTypedArray())
+            }
+            return children.size
+        }
+
+        override fun disposeChildren(p0: Array<out LighterASTNode>?, p1: Int) {
+        }
+
+        override fun getStartOffset(node: LighterASTNode): Int = node.unwrap().startOffset
+
+        override fun getEndOffset(node: LighterASTNode): Int = node.unwrap().let { it.startOffset + it.textLength - 1 }
+    }
 }
 
 class FirRealPsiSourceElement<out P : PsiElement>(psi: P) : FirPsiSourceElement<P>(psi) {
@@ -179,7 +213,7 @@ class FirFakeSourceElement<out P : PsiElement>(psi: P, override val kind: FirFak
 
 fun FirSourceElement.fakeElement(newKind: FirFakeSourceElementKind): FirSourceElement {
     return when (this) {
-        is FirLightSourceElement -> FirLightSourceElement(lighterASTNode, startOffset, endOffset, tree, newKind)
+        is FirLightSourceElement -> FirLightSourceElement(lighterASTNode, startOffset, endOffset, treeStructure, newKind)
         is FirPsiSourceElement<*> -> FirFakeSourceElement(psi, newKind)
     }
 }
@@ -188,7 +222,7 @@ class FirLightSourceElement(
     override val lighterASTNode: LighterASTNode,
     override val startOffset: Int,
     override val endOffset: Int,
-    val tree: FlyweightCapableTreeStructure<LighterASTNode>,
+    override val treeStructure: FlyweightCapableTreeStructure<LighterASTNode>,
     override val kind: FirSourceElementKind = FirRealSourceElementKind,
 ) : FirSourceElement() {
     override val elementType: IElementType