FIR checkers: simplify hasVal / hasVar source element checks

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

@@ -5,18 +5,11 @@
 
 package org.jetbrains.kotlin.fir.analysis.checkers.declaration
 
-import com.intellij.lang.LighterASTNode
-import com.intellij.openapi.util.Ref
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.descriptors.ClassKind.ANNOTATION_CLASS
 import org.jetbrains.kotlin.descriptors.ClassKind.ENUM_CLASS
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
-import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
-import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnosticFactory0
-import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.FirLightSourceElement
-import org.jetbrains.kotlin.fir.FirPsiSourceElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.resolve.toSymbol
@@ -27,10 +20,8 @@ import org.jetbrains.kotlin.fir.symbols.StandardClassIds.unsignedTypes
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.KtNodeTypes.FUN
 import org.jetbrains.kotlin.KtNodeTypes.VALUE_PARAMETER
-import org.jetbrains.kotlin.lexer.KtTokens.VAL_KEYWORD
-import org.jetbrains.kotlin.lexer.KtTokens.VAR_KEYWORD
+import org.jetbrains.kotlin.fir.analysis.diagnostics.*
 import org.jetbrains.kotlin.name.ClassId
-import org.jetbrains.kotlin.psi.KtParameter
 
 object FirAnnotationClassDeclarationChecker : FirBasicDeclarationChecker() {
     override fun check(declaration: FirDeclaration, context: CheckerContext, reporter: DiagnosticReporter) {
@@ -42,24 +33,11 @@ object FirAnnotationClassDeclarationChecker : FirBasicDeclarationChecker() {
             when {
                 it is FirConstructor && it.isPrimary -> {
                     for (parameter in it.valueParameters) {
-                        when (val parameterSourceElement = parameter.source) {
-                            is FirPsiSourceElement<*> -> {
-                                val parameterPsiElement = parameterSourceElement.psi as KtParameter
-
-                                if (!parameterPsiElement.hasValOrVar())
-                                    reporter.report(parameterSourceElement, FirErrors.MISSING_VAL_ON_ANNOTATION_PARAMETER)
-                                else if (parameterPsiElement.isMutable)
-                                    reporter.report(parameterSourceElement, FirErrors.VAR_ANNOTATION_PARAMETER)
-                            }
-                            is FirLightSourceElement -> {
-                                val kidsRef = Ref<Array<LighterASTNode?>>()
-                                parameterSourceElement.treeStructure.getChildren(parameterSourceElement.lighterASTNode, kidsRef)
-
-                                if (kidsRef.get().any { it?.tokenType == VAR_KEYWORD })
-                                    reporter.report(parameterSourceElement, FirErrors.VAR_ANNOTATION_PARAMETER)
-                                else if (kidsRef.get().all { it?.tokenType != VAL_KEYWORD })
-                                    reporter.report(parameterSourceElement, FirErrors.MISSING_VAL_ON_ANNOTATION_PARAMETER)
-                            }
+                        val source = parameter.source ?: continue
+                        if (!source.hasValOrVar()) {
+                            reporter.report(source, FirErrors.MISSING_VAL_ON_ANNOTATION_PARAMETER)
+                        } else if (source.hasVar()) {
+                            reporter.report(source, FirErrors.VAR_ANNOTATION_PARAMETER)
                         }
 
                         val typeRef = parameter.returnTypeRef

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/LightTreePositioningStrategies.kt

@@ -12,6 +12,7 @@ import com.intellij.psi.tree.IElementType
 import com.intellij.psi.tree.TokenSet
 import com.intellij.util.diff.FlyweightCapableTreeStructure
 import org.jetbrains.kotlin.KtNodeTypes
+import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtParameter.VAL_VAR_TOKEN_SET
 
@@ -37,7 +38,7 @@ object LightTreePositioningStrategies {
 
     val VAL_OR_VAR_NODE: LightTreePositioningStrategy = object : LightTreePositioningStrategy() {
         override fun mark(node: LighterASTNode, tree: FlyweightCapableTreeStructure<LighterASTNode>): List<TextRange> {
-            val target = tree.findChildByType(node, VAL_VAR_TOKEN_SET) ?: node
+            val target = tree.valOrVarKeyword(node) ?: node
             return markElement(target, tree)
         }
     }
@@ -143,6 +144,12 @@ object LightTreePositioningStrategies {
     }
 }
 
+fun FirSourceElement.hasValOrVar(): Boolean =
+    treeStructure.valOrVarKeyword(lighterASTNode) != null
+
+fun FirSourceElement.hasVar(): Boolean =
+    treeStructure.findChildByType(lighterASTNode, KtTokens.VAR_KEYWORD) != null
+
 private fun FlyweightCapableTreeStructure<LighterASTNode>.constructorKeyword(node: LighterASTNode): LighterASTNode? =
     findChildByType(node, KtTokens.CONSTRUCTOR_KEYWORD)
 
@@ -158,6 +165,9 @@ private fun FlyweightCapableTreeStructure<LighterASTNode>.rightParenthesis(node:
 private fun FlyweightCapableTreeStructure<LighterASTNode>.objectKeyword(node: LighterASTNode): LighterASTNode? =
     findChildByType(node, KtTokens.OBJECT_KEYWORD)
 
+private fun FlyweightCapableTreeStructure<LighterASTNode>.valOrVarKeyword(node: LighterASTNode): LighterASTNode? =
+    findChildByType(node, VAL_VAR_TOKEN_SET)
+
 private fun FlyweightCapableTreeStructure<LighterASTNode>.accessorNamePlaceholder(node: LighterASTNode): LighterASTNode =
     findChildByType(node, KtTokens.GET_KEYWORD) ?: findChildByType(node, KtTokens.SET_KEYWORD)!!
 
@@ -186,15 +196,15 @@ private fun FlyweightCapableTreeStructure<LighterASTNode>.receiverTypeReference(
 }
 
 private fun FlyweightCapableTreeStructure<LighterASTNode>.findChildByType(node: LighterASTNode, type: IElementType): LighterASTNode? {
-    val childrenRef = Ref<Array<LighterASTNode>>()
+    val childrenRef = Ref<Array<LighterASTNode?>>()
     getChildren(node, childrenRef)
-    return childrenRef.get()?.firstOrNull { it.tokenType == type }
+    return childrenRef.get()?.firstOrNull { it?.tokenType == type }
 }
 
 private fun FlyweightCapableTreeStructure<LighterASTNode>.findChildByType(node: LighterASTNode, type: TokenSet): LighterASTNode? {
-    val childrenRef = Ref<Array<LighterASTNode>>()
+    val childrenRef = Ref<Array<LighterASTNode?>>()
     getChildren(node, childrenRef)
-    return childrenRef.get()?.firstOrNull { it.tokenType in type }
+    return childrenRef.get()?.firstOrNull { it?.tokenType in type }
 }
 
 private fun FlyweightCapableTreeStructure<LighterASTNode>.findParentOfType(node: LighterASTNode, type: IElementType): LighterASTNode? {