[FIR] Implement light tree VAL_VAR strategy as an example

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

@@ -11,36 +11,52 @@ import org.jetbrains.kotlin.fir.FirSourceElement
 import kotlin.properties.ReadOnlyProperty
 import kotlin.reflect.KProperty
 
-fun <E : FirSourceElement, P : PsiElement> warning0(): DiagnosticFactory0DelegateProvider<E, P> {
+fun <E : FirSourceElement, P : PsiElement> warning0(
-    return DiagnosticFactory0DelegateProvider(Severity.WARNING)
+    positioningStrategy: LightTreePositioningStrategy = LightTreePositioningStrategy.DEFAULT
+): DiagnosticFactory0DelegateProvider<E, P> {
+    return DiagnosticFactory0DelegateProvider(Severity.WARNING, positioningStrategy)
 }
 
-fun <E : FirSourceElement, P : PsiElement, A : Any> warning1(): DiagnosticFactory1DelegateProvider<E, P, A> {
+fun <E : FirSourceElement, P : PsiElement, A : Any> warning1(
-    return DiagnosticFactory1DelegateProvider(Severity.WARNING)
+    positioningStrategy: LightTreePositioningStrategy = LightTreePositioningStrategy.DEFAULT
+): DiagnosticFactory1DelegateProvider<E, P, A> {
+    return DiagnosticFactory1DelegateProvider(Severity.WARNING, positioningStrategy)
 }
 
-fun <E : FirSourceElement, P : PsiElement, A : Any, B : Any> warning2(): DiagnosticFactory2DelegateProvider<E, P, A, B> {
+fun <E : FirSourceElement, P : PsiElement, A : Any, B : Any> warning2(
-    return DiagnosticFactory2DelegateProvider(Severity.WARNING)
+    positioningStrategy: LightTreePositioningStrategy = LightTreePositioningStrategy.DEFAULT
+): DiagnosticFactory2DelegateProvider<E, P, A, B> {
+    return DiagnosticFactory2DelegateProvider(Severity.WARNING, positioningStrategy)
 }
 
-fun <E : FirSourceElement, P : PsiElement, A : Any, B : Any, C : Any> warning3(): DiagnosticFactory3DelegateProvider<E, P, A, B, C> {
+fun <E : FirSourceElement, P : PsiElement, A : Any, B : Any, C : Any> warning3(
-    return DiagnosticFactory3DelegateProvider(Severity.WARNING)
+    positioningStrategy: LightTreePositioningStrategy = LightTreePositioningStrategy.DEFAULT
+): DiagnosticFactory3DelegateProvider<E, P, A, B, C> {
+    return DiagnosticFactory3DelegateProvider(Severity.WARNING, positioningStrategy)
 }
 
-fun <E : FirSourceElement, P : PsiElement> error0(): DiagnosticFactory0DelegateProvider<E, P> {
+fun <E : FirSourceElement, P : PsiElement> error0(
-    return DiagnosticFactory0DelegateProvider(Severity.ERROR)
+    positioningStrategy: LightTreePositioningStrategy = LightTreePositioningStrategy.DEFAULT
+): DiagnosticFactory0DelegateProvider<E, P> {
+    return DiagnosticFactory0DelegateProvider(Severity.ERROR, positioningStrategy)
 }
 
-fun <E : FirSourceElement, P : PsiElement, A : Any> error1(): DiagnosticFactory1DelegateProvider<E, P, A> {
+fun <E : FirSourceElement, P : PsiElement, A : Any> error1(
-    return DiagnosticFactory1DelegateProvider(Severity.ERROR)
+    positioningStrategy: LightTreePositioningStrategy = LightTreePositioningStrategy.DEFAULT
+): DiagnosticFactory1DelegateProvider<E, P, A> {
+    return DiagnosticFactory1DelegateProvider(Severity.ERROR, positioningStrategy)
 }
 
-fun <E : FirSourceElement, P : PsiElement, A : Any, B : Any> error2(): DiagnosticFactory2DelegateProvider<E, P, A, B> {
+fun <E : FirSourceElement, P : PsiElement, A : Any, B : Any> error2(
-    return DiagnosticFactory2DelegateProvider(Severity.ERROR)
+    positioningStrategy: LightTreePositioningStrategy = LightTreePositioningStrategy.DEFAULT
+): DiagnosticFactory2DelegateProvider<E, P, A, B> {
+    return DiagnosticFactory2DelegateProvider(Severity.ERROR, positioningStrategy)
 }
 
-fun <E : FirSourceElement, P : PsiElement, A : Any, B : Any, C : Any> error3(): DiagnosticFactory3DelegateProvider<E, P, A, B, C> {
+fun <E : FirSourceElement, P : PsiElement, A : Any, B : Any, C : Any> error3(
-    return DiagnosticFactory3DelegateProvider(Severity.ERROR)
+    positioningStrategy: LightTreePositioningStrategy = LightTreePositioningStrategy.DEFAULT
+): DiagnosticFactory3DelegateProvider<E, P, A, B, C> {
+    return DiagnosticFactory3DelegateProvider(Severity.ERROR, positioningStrategy)
 }
 
 /**
@@ -48,52 +64,56 @@ fun <E : FirSourceElement, P : PsiElement, A : Any, B : Any, C : Any> error3():
  *   that takes `PsiElement` as first type parameter
  */
 fun <E : FirSourceElement, P : PsiElement> existing0(): DiagnosticFactory0DelegateProvider<E, P> {
-    return DiagnosticFactory0DelegateProvider(Severity.ERROR)
+    return DiagnosticFactory0DelegateProvider(Severity.ERROR, LightTreePositioningStrategy.DEFAULT)
 }
 
 fun <E : FirSourceElement, P : PsiElement, A : Any> existing1(): DiagnosticFactory1DelegateProvider<E, P, A> {
-    return DiagnosticFactory1DelegateProvider(Severity.ERROR)
+    return DiagnosticFactory1DelegateProvider(Severity.ERROR, LightTreePositioningStrategy.DEFAULT)
 }
 
 fun <E : FirSourceElement, P : PsiElement, A : Any, B : Any> existing2(): DiagnosticFactory2DelegateProvider<E, P, A, B> {
-    return DiagnosticFactory2DelegateProvider(Severity.ERROR)
+    return DiagnosticFactory2DelegateProvider(Severity.ERROR, LightTreePositioningStrategy.DEFAULT)
 }
 
 fun <E : FirSourceElement, P : PsiElement, A : Any, B : Any, C : Any> existing3(): DiagnosticFactory3DelegateProvider<E, P, A, B, C> {
-    return DiagnosticFactory3DelegateProvider(Severity.ERROR)
+    return DiagnosticFactory3DelegateProvider(Severity.ERROR, LightTreePositioningStrategy.DEFAULT)
 }
 
 // ------------------------------ Providers ------------------------------
 
 class DiagnosticFactory0DelegateProvider<E : FirSourceElement, P : PsiElement>(
-    private val severity: Severity
+    private val severity: Severity,
+    private val positioningStrategy: LightTreePositioningStrategy
 ) {
     operator fun provideDelegate(thisRef: Any?, prop: KProperty<*>): ReadOnlyProperty<Any?, FirDiagnosticFactory0<E, P>> {
-        return DummyDelegate(FirDiagnosticFactory0(prop.name, severity))
+        return DummyDelegate(FirDiagnosticFactory0(prop.name, severity, positioningStrategy))
     }
 }
 
 class DiagnosticFactory1DelegateProvider<E : FirSourceElement, P : PsiElement, A : Any>(
-    private val severity: Severity
+    private val severity: Severity,
+    private val positioningStrategy: LightTreePositioningStrategy
 ) {
     operator fun provideDelegate(thisRef: Any?, prop: KProperty<*>): ReadOnlyProperty<Any?, FirDiagnosticFactory1<E, P, A>> {
-        return DummyDelegate(FirDiagnosticFactory1(prop.name, severity))
+        return DummyDelegate(FirDiagnosticFactory1(prop.name, severity, positioningStrategy))
     }
 }
 
 class DiagnosticFactory2DelegateProvider<E : FirSourceElement, P : PsiElement, A : Any, B : Any>(
-    private val severity: Severity
+    private val severity: Severity,
+    private val positioningStrategy: LightTreePositioningStrategy
 ) {
     operator fun provideDelegate(thisRef: Any?, prop: KProperty<*>): ReadOnlyProperty<Any?, FirDiagnosticFactory2<E, P, A, B>> {
-        return DummyDelegate(FirDiagnosticFactory2(prop.name, severity))
+        return DummyDelegate(FirDiagnosticFactory2(prop.name, severity, positioningStrategy))
     }
 }
 
 class DiagnosticFactory3DelegateProvider<E : FirSourceElement, P : PsiElement, A : Any, B : Any, C : Any>(
-    private val severity: Severity
+    private val severity: Severity,
+    private val positioningStrategy: LightTreePositioningStrategy
 ) {
     operator fun provideDelegate(thisRef: Any?, prop: KProperty<*>): ReadOnlyProperty<Any?, FirDiagnosticFactory3<E, P, A, B, C>> {
-        return DummyDelegate(FirDiagnosticFactory3(prop.name, severity))
+        return DummyDelegate(FirDiagnosticFactory3(prop.name, severity, positioningStrategy))
     }
 }
 

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

@@ -91,7 +91,7 @@ object FirErrors {
     val NON_CONST_VAL_USED_IN_CONSTANT_EXPRESSION by existing0<FirSourceElement, KtExpression>()
     val NOT_AN_ANNOTATION_CLASS by error1<FirSourceElement, PsiElement, String>()
     val NULLABLE_TYPE_OF_ANNOTATION_MEMBER by existing0<FirSourceElement, KtTypeReference>()
-    val VAR_ANNOTATION_PARAMETER by existing0<FirSourceElement, KtParameter>()
+    val VAR_ANNOTATION_PARAMETER by error0<FirSourceElement, KtParameter>(LightTreePositioningStrategies.VAL_OR_VAR_NODE)
 
     // Exposed visibility group
     val EXPOSED_TYPEALIAS_EXPANDED_TYPE by error3<FirSourceElement, PsiElement, FirEffectiveVisibility, FirMemberDeclaration, FirEffectiveVisibility>()

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

@@ -0,0 +1,35 @@
+/*
+ * 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.Ref
+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.psi.KtParameter.VAL_VAR_TOKEN_SET
+
+object LightTreePositioningStrategies {
+    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
+            return markElement(target, tree)
+        }
+    }
+}
+
+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? {
+    val childrenRef = Ref<Array<LighterASTNode>>()
+    getChildren(node, childrenRef)
+    return childrenRef.get()?.firstOrNull { it.tokenType in type }
+}

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

@@ -188,7 +188,13 @@ sealed class FirPsiSourceElement<out P : PsiElement>(val psi: P) : FirSourceElem
             node.unwrap().psi.parent?.node?.let { TreeBackedLighterAST.wrap(it) }
 
         override fun getChildren(node: LighterASTNode, nodesRef: Ref<Array<LighterASTNode>>): Int {
-            val children = node.unwrap().psi.children
+            val psi = node.unwrap().psi
+            val children = mutableListOf<PsiElement>()
+            var child = psi.firstChild
+            while (child != null) {
+                children += child
+                child = child.nextSibling
+            }
             if (children.isEmpty()) {
                 nodesRef.set(LighterASTNode.EMPTY_ARRAY)
             } else {

compiler/psi/src/org/jetbrains/kotlin/psi/KtParameter.java

@@ -129,7 +129,7 @@ public class KtParameter extends KtNamedDeclarationStub<KotlinParameterStub> imp
         return findChildByType(KtNodeTypes.DESTRUCTURING_DECLARATION);
     }
 
-    private static final TokenSet VAL_VAR_TOKEN_SET = TokenSet.create(KtTokens.VAL_KEYWORD, KtTokens.VAR_KEYWORD);
+    public static final TokenSet VAL_VAR_TOKEN_SET = TokenSet.create(KtTokens.VAL_KEYWORD, KtTokens.VAR_KEYWORD);
 
     @Override
     public ItemPresentation getPresentation() {