[FIR][LightTree] Position strategies fixes

Skip leading comments in DEFAULT strategy
Keep node offset 'padding' when getting child

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

@@ -36,6 +36,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)
-    return visitor.find(lighterASTNode)?.toFirLightSourceElement(treeStructure)
+    val childNode = visitor.find(lighterASTNode) ?: return null
+    //we need to keep 'padding' of parent node in child node
+    val offsetDelta = startOffset - lighterASTNode.startOffset
+    return childNode.toFirLightSourceElement(
+        treeStructure,
+        startOffset = childNode.startOffset + offsetDelta,
+        endOffset = childNode.endOffset + offsetDelta
+    )
 }
 

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

@@ -125,7 +125,12 @@ object LightTreePositioningStrategies {
             startOffset: Int,
             endOffset: Int,
             tree: FlyweightCapableTreeStructure<LighterASTNode>
-        ): List<TextRange> {
+        ): List<TextRange> = markElement(getElementToMark(node, tree), startOffset, endOffset, tree, node)
+
+        override fun isValid(node: LighterASTNode, tree: FlyweightCapableTreeStructure<LighterASTNode>): Boolean =
+            super.isValid(getElementToMark(node, tree), tree)
+
+        private fun getElementToMark(node: LighterASTNode, tree: FlyweightCapableTreeStructure<LighterASTNode>): LighterASTNode {
             val (returnTypeRef, nameIdentifierOrPlaceHolder) = when {
                 node.tokenType == KtNodeTypes.PROPERTY_ACCESSOR ->
                     tree.typeReference(node) to tree.accessorNamePlaceholder(node)
@@ -134,13 +139,7 @@ object LightTreePositioningStrategies {
                 else ->
                     null to null
             }
-            if (returnTypeRef != null) {
+            return returnTypeRef ?: (nameIdentifierOrPlaceHolder ?: node)
-                return markElement(returnTypeRef, startOffset, endOffset, tree, node)
-            }
-            if (nameIdentifierOrPlaceHolder != null) {
-                return markElement(nameIdentifierOrPlaceHolder, startOffset, endOffset, tree, node)
-            }
-            return DEFAULT.mark(node, startOffset, endOffset, tree)
         }
     }
 
@@ -621,7 +620,8 @@ private fun FlyweightCapableTreeStructure<LighterASTNode>.referenceExpression(
     getChildren(node, childrenRef)
     var result = childrenRef.get()?.firstOrNull {
         it?.tokenType == KtNodeTypes.REFERENCE_EXPRESSION || it?.tokenType == KtNodeTypes.CONSTRUCTOR_DELEGATION_REFERENCE ||
-                it?.tokenType == KtNodeTypes.SUPER_EXPRESSION || it?.tokenType == KtNodeTypes.PARENTHESIZED
+                it?.tokenType == KtNodeTypes.SUPER_EXPRESSION || it?.tokenType == KtNodeTypes.PARENTHESIZED ||
+                it?.tokenType == KtNodeTypes.ARRAY_ACCESS_EXPRESSION
     }
     while (locateReferencedName && result != null && result.tokenType == KtNodeTypes.PARENTHESIZED) {
         result = referenceExpression(result, locateReferencedName = true)

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

@@ -36,12 +36,7 @@ fun markElement(
     endOffset: Int,
     tree: FlyweightCapableTreeStructure<LighterASTNode>,
     originalNode: LighterASTNode = node,
-): List<TextRange> {
+): List<TextRange> = markRange(node, node, startOffset, endOffset, tree, originalNode)
-    if (node === originalNode) return listOf(TextRange(startOffset, endOffset))
-    val startDelta = tree.getStartOffset(node) - tree.getStartOffset(originalNode)
-    val endDelta = tree.getEndOffset(node) - tree.getEndOffset(originalNode)
-    return listOf(TextRange(startDelta + startOffset, endDelta + endOffset))
-}
 
 fun markRange(
     from: LighterASTNode,
@@ -62,8 +57,16 @@ private val DOC_AND_COMMENT_TOKENS = setOf(
     KtTokens.EOL_COMMENT, KtTokens.BLOCK_COMMENT, KtTokens.SHEBANG_COMMENT, KtTokens.DOC_COMMENT
 )
 
+private val FILLER_TOKENS = setOf(
+    KtTokens.WHITE_SPACE,
+    KtTokens.EOL_COMMENT,
+    KtTokens.BLOCK_COMMENT,
+    KtTokens.SHEBANG_COMMENT,
+    KtTokens.DOC_COMMENT,
+)
+
 private fun LighterASTNode.nonFillerChildOrSelf(tree: FlyweightCapableTreeStructure<LighterASTNode>): LighterASTNode =
-    getChildren(tree).firstOrNull { it != null && it.tokenType !in DOC_AND_COMMENT_TOKENS } ?: this
+    getChildren(tree).firstOrNull { it != null && it.tokenType !in FILLER_TOKENS } ?: this
 
 
 private fun hasSyntaxErrors(node: LighterASTNode, tree: FlyweightCapableTreeStructure<LighterASTNode>): Boolean {