From 42c59f7383ce3eaee9d40bc82931fb4fc869910d Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Tue, 17 Nov 2020 15:15:11 +0300
Subject: [PATCH] [FIR] Enhance light tree DEFAULT strategy for objects to
 cover header only

---
 .../LightTreePositioningStrategies.kt         | 19 +++++++++++++++++++
 .../LightTreePositioningStrategy.kt           | 11 +----------
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/LightTreePositioningStrategies.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/LightTreePositioningStrategies.kt
index 63353c5d268..c2c9a105958 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/LightTreePositioningStrategies.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/LightTreePositioningStrategies.kt
@@ -16,6 +16,25 @@ import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtParameter.VAL_VAR_TOKEN_SET
 
 object LightTreePositioningStrategies {
+    internal val DEFAULT = object : LightTreePositioningStrategy() {
+        override fun mark(node: LighterASTNode, tree: FlyweightCapableTreeStructure<LighterASTNode>): List<TextRange> {
+            when (node.tokenType) {
+                KtNodeTypes.OBJECT_DECLARATION -> {
+                    val objectKeyword = tree.findChildByType(node, KtTokens.OBJECT_KEYWORD)!!
+                    return markRange(
+                        from = objectKeyword,
+                        to = tree.findChildByType(node, KtTokens.IDENTIFIER) ?: objectKeyword,
+                        tree
+                    )
+                }
+                KtNodeTypes.CONSTRUCTOR_DELEGATION_CALL -> {
+                    return SECONDARY_CONSTRUCTOR_DELEGATION_CALL.mark(node, tree)
+                }
+            }
+            return super.mark(node, tree)
+        }
+    }
+
     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
diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/LightTreePositioningStrategy.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/LightTreePositioningStrategy.kt
index 1b11a5051cd..6e6bdf7bbe2 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/LightTreePositioningStrategy.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/LightTreePositioningStrategy.kt
@@ -21,16 +21,7 @@ open class LightTreePositioningStrategy {
     }
 
     companion object {
-        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)
-            }
-        }
+        val DEFAULT = LightTreePositioningStrategies.DEFAULT
     }
 }