[FIR][LightTree] Better expression identification

And some refactoring Regenerate IDE parts
2021-04-19 20:20:28 +03:00
parent 2f7d6da22f
commit 5fc1514104
6 changed files with 34 additions and 23 deletions
@@ -8,7 +8,6 @@ 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.impl.source.tree.ElementType
 import com.intellij.psi.tree.IElementType
 import com.intellij.psi.tree.TokenSet
 import com.intellij.util.diff.FlyweightCapableTreeStructure
@@ -19,8 +18,8 @@ import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.lexer.KtTokens.MODALITY_MODIFIERS
 import org.jetbrains.kotlin.lexer.KtTokens.VISIBILITY_MODIFIERS
 import org.jetbrains.kotlin.psi.KtParameter.VAL_VAR_TOKEN_SET
-import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.stubs.elements.KtConstantExpressionElementType
+import org.jetbrains.kotlin.psi.stubs.elements.KtStringTemplateExpressionElementType
 import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult

 object LightTreePositioningStrategies {
@@ -644,20 +643,23 @@ internal fun FlyweightCapableTreeStructure<LighterASTNode>.nameIdentifier(node:
 private fun FlyweightCapableTreeStructure<LighterASTNode>.operationReference(node: LighterASTNode): LighterASTNode? =
    findChildByType(node, KtNodeTypes.OPERATION_REFERENCE)

-private val EXPRESSION_NODE_TYPES = setOf(
+private val EXPRESSIONS_SET = listOf(
    KtNodeTypes.REFERENCE_EXPRESSION,
-    KtNodeTypes.CONSTRUCTOR_DELEGATION_REFERENCE,
-    KtNodeTypes.SUPER_EXPRESSION,
-    KtNodeTypes.ARRAY_ACCESS_EXPRESSION,
-    KtNodeTypes.CALL_EXPRESSION,
-    KtNodeTypes.LABELED_EXPRESSION,
    KtNodeTypes.DOT_QUALIFIED_EXPRESSION,
-    KtNodeTypes.FUN,
-    KtNodeTypes.IF,
-    KtNodeTypes.STRING_TEMPLATE,
    KtNodeTypes.LAMBDA_EXPRESSION,
+    KtNodeTypes.FUN
 )

+fun LighterASTNode.isExpression(): Boolean {
+    return when (this.tokenType) {
+        is KtNodeType,
+        is KtConstantExpressionElementType,
+        is KtStringTemplateExpressionElementType,
+        in EXPRESSIONS_SET -> true
+        else -> false
+    }
+}
+
 /**
 * @param locateReferencedName whether to remove any nested parentheses while locating the reference element. This is useful for diagnostics
 * on super and unresolved references. For example, with the following, only the part inside the parentheses should be highlighted.
@@ -678,9 +680,7 @@ private fun FlyweightCapableTreeStructure<LighterASTNode>.referenceExpression(
    val childrenRef = Ref<Array<LighterASTNode?>>()
    getChildren(node, childrenRef)
    var result = childrenRef.get()?.firstOrNull {
-        it?.tokenType in EXPRESSION_NODE_TYPES ||
-                it?.tokenType is KtConstantExpressionElementType ||
-                it?.tokenType == KtNodeTypes.PARENTHESIZED
+        it?.isExpression() == true || it?.tokenType == KtNodeTypes.PARENTHESIZED
    }
    while (locateReferencedName && result != null && result.tokenType == KtNodeTypes.PARENTHESIZED) {
        result = referenceExpression(result, locateReferencedName = true)
@@ -689,7 +689,7 @@ private fun FlyweightCapableTreeStructure<LighterASTNode>.referenceExpression(
 }

 private fun FlyweightCapableTreeStructure<LighterASTNode>.findReferenceExpressionDeep(node: LighterASTNode): LighterASTNode? =
-    findDescendantByTypes(node, EXPRESSION_NODE_TYPES)
+    findFirstDescendant(node) { it.isExpression() }

 private fun FlyweightCapableTreeStructure<LighterASTNode>.rightParenthesis(node: LighterASTNode): LighterASTNode? =
    findChildByType(node, KtTokens.RPAR)
@@ -802,6 +802,16 @@ fun FlyweightCapableTreeStructure<LighterASTNode>.findDescendantByTypes(node: Li
        ?.firstNotNullResult { child -> child?.let { findDescendantByTypes(it, types) } }
 }

+fun FlyweightCapableTreeStructure<LighterASTNode>.findFirstDescendant(
+    node: LighterASTNode,
+    predicate: (LighterASTNode) -> Boolean
+): LighterASTNode? {
+    val childrenRef = Ref<Array<LighterASTNode?>>()
+    getChildren(node, childrenRef)
+    return childrenRef.get()?.firstOrNull { it != null && predicate(it) }
+        ?: childrenRef.get()?.firstNotNullResult { child -> child?.let { findFirstDescendant(it, predicate) } }
+}
+
 private fun FlyweightCapableTreeStructure<LighterASTNode>.findChildByType(node: LighterASTNode, type: TokenSet): LighterASTNode? {
    val childrenRef = Ref<Array<LighterASTNode?>>()
    getChildren(node, childrenRef)
@@ -16,7 +16,7 @@ open enum class Z : Enum<Z> {
    fun test() {
    }

-    inner class A {
+    local inner class A {
      constructor() /* primary */ {
        super/*Any*/()
        /* <init>() */
@@ -36,3 +36,4 @@ open enum class Z : Enum<Z> {
  fun valueOf(value: String): Z /* Synthetic body for ENUM_VALUEOF */

 }
+
@@ -113,8 +113,8 @@ class FirDiagnosticsHandler(testServices: TestServices) : FirAnalysisHandler(tes
                FirErrors.SYNTAX.on(FirRealPsiSourceElement(it)).toMetaInfo(testFile, lightTreeEnabled, lightTreeComparingModeEnabled)
            }
        } else {
-            collectLightTreeSyntaxErrors(firFile).map { node ->
-                FirErrors.SYNTAX.on(node).toMetaInfo(testFile, lightTreeEnabled, lightTreeComparingModeEnabled)
+            collectLightTreeSyntaxErrors(firFile).map { sourceElement ->
+                FirErrors.SYNTAX.on(sourceElement).toMetaInfo(testFile, lightTreeEnabled, lightTreeComparingModeEnabled)
            }
        }

@@ -48,12 +48,12 @@ private class LightTreeErrorsCollector(private val tree: Tree) {

 }

-private data class OffsetTree(val tree: Tree, val offset: Int)
+private data class TreeWithOffset(val tree: Tree, val offset: Int)

 private data class VisitorState(
    val lastTree: Tree? = null,
    val visitedTrees: MutableSet<Tree> = mutableSetOf(),
-    val result: MutableList<OffsetTree> = mutableListOf()
+    val result: MutableList<TreeWithOffset> = mutableListOf()
 )

 private object FirTreesExtractVisitor : FirVisitor<Unit, VisitorState>() {
@@ -62,7 +62,7 @@ private object FirTreesExtractVisitor : FirVisitor<Unit, VisitorState>() {
        val currentTree = source.treeStructure
        val (lastTree, visitedTrees, result) = data
        val newData = if (lastTree !== currentTree && visitedTrees.add(currentTree)) {
-            val newTreeWithOffset = OffsetTree(currentTree, element.source?.startOffset ?: 0)
+            val newTreeWithOffset = TreeWithOffset(currentTree, element.source?.startOffset ?: 0)
            result.add(newTreeWithOffset)
            data.copy(lastTree = lastTree)
        } else {
@@ -759,7 +759,7 @@ sealed class KtFirDiagnostic<PSI: PsiElement> : KtDiagnosticWithPsi<PSI> {
        override val diagnosticClass get() = CyclicGenericUpperBound::class
    }

-    abstract class DeprecatedTypeParameterSyntax : KtFirDiagnostic<KtTypeParameterList>() {
+    abstract class DeprecatedTypeParameterSyntax : KtFirDiagnostic<KtDeclaration>() {
        override val diagnosticClass get() = DeprecatedTypeParameterSyntax::class
    }

@@ -1228,7 +1228,7 @@ internal class CyclicGenericUpperBoundImpl(
 internal class DeprecatedTypeParameterSyntaxImpl(
    firDiagnostic: FirPsiDiagnostic<*>,
    override val token: ValidityToken,
-) : KtFirDiagnostic.DeprecatedTypeParameterSyntax(), KtAbstractFirDiagnostic<KtTypeParameterList> {
+) : KtFirDiagnostic.DeprecatedTypeParameterSyntax(), KtAbstractFirDiagnostic<KtDeclaration> {
    override val firDiagnostic: FirPsiDiagnostic<*> by weakRef(firDiagnostic)
 }