FirSourceElement: introduce common 'lighterASTNode' property

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

@@ -37,5 +37,5 @@ 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)
 
-    return visitor.find(lightNode)?.let { it.toFirLightSourceElement(it.startOffset, it.endOffset, tree) }
+    return visitor.find(lighterASTNode)?.let { it.toFirLightSourceElement(it.startOffset, it.endOffset, tree) }
 }

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

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

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

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

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.element, kidsRef)
+                                parameterSourceElement.tree.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

@@ -19,7 +19,6 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
-import org.jetbrains.kotlin.fir.lightNode
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.psi.KtPrimaryConstructor
 
@@ -36,11 +35,11 @@ object FirConstructorInInterfaceChecker : FirBasicDeclarationChecker() {
 
     private fun FirSourceElement.hasPrimaryConstructor(): Boolean {
         val localPsi = psi
-        val localLightNode = lightNode
+        val localLightNode = lighterASTNode
 
         if (localPsi != null && localPsi !is PsiErrorElement) {
             return localPsi.hasPrimaryConstructor()
-        } else if (localLightNode != null && this is FirLightSourceElement) {
+        } else if (this is FirLightSourceElement) {
             return localLightNode.hasPrimaryConstructor(tree)
         }
 

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

@@ -19,7 +19,6 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
-import org.jetbrains.kotlin.fir.lightNode
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.psi.KtDelegatedSuperTypeEntry
 
@@ -36,11 +35,11 @@ object FirDelegationInInterfaceChecker : FirMemberDeclarationChecker() {
 
     private fun FirSourceElement.findSuperTypeDelegation(): Int {
         val localPsi = psi
-        val localLightNode = lightNode
+        val localLightNode = lighterASTNode
 
         if (localPsi != null && localPsi !is PsiErrorElement) {
             return localPsi.findSuperTypeDelegation()
-        } else if (localLightNode != null && this is FirLightSourceElement) {
+        } else if (this is FirLightSourceElement) {
             return localLightNode.findSuperTypeDelegation(tree)
         }
 

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

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

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

@@ -19,7 +19,6 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
-import org.jetbrains.kotlin.fir.lightNode
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.psi.KtSuperTypeCallEntry
 
@@ -36,11 +35,11 @@ object FirSupertypeInitializedInInterfaceChecker : FirMemberDeclarationChecker()
 
     private fun FirSourceElement.findSuperTypeCall(): Int {
         val localPsi = psi
-        val localLightNode = lightNode
+        val localLightNode = lighterASTNode
 
         if (localPsi != null && localPsi !is PsiErrorElement) {
             return localPsi.findSuperTypeCall()
-        } else if (localLightNode != null && this is FirLightSourceElement) {
+        } else if (this is FirLightSourceElement) {
             return localLightNode.findSuperTypeCall(tree)
         }
 

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

@@ -20,7 +20,6 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.declarations.FirConstructor
 import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-import org.jetbrains.kotlin.fir.lightNode
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.psi.KtSuperTypeCallEntry
 
@@ -40,11 +39,11 @@ object FirSupertypeInitializedWithoutPrimaryConstructor : FirMemberDeclarationCh
 
     private fun FirSourceElement.anySupertypeHasConstructorParentheses(): Boolean {
         val localPsi = psi
-        val localLightNode = lightNode
+        val localLightNode = lighterASTNode
 
         if (localPsi != null && localPsi !is PsiErrorElement) {
             return localPsi.anySupertypeHasConstructorParentheses()
-        } else if (localLightNode != null && this is FirLightSourceElement) {
+        } else if (this is FirLightSourceElement) {
             return localLightNode.anySupertypeHasConstructorParentheses(tree)
         }
 

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

@@ -18,7 +18,6 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
 import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
-import org.jetbrains.kotlin.fir.lightNode
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.psi.KtTypeArgumentList
 
@@ -38,11 +37,11 @@ object FirTypeArgumentsNotAllowedExpressionChecker : FirQualifiedAccessChecker()
 
     private fun FirSourceElement.hasAnyArguments(): Boolean {
         val localPsi = this.psi
-        val localLight = this.lightNode
+        val localLight = this.lighterASTNode
 
         if (localPsi != null && localPsi !is PsiErrorElement) {
             return localPsi.hasAnyArguments()
-        } else if (localLight != null && this is FirLightSourceElement) {
+        } else if (this is FirLightSourceElement) {
             return localLight.hasAnyArguments(this.tree)
         }
 

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

@@ -51,8 +51,8 @@ object CanBeReplacedWithOperatorAssignmentChecker : FirExpressionChecker<FirVari
                 needToReport = true
             }
         } else if (assignmentSource is FirLightSourceElement) {
-            val lValueLightTree = (lValue.source as FirLightSourceElement).element
+            val lValueLightTree = lValue.source!!.lighterASTNode
-            val rValueLightTree = (rValue.source as FirLightSourceElement).element
+            val rValueLightTree = rValue.source!!.lighterASTNode
             if (lightTreeMatcher(lValueLightTree, rValueLightTree, assignmentSource)) {
                 needToReport = true
             }

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

@@ -104,7 +104,7 @@ object CanBeValChecker : AbstractFirPropertyInitializationChecker() {
             val source = fir.source as FirLightSourceElement
             val tree = (fir.source as FirLightSourceElement).tree
             val children = Ref<Array<LighterASTNode?>>()
-            tree.getChildren(source.element, children)
+            tree.getChildren(source.lighterASTNode, children)
             children.get().filterNotNull().filter { it.tokenType == KtNodeTypes.DESTRUCTURING_DECLARATION_ENTRY }.size
         }
         else -> null

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

@@ -98,7 +98,7 @@ object RedundantExplicitTypeChecker : FirMemberDeclarationChecker() {
                     source.psi?.text
                 }
                 is FirLightSourceElement -> {
-                    (source as FirLightSourceElement).element.toString()
+                    source?.lighterASTNode?.toString()
                 }
                 else -> null
             }

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

@@ -19,7 +19,6 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.REDUNDANT_SINGLE_EXPRESSION_STRING_TEMPLATE
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.symbols.StandardClassIds
 import org.jetbrains.kotlin.fir.types.classId
 import org.jetbrains.kotlin.fir.types.coneType
@@ -38,12 +37,12 @@ object RedundantSingleExpressionStringTemplateChecker : FirBasicExpressionChecke
     }
 
     private fun FirStatement.stringParentChildrenCount(): Int? {
-        return when (source) {
+        return when (val source = source) {
             is FirPsiSourceElement<*> -> {
-                source.psi?.stringParentChildrenCount()
+                source.psi.stringParentChildrenCount()
             }
             is FirLightSourceElement -> {
-                (source as FirLightSourceElement).element.stringParentChildrenCount(source as FirLightSourceElement)
+                source.lighterASTNode.stringParentChildrenCount(source)
             }
             else -> null
         }

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

@@ -486,7 +486,7 @@ class DeclarationsConverter(
 
                     //parse data class
                     if (modifiers.isDataClass() && firPrimaryConstructor != null) {
-                        val zippedParameters = properties.map { it.source?.lightNode!! to it }
+                        val zippedParameters = properties.map { it.source!!.lighterASTNode to it }
                         DataClassMembersGenerator(
                             baseSession,
                             classNode,

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

@@ -44,7 +44,7 @@ class ValueParameter(
 
         return buildProperty {
             val parameterSource = firValueParameter.source as? FirLightSourceElement
-            val parameterNode = parameterSource?.lightNode
+            val parameterNode = parameterSource?.lighterASTNode
             source = parameterNode?.toFirLightSourceElement(
                 parameterSource.startOffset, parameterSource.endOffset, parameterSource.tree,
                 FirFakeSourceElementKind.PropertyFromParameter

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir
 
 import com.intellij.lang.LighterASTNode
+import com.intellij.lang.TreeBackedLighterAST
 import com.intellij.psi.PsiElement
 import com.intellij.psi.tree.IElementType
 import com.intellij.util.diff.FlyweightCapableTreeStructure
@@ -152,9 +153,11 @@ sealed class FirSourceElement {
     abstract val startOffset: Int
     abstract val endOffset: Int
     abstract val kind: FirSourceElementKind
+    abstract val lighterASTNode: LighterASTNode
 }
 
-
+// NB: in certain situations, psi.node could be null
+// Potentially exceptions can be provoked by elementType / lighterASTNode
 sealed class FirPsiSourceElement<out P : PsiElement>(val psi: P) : FirSourceElement() {
     override val elementType: IElementType
         get() = psi.node.elementType
@@ -164,6 +167,8 @@ sealed class FirPsiSourceElement<out P : PsiElement>(val psi: P) : FirSourceElem
 
     override val endOffset: Int
         get() = psi.textRange.endOffset
+
+    override val lighterASTNode by lazy { TreeBackedLighterAST.wrap(psi.node) }
 }
 
 class FirRealPsiSourceElement<out P : PsiElement>(psi: P) : FirPsiSourceElement<P>(psi) {
@@ -174,20 +179,20 @@ class FirFakeSourceElement<out P : PsiElement>(psi: P, override val kind: FirFak
 
 fun FirSourceElement.fakeElement(newKind: FirFakeSourceElementKind): FirSourceElement {
     return when (this) {
-        is FirLightSourceElement -> FirLightSourceElement(element, startOffset, endOffset, tree, newKind)
+        is FirLightSourceElement -> FirLightSourceElement(lighterASTNode, startOffset, endOffset, tree, newKind)
         is FirPsiSourceElement<*> -> FirFakeSourceElement(psi, newKind)
     }
 }
 
 class FirLightSourceElement(
-    val element: LighterASTNode,
+    override val lighterASTNode: LighterASTNode,
     override val startOffset: Int,
     override val endOffset: Int,
     val tree: FlyweightCapableTreeStructure<LighterASTNode>,
     override val kind: FirSourceElementKind = FirRealSourceElementKind,
 ) : FirSourceElement() {
     override val elementType: IElementType
-        get() = element.tokenType
+        get() = lighterASTNode.tokenType
 }
 
 val FirSourceElement?.psi: PsiElement? get() = (this as? FirPsiSourceElement<*>)?.psi
@@ -207,5 +212,3 @@ inline fun LighterASTNode.toFirLightSourceElement(
     tree: FlyweightCapableTreeStructure<LighterASTNode>,
     kind: FirSourceElementKind = FirRealSourceElementKind
 ): FirLightSourceElement = FirLightSourceElement(this, startOffset, endOffset, tree, kind)
-
-val FirSourceElement?.lightNode: LighterASTNode? get() = (this as? FirLightSourceElement)?.element