[FIR] Add positioning strategy for return keyword

compiler/fir/checkers/checkers-component-generator/src/org/jetbrains/kotlin/fir/checkers/generator/diagnostics/DiagnosticData.kt

@@ -54,6 +54,8 @@ enum class PositioningStrategy(private val strategy: String? = null) {
     NAME_OF_NAMED_ARGUMENT,
     VALUE_ARGUMENTS,
     SUPERTYPES_LIST,
+    RETURN_KEYWORD,
+
     ;
 
     val expressionToCreate get() = "SourceElementPositioningStrategies.${strategy ?: name}"

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

@@ -496,6 +496,17 @@ object LightTreePositioningStrategies {
             return markElement(tree.findChildByType(node, KtNodeTypes.INDICES)!!, startOffset, endOffset, tree, node)
         }
     }
+
+    val RETURN_KEYWORD = object : LightTreePositioningStrategy() {
+        override fun mark(
+            node: LighterASTNode,
+            startOffset: Int,
+            endOffset: Int,
+            tree: FlyweightCapableTreeStructure<LighterASTNode>
+        ): List<TextRange> {
+            return markElement(tree.returnKeyword(node) ?: node, startOffset, endOffset, tree)
+        }
+    }
 }
 
 fun FirSourceElement.hasValOrVar(): Boolean =
@@ -525,6 +536,9 @@ private fun FlyweightCapableTreeStructure<LighterASTNode>.whenKeyword(node: Ligh
 private fun FlyweightCapableTreeStructure<LighterASTNode>.ifKeyword(node: LighterASTNode): LighterASTNode? =
     findChildByType(node, KtTokens.IF_KEYWORD)
 
+private fun FlyweightCapableTreeStructure<LighterASTNode>.returnKeyword(node: LighterASTNode): LighterASTNode? =
+    findChildByType(node, KtTokens.RETURN_KEYWORD)
+
 private fun FlyweightCapableTreeStructure<LighterASTNode>.nameIdentifier(node: LighterASTNode): LighterASTNode? =
     findChildByType(node, KtTokens.IDENTIFIER)
 

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

@@ -167,4 +167,9 @@ object SourceElementPositioningStrategies {
         LightTreePositioningStrategies.ARRAY_ACCESS,
         PositioningStrategies.ARRAY_ACCESS
     )
+
+    val RETURN_KEYWORD = SourceElementPositioningStrategy(
+        LightTreePositioningStrategies.RETURN_KEYWORD,
+        PositioningStrategies.RETURN_KEYWORD
+    )
 }

compiler/frontend/src/org/jetbrains/kotlin/diagnostics/PositioningStrategies.kt

@@ -421,6 +421,12 @@ object PositioningStrategies {
         }
     }
 
+    val RETURN_KEYWORD: PositioningStrategy<KtReturnExpression> = object : PositioningStrategy<KtReturnExpression>() {
+        override fun mark(element: KtReturnExpression): List<TextRange> {
+            return markElement(element.returnKeyword)
+        }
+    }
+
     private class ModifierSetBasedPositioningStrategy(private val modifierSet: TokenSet) : PositioningStrategy<KtModifierListOwner>() {
         override fun mark(element: KtModifierListOwner): List<TextRange> {
             val modifierList = element.modifierList

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/diagnostics/KtFirDataClassConverters.kt

@@ -34,6 +34,7 @@ import org.jetbrains.kotlin.psi.KtPrimaryConstructor
 import org.jetbrains.kotlin.psi.KtProperty
 import org.jetbrains.kotlin.psi.KtPropertyAccessor
 import org.jetbrains.kotlin.psi.KtPropertyDelegate
+import org.jetbrains.kotlin.psi.KtReturnExpression
 import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 import org.jetbrains.kotlin.psi.KtTypeAlias
 import org.jetbrains.kotlin.psi.KtTypeParameter