[FIR] Implement light tree CONSTRUCTOR_DELEGATION_CALL strategy

Default strategy sometimes delegates to the CONSTRUCTOR_DELEGATION_CALL
strategy, so we add it in this commit to process some related cases
properly.

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

@@ -71,11 +71,11 @@ object FirErrors {
     val NON_PRIVATE_CONSTRUCTOR_IN_ENUM by existing0<FirSourceElement, PsiElement>()
     val NON_PRIVATE_CONSTRUCTOR_IN_SEALED by existing0<FirSourceElement, PsiElement>()
     val CYCLIC_CONSTRUCTOR_DELEGATION_CALL by warning0<FirSourceElement, PsiElement>()
-    val PRIMARY_CONSTRUCTOR_DELEGATION_CALL_EXPECTED by warning0<FirSourceElement, PsiElement>()
+    val PRIMARY_CONSTRUCTOR_DELEGATION_CALL_EXPECTED by warning0<FirSourceElement, PsiElement>(LightTreePositioningStrategies.SECONDARY_CONSTRUCTOR_DELEGATION_CALL)
     val SUPERTYPE_INITIALIZED_WITHOUT_PRIMARY_CONSTRUCTOR by warning0<FirSourceElement, PsiElement>()
     val DELEGATION_SUPER_CALL_IN_ENUM_CONSTRUCTOR by warning0<FirSourceElement, PsiElement>()
     val PRIMARY_CONSTRUCTOR_REQUIRED_FOR_DATA_CLASS by warning0<FirSourceElement, PsiElement>()
-    val EXPLICIT_DELEGATION_CALL_REQUIRED by warning0<FirSourceElement, PsiElement>()
+    val EXPLICIT_DELEGATION_CALL_REQUIRED by warning0<FirSourceElement, PsiElement>(LightTreePositioningStrategies.SECONDARY_CONSTRUCTOR_DELEGATION_CALL)
     val SEALED_CLASS_CONSTRUCTOR_CALL by error0<FirSourceElement, PsiElement>()
 
     // Annotations

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

@@ -11,6 +11,8 @@ import com.intellij.openapi.util.TextRange
 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.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtParameter.VAL_VAR_TOKEN_SET
 
 object LightTreePositioningStrategies {
@@ -20,16 +22,65 @@ object LightTreePositioningStrategies {
             return markElement(target, tree)
         }
     }
+
+    val SECONDARY_CONSTRUCTOR_DELEGATION_CALL: LightTreePositioningStrategy = object : LightTreePositioningStrategy() {
+        override fun mark(node: LighterASTNode, tree: FlyweightCapableTreeStructure<LighterASTNode>): List<TextRange> {
+            when (node.tokenType) {
+                KtNodeTypes.SECONDARY_CONSTRUCTOR -> {
+                    val valueParameterList = tree.findChildByType(node, KtNodeTypes.VALUE_PARAMETER_LIST) ?: return markElement(node, tree)
+                    return markRange(
+                        tree.findChildByType(node, KtTokens.CONSTRUCTOR_KEYWORD)!!,
+                        tree.lastChild(valueParameterList)!!, tree
+                    )
+                }
+                KtNodeTypes.CONSTRUCTOR_DELEGATION_CALL -> {
+                    val delegationReference = tree.findChildByType(node, KtNodeTypes.CONSTRUCTOR_DELEGATION_REFERENCE)
+                    if (delegationReference != null && tree.firstChild(delegationReference) == null) {
+                        val constructor = tree.findParentOfType(node, KtNodeTypes.SECONDARY_CONSTRUCTOR)!!
+                        val valueParameterList = tree.findChildByType(constructor, KtNodeTypes.VALUE_PARAMETER_LIST)
+                            ?: return markElement(constructor, tree)
+                        return markRange(
+                            tree.findChildByType(constructor, KtTokens.CONSTRUCTOR_KEYWORD)!!,
+                            tree.lastChild(valueParameterList)!!, tree
+                        )
+                    }
+                    return markElement(delegationReference ?: node, tree)
+                }
+                else -> error("unexpected element $node")
+            }
+        }
+    }
 }
 
-fun FlyweightCapableTreeStructure<LighterASTNode>.findChildByType(node: LighterASTNode, type: IElementType): LighterASTNode? {
+private fun FlyweightCapableTreeStructure<LighterASTNode>.findChildByType(node: LighterASTNode, type: IElementType): LighterASTNode? {
     val childrenRef = Ref<Array<LighterASTNode>>()
     getChildren(node, childrenRef)
     return childrenRef.get()?.firstOrNull { it.tokenType == type }
 }
 
-fun FlyweightCapableTreeStructure<LighterASTNode>.findChildByType(node: LighterASTNode, type: TokenSet): LighterASTNode? {
+private fun FlyweightCapableTreeStructure<LighterASTNode>.findChildByType(node: LighterASTNode, type: TokenSet): LighterASTNode? {
     val childrenRef = Ref<Array<LighterASTNode>>()
     getChildren(node, childrenRef)
     return childrenRef.get()?.firstOrNull { it.tokenType in type }
 }
+
+private fun FlyweightCapableTreeStructure<LighterASTNode>.findParentOfType(node: LighterASTNode, type: IElementType): LighterASTNode? {
+    var parent = getParent(node)
+    while (parent != null) {
+        if (parent.tokenType == type) return parent
+        parent = getParent(parent)
+    }
+    return null
+}
+
+private fun FlyweightCapableTreeStructure<LighterASTNode>.firstChild(node: LighterASTNode): LighterASTNode? {
+    val childrenRef = Ref<Array<LighterASTNode>>()
+    getChildren(node, childrenRef)
+    return childrenRef.get()?.firstOrNull()
+}
+
+private fun FlyweightCapableTreeStructure<LighterASTNode>.lastChild(node: LighterASTNode): LighterASTNode? {
+    val childrenRef = Ref<Array<LighterASTNode>>()
+    getChildren(node, childrenRef)
+    return childrenRef.get()?.lastOrNull()
+}

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.analysis.diagnostics
 import com.intellij.lang.LighterASTNode
 import com.intellij.openapi.util.TextRange
 import com.intellij.util.diff.FlyweightCapableTreeStructure
+import org.jetbrains.kotlin.KtNodeTypes
 
 open class LightTreePositioningStrategy {
     open fun markDiagnostic(diagnostic: FirDiagnostic<*>): List<TextRange> {
@@ -20,10 +21,23 @@ open class LightTreePositioningStrategy {
     }
 
     companion object {
-        val DEFAULT = LightTreePositioningStrategy()
+        val DEFAULT = object : LightTreePositioningStrategy() {
+            override fun mark(node: LighterASTNode, tree: FlyweightCapableTreeStructure<LighterASTNode>): List<TextRange> {
+                when (node.tokenType) {
+                    KtNodeTypes.CONSTRUCTOR_DELEGATION_CALL -> {
+                        return LightTreePositioningStrategies.SECONDARY_CONSTRUCTOR_DELEGATION_CALL.mark(node, tree)
+                    }
+                }
+                return super.mark(node, tree)
+            }
+        }
     }
 }
 
 fun markElement(node: LighterASTNode, tree: FlyweightCapableTreeStructure<LighterASTNode>): List<TextRange> {
     return listOf(TextRange(tree.getStartOffset(node), tree.getEndOffset(node)))
 }
+
+fun markRange(from: LighterASTNode, to: LighterASTNode, tree: FlyweightCapableTreeStructure<LighterASTNode>): List<TextRange> {
+    return listOf(TextRange(tree.getStartOffset(from), tree.getEndOffset(to)))
+}

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

@@ -8,8 +8,10 @@ package org.jetbrains.kotlin.fir
 import com.intellij.lang.LighterASTNode
 import com.intellij.lang.TreeBackedLighterAST
 import com.intellij.openapi.util.Ref
+import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
+import com.intellij.psi.PsiWhiteSpace
 import com.intellij.psi.tree.IElementType
 import com.intellij.util.diff.FlyweightCapableTreeStructure
 
@@ -206,9 +208,39 @@ sealed class FirPsiSourceElement<out P : PsiElement>(val psi: P) : FirSourceElem
         override fun disposeChildren(p0: Array<out LighterASTNode>?, p1: Int) {
         }
 
-        override fun getStartOffset(node: LighterASTNode): Int = node.unwrap().startOffset
+        override fun getStartOffset(node: LighterASTNode): Int {
+            return getStartOffset(node.unwrap().psi)
+        }
 
-        override fun getEndOffset(node: LighterASTNode): Int = node.unwrap().let { it.startOffset + it.textLength }
+        private fun getStartOffset(element: PsiElement): Int {
+            var child = element.firstChild
+            if (child != null) {
+                while (child is PsiComment || child is PsiWhiteSpace) {
+                    child = child.nextSibling
+                }
+                if (child != null) {
+                    return getStartOffset(child)
+                }
+            }
+            return element.textRange.startOffset
+        }
+
+        override fun getEndOffset(node: LighterASTNode): Int {
+            return getEndOffset(node.unwrap().psi)
+        }
+
+        private fun getEndOffset(element: PsiElement): Int {
+            var child = element.lastChild
+            if (child != null) {
+                while (child is PsiComment || child is PsiWhiteSpace) {
+                    child = child.prevSibling
+                }
+                if (child != null) {
+                    return getEndOffset(child)
+                }
+            }
+            return element.textRange.endOffset
+        }
     }
 }