From d942780c14faa3ad5f85c8f9dc84179849954bac Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Tue, 17 Nov 2020 10:40:13 +0300
Subject: [PATCH] [FIR] Introduce LightTreePositioningStrategy

---
 .../diagnostics/FirDiagnosticFactory.kt       | 33 ++++++++++++++-----
 .../LightTreePositioningStrategy.kt           | 29 ++++++++++++++++
 2 files changed, 54 insertions(+), 8 deletions(-)
 create mode 100644 compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/LightTreePositioningStrategy.kt

diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnosticFactory.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnosticFactory.kt
index 9cfdda9721c..22eeb664358 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnosticFactory.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/FirDiagnosticFactory.kt
@@ -7,6 +7,7 @@
 
 package org.jetbrains.kotlin.fir.analysis.diagnostics
 
+import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.diagnostics.*
 import org.jetbrains.kotlin.fir.FirLightSourceElement
@@ -16,19 +17,26 @@ import org.jetbrains.kotlin.fir.FirSourceElement
 sealed class AbstractFirDiagnosticFactory<out E : FirSourceElement, D : FirDiagnostic<E>>(
     val name: String,
     val severity: Severity,
+    val positioningStrategy: LightTreePositioningStrategy,
 ) {
     abstract val psiDiagnosticFactory: DiagnosticFactoryWithPsiElement<*, *>
 
     abstract val defaultRenderer: FirDiagnosticRenderer<*>
 
+    fun getTextRanges(diagnostic: FirDiagnostic<*>): List<TextRange> =
+        positioningStrategy.markDiagnostic(diagnostic)
+
     override fun toString(): String {
         return name
     }
 }
 
 class FirDiagnosticFactory0<E : FirSourceElement, P : PsiElement>(
-    name: String, severity: Severity, override val psiDiagnosticFactory: DiagnosticFactory0<P>
-) : AbstractFirDiagnosticFactory<E, FirSimpleDiagnostic<E>>(name, severity) {
+    name: String,
+    severity: Severity,
+    override val psiDiagnosticFactory: DiagnosticFactory0<P>,
+    positioningStrategy: LightTreePositioningStrategy = LightTreePositioningStrategy.DEFAULT,
+) : AbstractFirDiagnosticFactory<E, FirSimpleDiagnostic<E>>(name, severity, positioningStrategy) {
     companion object {
         private val DefaultRenderer = SimpleFirDiagnosticRenderer("")
     }
@@ -48,8 +56,11 @@ class FirDiagnosticFactory0<E : FirSourceElement, P : PsiElement>(
 }
 
 class FirDiagnosticFactory1<E : FirSourceElement, P : PsiElement, A : Any>(
-    name: String, severity: Severity, override val psiDiagnosticFactory: DiagnosticFactory1<P, A>
-) : AbstractFirDiagnosticFactory<E, FirDiagnosticWithParameters1<E, A>>(name, severity) {
+    name: String,
+    severity: Severity,
+    override val psiDiagnosticFactory: DiagnosticFactory1<P, A>,
+    positioningStrategy: LightTreePositioningStrategy = LightTreePositioningStrategy.DEFAULT,
+) : AbstractFirDiagnosticFactory<E, FirDiagnosticWithParameters1<E, A>>(name, severity, positioningStrategy) {
     companion object {
         private val DefaultRenderer = FirDiagnosticWithParameters1Renderer(
             "{0}",
@@ -72,8 +83,11 @@ class FirDiagnosticFactory1<E : FirSourceElement, P : PsiElement, A : Any>(
 }
 
 class FirDiagnosticFactory2<E : FirSourceElement, P : PsiElement, A : Any, B : Any>(
-    name: String, severity: Severity, override val psiDiagnosticFactory: DiagnosticFactory2<P, A, B>
-) : AbstractFirDiagnosticFactory<E, FirDiagnosticWithParameters2<E, A, B>>(name, severity) {
+    name: String,
+    severity: Severity,
+    override val psiDiagnosticFactory: DiagnosticFactory2<P, A, B>,
+    positioningStrategy: LightTreePositioningStrategy = LightTreePositioningStrategy.DEFAULT,
+) : AbstractFirDiagnosticFactory<E, FirDiagnosticWithParameters2<E, A, B>>(name, severity, positioningStrategy) {
     companion object {
         private val DefaultRenderer = FirDiagnosticWithParameters2Renderer(
             "{0}, {1}",
@@ -97,8 +111,11 @@ class FirDiagnosticFactory2<E : FirSourceElement, P : PsiElement, A : Any, B : A
 }
 
 class FirDiagnosticFactory3<E : FirSourceElement, P : PsiElement, A : Any, B : Any, C : Any>(
-    name: String, severity: Severity, override val psiDiagnosticFactory: DiagnosticFactory3<P, A, B, C>
-) : AbstractFirDiagnosticFactory<E, FirDiagnosticWithParameters3<E, A, B, C>>(name, severity) {
+    name: String,
+    severity: Severity,
+    override val psiDiagnosticFactory: DiagnosticFactory3<P, A, B, C>,
+    positioningStrategy: LightTreePositioningStrategy = LightTreePositioningStrategy.DEFAULT,
+) : AbstractFirDiagnosticFactory<E, FirDiagnosticWithParameters3<E, A, B, C>>(name, severity, positioningStrategy) {
     companion object {
         private val DefaultRenderer = FirDiagnosticWithParameters3Renderer(
             "{0}, {1}, {2}",
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
new file mode 100644
index 00000000000..9c3206c5c24
--- /dev/null
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/diagnostics/LightTreePositioningStrategy.kt
@@ -0,0 +1,29 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.analysis.diagnostics
+
+import com.intellij.lang.LighterASTNode
+import com.intellij.openapi.util.TextRange
+import com.intellij.util.diff.FlyweightCapableTreeStructure
+
+open class LightTreePositioningStrategy {
+    open fun markDiagnostic(diagnostic: FirDiagnostic<*>): List<TextRange> {
+        val element = diagnostic.element
+        return mark(element.lighterASTNode, element.treeStructure)
+    }
+
+    open fun mark(node: LighterASTNode, tree: FlyweightCapableTreeStructure<LighterASTNode>): List<TextRange> {
+        return markElement(node, tree)
+    }
+
+    companion object {
+        val DEFAULT = LightTreePositioningStrategy()
+    }
+}
+
+fun markElement(node: LighterASTNode, tree: FlyweightCapableTreeStructure<LighterASTNode>): List<TextRange> {
+    return listOf(TextRange(tree.getStartOffset(node), tree.getEndOffset(node)))
+}