[FIR] LT positioning strategy for UNREACHABLE_CODE

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

@@ -23,10 +23,6 @@ sealed class AbstractFirDiagnosticFactory(
     override fun toString(): String {
         return name
     }
-
-    protected fun checkMyDiagnostic(diagnostic: FirDiagnostic) {
-        require(diagnostic.factory === this) { "$diagnostic should be of this factory" }
-    }
 }
 
 class FirDiagnosticFactory0(
@@ -83,11 +79,6 @@ class FirDiagnosticFactory1<A>(
             else -> incorrectElement(element)
         }
     }
-
-    fun extractArgument(diagnostic: FirDiagnostic): A {
-        checkMyDiagnostic(diagnostic)
-        return (diagnostic as FirDiagnosticWithParameters1<A>).a
-    }
 }
 
 class FirDiagnosticFactory2<A, B>(
@@ -124,12 +115,6 @@ class FirDiagnosticFactory2<A, B>(
             else -> incorrectElement(element)
         }
     }
-
-    fun extractArguments(diagnostic: FirDiagnostic): Pair<A, B> {
-        checkMyDiagnostic(diagnostic)
-        val typed = diagnostic as FirDiagnosticWithParameters2<A, B>
-        return Pair(typed.a, typed.b)
-    }
 }
 
 class FirDiagnosticFactory3<A, B, C>(

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

@@ -833,6 +833,33 @@ object LightTreePositioningStrategies {
         private fun lastSymbol(range: TextRange): TextRange =
             if (range.isEmpty) range else TextRange.create(range.endOffset - 1, range.endOffset)
     }
+
+    val UNREACHABLE_CODE: LightTreePositioningStrategy = object : LightTreePositioningStrategy() {
+        override fun markFirDiagnostic(element: FirSourceElement, diagnostic: FirDiagnostic): List<TextRange> {
+            @Suppress("UNCHECKED_CAST")
+            val typed = diagnostic as FirDiagnosticWithParameters2<Set<FirSourceElement>, Set<FirSourceElement>>
+            with(UnreachableCodeLightTreeHelper(element.treeStructure)) {
+                val reachable = typed.a.map { it.lighterASTNode }.toSet()
+                val unreachable = typed.b.map { it.lighterASTNode }.toSet()
+                if (!element.lighterASTNode.hasChildrenInSet(reachable)) {
+                    return super.markFirDiagnostic(element, diagnostic)
+                }
+
+                val nodesToMark = element.lighterASTNode.getLeavesOrReachableChildren(reachable, unreachable)
+                    .removeReachableElementsWithMeaninglessSiblings(reachable)
+
+                if (nodesToMark.isEmpty()) {
+                    return super.markFirDiagnostic(element, diagnostic)
+                }
+
+                val ranges = nodesToMark.flatMap {
+                    markElement(it, element.startOffset, element.endOffset, element.treeStructure, element.lighterASTNode)
+                }
+
+                return ranges.mergeAdjacentTextRanges()
+            }
+        }
+    }
 }
 
 fun FirSourceElement.hasValOrVar(): Boolean =
@@ -1100,6 +1127,29 @@ fun FlyweightCapableTreeStructure<LighterASTNode>.collectDescendantsOfType(
     return result
 }
 
+fun FlyweightCapableTreeStructure<LighterASTNode>.traverseDescendants(
+    node: LighterASTNode,
+    acceptor: (LighterASTNode) -> Boolean
+) {
+    fun FlyweightCapableTreeStructure<LighterASTNode>.traverse(node: LighterASTNode) {
+        val childrenRef = Ref<Array<LighterASTNode?>>()
+        getChildren(node, childrenRef)
+
+        val childrenRefGet = childrenRef.get()
+        if (childrenRefGet != null) {
+            for (child in childrenRefGet) {
+                if (child != null) {
+                    if (acceptor(child)) {
+                        traverse(child)
+                    }
+                }
+            }
+        }
+    }
+
+    traverse(node)
+}
+
 fun FlyweightCapableTreeStructure<LighterASTNode>.findChildByType(node: LighterASTNode, type: TokenSet): LighterASTNode? {
     val childrenRef = Ref<Array<LighterASTNode?>>()
     getChildren(node, childrenRef)

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

@@ -88,11 +88,12 @@ private val FILLER_TOKENS = setOf(
 )
 
 private fun LighterASTNode.nonFillerFirstChildOrSelf(tree: FlyweightCapableTreeStructure<LighterASTNode>): LighterASTNode =
-    getChildren(tree).firstOrNull { it != null && it.tokenType !in FILLER_TOKENS } ?: this
+    getChildren(tree).firstOrNull { it != null && !it.isFiller() } ?: this
 
 internal fun LighterASTNode.nonFillerLastChildOrSelf(tree: FlyweightCapableTreeStructure<LighterASTNode>): LighterASTNode =
-    getChildren(tree).lastOrNull { it != null && it.tokenType !in FILLER_TOKENS } ?: this
+    getChildren(tree).lastOrNull { it != null && !it.isFiller() } ?: this
 
+internal fun LighterASTNode.isFiller() = tokenType in FILLER_TOKENS
 
 private fun hasSyntaxErrors(node: LighterASTNode, tree: FlyweightCapableTreeStructure<LighterASTNode>): Boolean {
     if (node.tokenType == TokenType.ERROR_ELEMENT) return true

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

@@ -274,7 +274,7 @@ object SourceElementPositioningStrategies {
     )
 
     val UNREACHABLE_CODE = SourceElementPositioningStrategy(
-        LightTreePositioningStrategies.DEFAULT, //todo
+        LightTreePositioningStrategies.UNREACHABLE_CODE,
         FirPsiPositioningStrategies.UNREACHABLE_CODE
     )
 

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

@@ -0,0 +1,86 @@
+/*
+ * Copyright 2010-2021 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
+import org.jetbrains.kotlin.fir.analysis.checkers.getChildren
+import org.jetbrains.kotlin.lexer.KtTokens
+
+typealias Node = LighterASTNode
+
+class UnreachableCodeLightTreeHelper(val tree: FlyweightCapableTreeStructure<Node>) {
+
+    fun Node.hasChildrenInSet(set: Set<Node>): Boolean {
+        var result = false
+        tree.traverseDescendants(this) {
+            if (!result && it != this && it in set) {
+                result = true
+            }
+            !result
+        }
+        return result
+    }
+
+    fun Node.getLeavesOrReachableChildren(reachable: Set<Node>, unreachable: Set<Node>): List<Node> {
+        val result = mutableListOf<Node>()
+        tree.traverseDescendants(this) { element ->
+            val isReachable = element in reachable && !element.hasChildrenInSet(unreachable)
+            if (isReachable || element.getChildren(tree).isEmpty()) {
+                result.add(element)
+                false
+            } else {
+                true
+            }
+        }
+        return result
+    }
+
+    fun List<Node>.removeReachableElementsWithMeaninglessSiblings(reachableElements: Set<Node>): List<Node> {
+        val childrenToRemove = mutableSetOf<Node>()
+        fun collectSiblingsIfMeaningless(elementIndex: Int, direction: Int) {
+            val index = elementIndex + direction
+            if (index !in 0 until size) return
+
+            val element = this[index]
+            if (element.isFiller() || element.tokenType == KtTokens.COMMA) {
+                childrenToRemove.add(element)
+                collectSiblingsIfMeaningless(index, direction)
+            }
+        }
+        for ((index, element) in this.withIndex()) {
+            if (reachableElements.contains(element)) {
+                childrenToRemove.add(element)
+                collectSiblingsIfMeaningless(index, -1)
+                collectSiblingsIfMeaningless(index, 1)
+            }
+        }
+        return filter { it !in childrenToRemove }
+    }
+
+    fun List<TextRange>.mergeAdjacentTextRanges(): List<TextRange> {
+        val result = ArrayList<TextRange>()
+        val lastRange = fold(null as TextRange?) { currentTextRange, elementRange ->
+            when {
+                currentTextRange == null -> {
+                    elementRange
+                }
+                currentTextRange.endOffset == elementRange.startOffset -> {
+                    currentTextRange.union(elementRange)
+                }
+                else -> {
+                    result.add(currentTextRange)
+                    elementRange
+                }
+            }
+        }
+        if (lastRange != null) {
+            result.add(lastRange)
+        }
+        return result
+    }
+}