[FIR][LightTree] More fixes in diagnostic positioning

Keep tree padding for modifiers
More REFERENCE_EXPRESSIONS kinds
Fix source for some call expressions

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.analysis
 
+import com.intellij.lang.LighterASTNode
 import com.intellij.psi.tree.IElementType
 import com.intellij.psi.tree.TokenSet
 import org.jetbrains.kotlin.fir.*
@@ -37,7 +38,13 @@ private fun FirPsiSourceElement<*>.getChild(types: Set<IElementType>, index: Int
 private fun FirLightSourceElement.getChild(types: Set<IElementType>, index: Int, depth: Int): FirSourceElement? {
     val visitor = LighterTreeElementFinderByType(treeStructure, types, index, depth)
     val childNode = visitor.find(lighterASTNode) ?: return null
-    //we need to keep 'padding' of parent node in child node
+    return buildChildSourceElement(childNode)
+}
+
+/**
+ * Keeps 'padding' of parent node in child node
+ */
+internal fun FirLightSourceElement.buildChildSourceElement(childNode: LighterASTNode): FirLightSourceElement {
     val offsetDelta = startOffset - lighterASTNode.startOffset
     return childNode.toFirLightSourceElement(
         treeStructure,

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

@@ -34,14 +34,15 @@ internal sealed class FirModifierList {
 
     class FirLightModifierList(
         val modifierList: LighterASTNode,
-        val tree: FlyweightCapableTreeStructure<LighterASTNode>
+        val tree: FlyweightCapableTreeStructure<LighterASTNode>,
+        private val offsetDelta: Int
     ) : FirModifierList() {
         override val modifiers: List<FirModifier.FirLightModifier>
             get() {
                 val modifierNodes = modifierList.getChildren(tree)
                 return modifierNodes.filterNotNull()
                     .filter { it.tokenType is KtModifierKeywordToken }
-                    .map { FirModifier.FirLightModifier(it, it.tokenType as KtModifierKeywordToken, tree) }
+                    .map { FirModifier.FirLightModifier(it, it.tokenType as KtModifierKeywordToken, tree, offsetDelta) }
             }
     }
 
@@ -65,10 +66,15 @@ internal sealed class FirModifier<Node : Any>(val node: Node, val token: KtModif
     class FirLightModifier(
         node: LighterASTNode,
         token: KtModifierKeywordToken,
-        val tree: FlyweightCapableTreeStructure<LighterASTNode>
+        val tree: FlyweightCapableTreeStructure<LighterASTNode>,
+        private val offsetDelta: Int
     ) : FirModifier<LighterASTNode>(node, token) {
         override val source: FirSourceElement
-            get() = node.toFirLightSourceElement(tree)
+            get() = node.toFirLightSourceElement(
+                tree,
+                startOffset = node.startOffset + offsetDelta,
+                endOffset = node.endOffset + offsetDelta
+            )
     }
 
     abstract val source: FirSourceElement
@@ -81,7 +87,8 @@ internal fun FirSourceElement?.getModifierList(): FirModifierList? {
         is FirLightSourceElement -> {
             val modifierListNode = lighterASTNode.getChildren(treeStructure).find { it?.tokenType == KtNodeTypes.MODIFIER_LIST }
                 ?: return null
-            FirModifierList.FirLightModifierList(modifierListNode, treeStructure)
+            val offsetDelta = startOffset - lighterASTNode.startOffset
+            FirModifierList.FirLightModifierList(modifierListNode, treeStructure, offsetDelta)
         }
     }
 }

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

@@ -20,6 +20,7 @@ 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.stubs.elements.KtConstantExpressionElementType
 import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
 
 object LightTreePositioningStrategies {
@@ -625,6 +626,9 @@ private val REFERENCE_EXPRESSIONS = setOf(
     KtNodeTypes.CONSTRUCTOR_DELEGATION_REFERENCE,
     KtNodeTypes.SUPER_EXPRESSION,
     KtNodeTypes.ARRAY_ACCESS_EXPRESSION,
+    KtNodeTypes.CALL_EXPRESSION,
+    KtNodeTypes.LABELED_EXPRESSION,
+    KtNodeTypes.DOT_QUALIFIED_EXPRESSION,
     KtNodeTypes.FUN,
 )
 
@@ -648,7 +652,9 @@ private fun FlyweightCapableTreeStructure<LighterASTNode>.referenceExpression(
     val childrenRef = Ref<Array<LighterASTNode?>>()
     getChildren(node, childrenRef)
     var result = childrenRef.get()?.firstOrNull {
-        it?.tokenType in REFERENCE_EXPRESSIONS || it?.tokenType == KtNodeTypes.PARENTHESIZED
+        it?.tokenType in REFERENCE_EXPRESSIONS ||
+                it?.tokenType is KtConstantExpressionElementType ||
+                it?.tokenType == KtNodeTypes.PARENTHESIZED
     }
     while (locateReferencedName && result != null && result.tokenType == KtNodeTypes.PARENTHESIZED) {
         result = referenceExpression(result, locateReferencedName = true)

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

@@ -19,7 +19,6 @@ import org.jetbrains.kotlin.fir.types.builder.buildImplicitTypeRef
 import org.jetbrains.kotlin.lexer.KtToken
 import org.jetbrains.kotlin.lexer.KtTokens.*
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.stubs.elements.KtFunctionElementType
 import kotlin.contracts.ExperimentalContracts
 
 abstract class BaseConverter(
@@ -81,6 +80,15 @@ abstract class BaseConverter(
         return null
     }
 
+    protected fun LighterASTNode.getFirstChildExpressionUnwrapped(): LighterASTNode? {
+        val expression = getFirstChildExpression() ?: return null
+        return if (expression.tokenType == KtNodeTypes.PARENTHESIZED) {
+            expression.getFirstChildExpressionUnwrapped()
+        } else {
+            expression
+        }
+    }
+
     private fun LighterASTNode.getLastChildExpression(): LighterASTNode? {
         var result: LighterASTNode? = null
         forEachChildren {

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

@@ -548,7 +548,7 @@ class ExpressionsConverter(
         val (calleeReference, explicitReceiver, isImplicitInvoke) = when {
             name != null -> CalleeAndReceiver(
                 buildSimpleNamedReference {
-                    this.source = source
+                    this.source = callSuffix.getFirstChildExpressionUnwrapped()?.toFirSourceElement() ?: source
                     this.name = name.nameAsSafeName()
                 }
             )

compiler/testData/diagnostics/tests/incompleteCode/diagnosticWithSyntaxError/checkBackingFieldException.kt

@@ -4,7 +4,8 @@ package h
 class Square() {
   var size : Double =
   <!UNRESOLVED_REFERENCE!>set<!>(<!UNRESOLVED_REFERENCE!>value<!>) {
-    <!SYNTAX!>$area<!> <!SYNTAX!>= size * size<!>
+    //in LT this LAMBDA_EXPRESSION get parsed lazyly, but doesn't got anywhere in FIR tree (as property doesn't have place for it)
+    <!SYNTAX{PSI}!>$area<!> <!SYNTAX{PSI}!>= size * size<!>
   }
 
   <!MUST_BE_INITIALIZED_OR_BE_ABSTRACT!>var area : Double<!>

compiler/testData/diagnostics/tests/inner/innerClassInEnumEntryClass_lv11.fir.kt

@@ -6,9 +6,9 @@ enum class Enum {
 
         <!NESTED_CLASS_NOT_ALLOWED!>class TestNested<!>
 
-        <!NESTED_CLASS_NOT_ALLOWED!>interface TestInterface<!>
+        <!LOCAL_INTERFACE_NOT_ALLOWED, NESTED_CLASS_NOT_ALLOWED!>interface TestInterface<!>
 
-        object TestObject
+        <!LOCAL_OBJECT_NOT_ALLOWED!>object TestObject<!>
 
         enum class TestEnumClass {
             OTHER_ENTRY

compiler/testData/diagnostics/tests/inner/innerClassInEnumEntryClass_lv12.fir.kt

@@ -6,9 +6,9 @@ enum class Enum {
 
         <!NESTED_CLASS_NOT_ALLOWED!>class TestNested<!>
 
-        <!NESTED_CLASS_NOT_ALLOWED!>interface TestInterface<!>
+        <!LOCAL_INTERFACE_NOT_ALLOWED, NESTED_CLASS_NOT_ALLOWED!>interface TestInterface<!>
 
-        object TestObject
+        <!LOCAL_OBJECT_NOT_ALLOWED!>object TestObject<!>
 
         enum class TestEnumClass {
             OTHER_ENTRY

compiler/testData/diagnostics/tests/inner/innerClassInEnumEntryClass_lv13.fir.kt

@@ -6,9 +6,9 @@ enum class Enum {
 
         <!NESTED_CLASS_NOT_ALLOWED!>class TestNested<!>
 
-        <!NESTED_CLASS_NOT_ALLOWED!>interface TestInterface<!>
+        <!LOCAL_INTERFACE_NOT_ALLOWED, NESTED_CLASS_NOT_ALLOWED!>interface TestInterface<!>
 
-        object TestObject
+        <!LOCAL_OBJECT_NOT_ALLOWED!>object TestObject<!>
 
         enum class TestEnumClass {
             OTHER_ENTRY

compiler/testData/diagnostics/tests/objects/Objects.fir.kt

@@ -6,13 +6,13 @@ package toplevelObjectDeclarations
 
   <!NO_VALUE_FOR_PARAMETER!>class T : <!SUPERTYPE_NOT_INITIALIZED!>Foo<!> {}<!>
 
-  object A : <!SUPERTYPE_NOT_INITIALIZED!>Foo<!> {
+  <!NO_VALUE_FOR_PARAMETER{LT}!>object A : <!SUPERTYPE_NOT_INITIALIZED!>Foo<!> {
     val x : Int = 2
 
     fun test() : Int {
-      return x + foo(<!NO_VALUE_FOR_PARAMETER!>)<!>
+      return x + foo(<!NO_VALUE_FOR_PARAMETER{PSI}!>)<!>
     }
-  }
+  }<!>
 
   object B : A {}
 

compiler/testData/diagnostics/tests/regressions/kt328.fir.kt

@@ -25,5 +25,5 @@ val z = z
 fun block(f : () -> Unit) = f()
 
 fun bar3() = block{ <!UNRESOLVED_REFERENCE!>foo3<!>() // <-- missing closing curly bracket
-fun foo3() = block{ bar3() }<!SYNTAX!><!>
+fun foo3() = block{ bar3() }<!SYNTAX{PSI}!><!>