[FIR][LightTree] Better expression identification

And some refactoring
Regenerate IDE parts

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

@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.fir.analysis.diagnostics
 import com.intellij.lang.LighterASTNode
 import com.intellij.openapi.util.Ref
 import com.intellij.openapi.util.TextRange
-import com.intellij.psi.impl.source.tree.ElementType
 import com.intellij.psi.tree.IElementType
 import com.intellij.psi.tree.TokenSet
 import com.intellij.util.diff.FlyweightCapableTreeStructure
@@ -19,8 +18,8 @@ import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.lexer.KtTokens.MODALITY_MODIFIERS
 import org.jetbrains.kotlin.lexer.KtTokens.VISIBILITY_MODIFIERS
 import org.jetbrains.kotlin.psi.KtParameter.VAL_VAR_TOKEN_SET
-import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.stubs.elements.KtConstantExpressionElementType
+import org.jetbrains.kotlin.psi.stubs.elements.KtStringTemplateExpressionElementType
 import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
 
 object LightTreePositioningStrategies {
@@ -644,20 +643,23 @@ internal fun FlyweightCapableTreeStructure<LighterASTNode>.nameIdentifier(node:
 private fun FlyweightCapableTreeStructure<LighterASTNode>.operationReference(node: LighterASTNode): LighterASTNode? =
     findChildByType(node, KtNodeTypes.OPERATION_REFERENCE)
 
-private val EXPRESSION_NODE_TYPES = setOf(
+private val EXPRESSIONS_SET = listOf(
     KtNodeTypes.REFERENCE_EXPRESSION,
-    KtNodeTypes.CONSTRUCTOR_DELEGATION_REFERENCE,
-    KtNodeTypes.SUPER_EXPRESSION,
-    KtNodeTypes.ARRAY_ACCESS_EXPRESSION,
-    KtNodeTypes.CALL_EXPRESSION,
-    KtNodeTypes.LABELED_EXPRESSION,
     KtNodeTypes.DOT_QUALIFIED_EXPRESSION,
-    KtNodeTypes.FUN,
-    KtNodeTypes.IF,
-    KtNodeTypes.STRING_TEMPLATE,
     KtNodeTypes.LAMBDA_EXPRESSION,
+    KtNodeTypes.FUN
 )
 
+fun LighterASTNode.isExpression(): Boolean {
+    return when (this.tokenType) {
+        is KtNodeType,
+        is KtConstantExpressionElementType,
+        is KtStringTemplateExpressionElementType,
+        in EXPRESSIONS_SET -> true
+        else -> false
+    }
+}
+
 /**
  * @param locateReferencedName whether to remove any nested parentheses while locating the reference element. This is useful for diagnostics
  * on super and unresolved references. For example, with the following, only the part inside the parentheses should be highlighted.
@@ -678,9 +680,7 @@ private fun FlyweightCapableTreeStructure<LighterASTNode>.referenceExpression(
     val childrenRef = Ref<Array<LighterASTNode?>>()
     getChildren(node, childrenRef)
     var result = childrenRef.get()?.firstOrNull {
-        it?.tokenType in EXPRESSION_NODE_TYPES ||
+        it?.isExpression() == true || it?.tokenType == KtNodeTypes.PARENTHESIZED
-                it?.tokenType is KtConstantExpressionElementType ||
-                it?.tokenType == KtNodeTypes.PARENTHESIZED
     }
     while (locateReferencedName && result != null && result.tokenType == KtNodeTypes.PARENTHESIZED) {
         result = referenceExpression(result, locateReferencedName = true)
@@ -689,7 +689,7 @@ private fun FlyweightCapableTreeStructure<LighterASTNode>.referenceExpression(
 }
 
 private fun FlyweightCapableTreeStructure<LighterASTNode>.findReferenceExpressionDeep(node: LighterASTNode): LighterASTNode? =
-    findDescendantByTypes(node, EXPRESSION_NODE_TYPES)
+    findFirstDescendant(node) { it.isExpression() }
 
 private fun FlyweightCapableTreeStructure<LighterASTNode>.rightParenthesis(node: LighterASTNode): LighterASTNode? =
     findChildByType(node, KtTokens.RPAR)
@@ -802,6 +802,16 @@ fun FlyweightCapableTreeStructure<LighterASTNode>.findDescendantByTypes(node: Li
         ?.firstNotNullResult { child -> child?.let { findDescendantByTypes(it, types) } }
 }
 
+fun FlyweightCapableTreeStructure<LighterASTNode>.findFirstDescendant(
+    node: LighterASTNode,
+    predicate: (LighterASTNode) -> Boolean
+): LighterASTNode? {
+    val childrenRef = Ref<Array<LighterASTNode?>>()
+    getChildren(node, childrenRef)
+    return childrenRef.get()?.firstOrNull { it != null && predicate(it) }
+        ?: childrenRef.get()?.firstNotNullResult { child -> child?.let { findFirstDescendant(it, predicate) } }
+}
+
 private fun FlyweightCapableTreeStructure<LighterASTNode>.findChildByType(node: LighterASTNode, type: TokenSet): LighterASTNode? {
     val childrenRef = Ref<Array<LighterASTNode?>>()
     getChildren(node, childrenRef)

compiler/testData/ir/irText/singletons/enumEntry.fir.kt.txt

@@ -16,7 +16,7 @@ open enum class Z : Enum<Z> {
     fun test() {
     }
 
-    inner class A {
+    local inner class A {
       constructor() /* primary */ {
         super/*Any*/()
         /* <init>() */
@@ -36,3 +36,4 @@ open enum class Z : Enum<Z> {
   fun valueOf(value: String): Z /* Synthetic body for ENUM_VALUEOF */
 
 }
+

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/frontend/fir/handlers/FirDiagnosticsHandler.kt

@@ -113,8 +113,8 @@ class FirDiagnosticsHandler(testServices: TestServices) : FirAnalysisHandler(tes
                 FirErrors.SYNTAX.on(FirRealPsiSourceElement(it)).toMetaInfo(testFile, lightTreeEnabled, lightTreeComparingModeEnabled)
             }
         } else {
-            collectLightTreeSyntaxErrors(firFile).map { node ->
+            collectLightTreeSyntaxErrors(firFile).map { sourceElement ->
-                FirErrors.SYNTAX.on(node).toMetaInfo(testFile, lightTreeEnabled, lightTreeComparingModeEnabled)
+                FirErrors.SYNTAX.on(sourceElement).toMetaInfo(testFile, lightTreeEnabled, lightTreeComparingModeEnabled)
             }
         }
 

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/frontend/fir/handlers/lightTreeDiagnostics.kt

@@ -48,12 +48,12 @@ private class LightTreeErrorsCollector(private val tree: Tree) {
 
 }
 
-private data class OffsetTree(val tree: Tree, val offset: Int)
+private data class TreeWithOffset(val tree: Tree, val offset: Int)
 
 private data class VisitorState(
     val lastTree: Tree? = null,
     val visitedTrees: MutableSet<Tree> = mutableSetOf(),
-    val result: MutableList<OffsetTree> = mutableListOf()
+    val result: MutableList<TreeWithOffset> = mutableListOf()
 )
 
 private object FirTreesExtractVisitor : FirVisitor<Unit, VisitorState>() {
@@ -62,7 +62,7 @@ private object FirTreesExtractVisitor : FirVisitor<Unit, VisitorState>() {
         val currentTree = source.treeStructure
         val (lastTree, visitedTrees, result) = data
         val newData = if (lastTree !== currentTree && visitedTrees.add(currentTree)) {
-            val newTreeWithOffset = OffsetTree(currentTree, element.source?.startOffset ?: 0)
+            val newTreeWithOffset = TreeWithOffset(currentTree, element.source?.startOffset ?: 0)
             result.add(newTreeWithOffset)
             data.copy(lastTree = lastTree)
         } else {

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/diagnostics/KtFirDiagnostics.kt

@@ -759,7 +759,7 @@ sealed class KtFirDiagnostic<PSI: PsiElement> : KtDiagnosticWithPsi<PSI> {
         override val diagnosticClass get() = CyclicGenericUpperBound::class
     }
 
-    abstract class DeprecatedTypeParameterSyntax : KtFirDiagnostic<KtTypeParameterList>() {
+    abstract class DeprecatedTypeParameterSyntax : KtFirDiagnostic<KtDeclaration>() {
         override val diagnosticClass get() = DeprecatedTypeParameterSyntax::class
     }
 

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/diagnostics/KtFirDiagnosticsImpl.kt

@@ -1228,7 +1228,7 @@ internal class CyclicGenericUpperBoundImpl(
 internal class DeprecatedTypeParameterSyntaxImpl(
     firDiagnostic: FirPsiDiagnostic<*>,
     override val token: ValidityToken,
-) : KtFirDiagnostic.DeprecatedTypeParameterSyntax(), KtAbstractFirDiagnostic<KtTypeParameterList> {
+) : KtFirDiagnostic.DeprecatedTypeParameterSyntax(), KtAbstractFirDiagnostic<KtDeclaration> {
     override val firDiagnostic: FirPsiDiagnostic<*> by weakRef(firDiagnostic)
 }