FIR checkers: extract getChildren(), simplify findSuperTypeDelegation()

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/FirModifierList.kt

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.fir.analysis.checkers
 
 import com.intellij.lang.ASTNode
 import com.intellij.lang.LighterASTNode
-import com.intellij.openapi.util.Ref
 import com.intellij.psi.PsiElement
 import com.intellij.psi.tree.TokenSet
 import com.intellij.util.diff.FlyweightCapableTreeStructure
@@ -27,9 +26,8 @@ fun FirSourceElement?.getModifierList(): FirModifierList? {
         null -> null
         is FirPsiSourceElement<*> -> (psi as? KtModifierListOwner)?.modifierList?.let { FirPsiModifierList(it) }
         is FirLightSourceElement -> {
-            val kidsRef = Ref<Array<LighterASTNode?>>()
+            val modifierListNode = lighterASTNode.getChildren(treeStructure).find { it?.tokenType == KtNodeTypes.MODIFIER_LIST }
-            treeStructure.getChildren(lighterASTNode, kidsRef)
+                ?: return null
-            val modifierListNode = kidsRef.get().find { it?.tokenType == KtNodeTypes.MODIFIER_LIST } ?: return null
             FirLightModifierList(modifierListNode, treeStructure)
         }
     }
@@ -48,9 +46,7 @@ class FirPsiModifierList(val modifierList: KtModifierList) : FirModifierList() {
 class FirLightModifierList(val modifierList: LighterASTNode, val tree: FlyweightCapableTreeStructure<LighterASTNode>) : FirModifierList() {
     override val modifiers: List<FirLightModifier>
         get() {
-            val kidsRef = Ref<Array<LighterASTNode?>>()
+            val modifierNodes = modifierList.getChildren(tree)
-            tree.getChildren(modifierList, kidsRef)
-            val modifierNodes = kidsRef.get()
             return modifierNodes.filterNotNull()
                 .filter { it.tokenType is KtModifierKeywordToken }
                 .map { FirLightModifier(it, it.tokenType as KtModifierKeywordToken, tree) }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/SourceHelpers.kt

@@ -0,0 +1,17 @@
+/*
+ * 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.checkers
+
+import com.intellij.lang.LighterASTNode
+import com.intellij.openapi.util.Ref
+import com.intellij.util.diff.FlyweightCapableTreeStructure
+
+internal fun LighterASTNode.getChildren(tree: FlyweightCapableTreeStructure<LighterASTNode>): List<LighterASTNode?> {
+    val children = Ref<Array<LighterASTNode?>>()
+    val count = tree.getChildren(this, children)
+    return if (count > 0) children.get().filterNotNull() else emptyList()
+}
+

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirDelegationInInterfaceChecker.kt

@@ -6,21 +6,16 @@
 package org.jetbrains.kotlin.fir.analysis.checkers.declaration
 
 import com.intellij.lang.LighterASTNode
-import com.intellij.openapi.util.Ref
-import com.intellij.psi.PsiElement
-import com.intellij.psi.PsiErrorElement
 import com.intellij.util.diff.FlyweightCapableTreeStructure
 import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.descriptors.ClassKind
-import org.jetbrains.kotlin.fir.FirLightSourceElement
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.getChildren
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
-import org.jetbrains.kotlin.fir.psi
-import org.jetbrains.kotlin.psi.KtDelegatedSuperTypeEntry
 
 object FirDelegationInInterfaceChecker : FirMemberDeclarationChecker() {
     override fun check(declaration: FirMemberDeclaration, context: CheckerContext, reporter: DiagnosticReporter) {
@@ -33,46 +28,21 @@ object FirDelegationInInterfaceChecker : FirMemberDeclarationChecker() {
         }
     }
 
-    private fun FirSourceElement.findSuperTypeDelegation(): Int {
+    private fun FirSourceElement.findSuperTypeDelegation(): Int =
-        val localPsi = psi
+        lighterASTNode.findSuperTypeDelegation(treeStructure)
-        val localLightNode = lighterASTNode
-
-        if (localPsi != null && localPsi !is PsiErrorElement) {
-            return localPsi.findSuperTypeDelegation()
-        } else if (this is FirLightSourceElement) {
-            return localLightNode.findSuperTypeDelegation(treeStructure)
-        }
-
-        return -1
-    }
-
-    private fun PsiElement.findSuperTypeDelegation(): Int {
-        val children = this.children // this is a method call and it collects children
-        return if (children.isNotEmpty() && children[0] !is PsiErrorElement) {
-            children[0].children.indexOfFirst { it is KtDelegatedSuperTypeEntry }
-        } else {
-            -1
-        }
-    }
 
     private fun LighterASTNode.findSuperTypeDelegation(tree: FlyweightCapableTreeStructure<LighterASTNode>): Int {
         val children = getChildren(tree)
         return if (children.isNotEmpty()) {
-            children.find { it.tokenType == KtNodeTypes.SUPER_TYPE_LIST }
+            children.find { it?.tokenType == KtNodeTypes.SUPER_TYPE_LIST }
                 ?.getChildren(tree)
-                ?.indexOfFirst { it.tokenType == KtNodeTypes.DELEGATED_SUPER_TYPE_ENTRY }
+                ?.indexOfFirst { it?.tokenType == KtNodeTypes.DELEGATED_SUPER_TYPE_ENTRY }
                 ?: -1
         } else {
             -1
         }
     }
 
-    private fun LighterASTNode.getChildren(tree: FlyweightCapableTreeStructure<LighterASTNode>): List<LighterASTNode> {
-        val children = Ref<Array<LighterASTNode?>>()
-        val count = tree.getChildren(this, children)
-        return if (count > 0) children.get().filterNotNull() else emptyList()
-    }
-
     private fun DiagnosticReporter.report(source: FirSourceElement?) {
         source?.let { report(FirErrors.DELEGATION_IN_INTERFACE.on(it)) }
     }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirExposedVisibilityDeclarationChecker.kt

@@ -5,13 +5,12 @@
 
 package org.jetbrains.kotlin.fir.analysis.checkers.declaration
 
-import com.intellij.lang.LighterASTNode
-import com.intellij.openapi.util.Ref
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiNameIdentifierOwner
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.getChildren
 import org.jetbrains.kotlin.fir.analysis.checkers.toRegularClass
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnosticFactory3
@@ -233,10 +232,8 @@ object FirExposedVisibilityDeclarationChecker : FirMemberDeclarationChecker() {
     private fun FirSourceElement.getIdentifierSource() = when (this) {
         is FirPsiSourceElement<*> -> (this.psi as? PsiNameIdentifierOwner)?.nameIdentifier?.toFirPsiSourceElement()
         is FirLightSourceElement -> {
-            val kidsRef = Ref<Array<LighterASTNode?>>()
+            val identifier = lighterASTNode.getChildren(treeStructure).find { it?.tokenType == KtTokens.IDENTIFIER }
-            this.treeStructure.getChildren(lighterASTNode, kidsRef)
+            identifier?.toFirLightSourceElement(treeStructure)
-            val identifier = kidsRef.get().find { it?.tokenType == KtTokens.IDENTIFIER }
-            identifier?.toFirLightSourceElement(this.treeStructure)
         }
     }
 

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirSupertypeInitializedInInterfaceChecker.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.fir.analysis.checkers.declaration
 
 import com.intellij.lang.LighterASTNode
-import com.intellij.openapi.util.Ref
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiErrorElement
 import com.intellij.util.diff.FlyweightCapableTreeStructure
@@ -15,6 +14,7 @@ import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.FirLightSourceElement
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.getChildren
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.declarations.FirClass
@@ -58,21 +58,15 @@ object FirSupertypeInitializedInInterfaceChecker : FirMemberDeclarationChecker()
     private fun LighterASTNode.findSuperTypeCall(tree: FlyweightCapableTreeStructure<LighterASTNode>): Int {
         val children = getChildren(tree)
         return if (children.isNotEmpty()) {
-            children.find { it.tokenType == KtNodeTypes.SUPER_TYPE_LIST }
+            children.find { it?.tokenType == KtNodeTypes.SUPER_TYPE_LIST }
                 ?.getChildren(tree)
-                ?.indexOfFirst { it.tokenType == KtNodeTypes.SUPER_TYPE_CALL_ENTRY }
+                ?.indexOfFirst { it?.tokenType == KtNodeTypes.SUPER_TYPE_CALL_ENTRY }
                 ?: -1
         } else {
             -1
         }
     }
 
-    private fun LighterASTNode.getChildren(tree: FlyweightCapableTreeStructure<LighterASTNode>): List<LighterASTNode> {
-        val children = Ref<Array<LighterASTNode?>>()
-        val count = tree.getChildren(this, children)
-        return if (count > 0) children.get().filterNotNull() else emptyList()
-    }
-
     private fun DiagnosticReporter.report(source: FirSourceElement?) {
         source?.let { report(FirErrors.SUPERTYPE_INITIALIZED_IN_INTERFACE.on(it)) }
     }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirSupertypeInitializedWithoutPrimaryConstructor.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.fir.analysis.checkers.declaration
 
 import com.intellij.lang.LighterASTNode
-import com.intellij.openapi.util.Ref
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiErrorElement
 import com.intellij.util.diff.FlyweightCapableTreeStructure
@@ -15,6 +14,7 @@ import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.FirLightSourceElement
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.getChildren
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.declarations.FirConstructor
@@ -56,16 +56,10 @@ object FirSupertypeInitializedWithoutPrimaryConstructor : FirMemberDeclarationCh
     }
 
     private fun LighterASTNode.anySupertypeHasConstructorParentheses(tree: FlyweightCapableTreeStructure<LighterASTNode>): Boolean {
-        val superTypes = getChildren(tree).find { it.tokenType == KtNodeTypes.SUPER_TYPE_LIST }
+        val superTypes = getChildren(tree).find { it?.tokenType == KtNodeTypes.SUPER_TYPE_LIST }
             ?: return false
 
-        return superTypes.getChildren(tree).any { it.tokenType == KtNodeTypes.SUPER_TYPE_CALL_ENTRY }
+        return superTypes.getChildren(tree).any { it?.tokenType == KtNodeTypes.SUPER_TYPE_CALL_ENTRY }
-    }
-
-    private fun LighterASTNode.getChildren(tree: FlyweightCapableTreeStructure<LighterASTNode>): List<LighterASTNode> {
-        val children = Ref<Array<LighterASTNode?>>()
-        val count = tree.getChildren(this, children)
-        return if (count > 0) children.get().filterNotNull() else emptyList()
     }
 
     private fun DiagnosticReporter.report(source: FirSourceElement?) {

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirTypeArgumentsNotAllowedExpressionChecker.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.KtNodeTypes.TYPE_ARGUMENT_LIST
 import org.jetbrains.kotlin.fir.FirLightSourceElement
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.analysis.checkers.getChildren
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
@@ -58,12 +59,6 @@ object FirTypeArgumentsNotAllowedExpressionChecker : FirQualifiedAccessChecker()
         return children.count { it != null } > 1 && children[1]?.tokenType == TYPE_ARGUMENT_LIST
     }
 
-    private fun LighterASTNode.getChildren(tree: FlyweightCapableTreeStructure<LighterASTNode>): Array<out LighterASTNode?> {
-        val childrenRef = Ref<Array<LighterASTNode>>()
-        val childCount = tree.getChildren(this, childrenRef)
-        return if (childCount > 0) childrenRef.get() else emptyArray()
-    }
-
     private fun DiagnosticReporter.report(source: FirSourceElement?) {
         source?.let {
             report(FirErrors.TYPE_ARGUMENTS_NOT_ALLOWED.on(it))

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/extended/CanBeReplacedWithOperatorAssignmentChecker.kt

@@ -6,12 +6,12 @@
 package org.jetbrains.kotlin.fir.analysis.checkers.extended
 
 import com.intellij.lang.LighterASTNode
-import com.intellij.openapi.util.Ref
 import com.intellij.psi.tree.IElementType
 import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirExpressionChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.getChildren
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.analysis.getChild
@@ -72,9 +72,7 @@ object CanBeReplacedWithOperatorAssignmentChecker : FirExpressionChecker<FirVari
         prevOperator: LighterASTNode? = null
     ): Boolean {
         val tree = source.treeStructure
-        val childrenNullable = Ref<Array<LighterASTNode?>>()
+        val children = expression.getChildren(tree).filterNotNull()
-        tree.getChildren(expression, childrenNullable)
-        val children = childrenNullable.get().filterNotNull()
 
         val operator = children.firstOrNull { it.tokenType == KtNodeTypes.OPERATION_REFERENCE }
         if (prevOperator != null && !isLightNodesHierarchicallyTrue(prevOperator, operator)) return false

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/extended/CanBeValChecker.kt

@@ -5,12 +5,11 @@
 
 package org.jetbrains.kotlin.fir.analysis.checkers.extended
 
-import com.intellij.lang.LighterASTNode
-import com.intellij.openapi.util.Ref
 import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.contracts.description.EventOccurrencesRange
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.analysis.cfa.*
+import org.jetbrains.kotlin.fir.analysis.checkers.getChildren
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.analysis.getChild
@@ -103,9 +102,8 @@ object CanBeValChecker : AbstractFirPropertyInitializationChecker() {
         is FirLightSourceElement -> {
             val source = fir.source as FirLightSourceElement
             val tree = (fir.source as FirLightSourceElement).treeStructure
-            val children = Ref<Array<LighterASTNode?>>()
+            val children = source.lighterASTNode.getChildren(tree)
-            tree.getChildren(source.lighterASTNode, children)
+            children.filter { it?.tokenType == KtNodeTypes.DESTRUCTURING_DECLARATION_ENTRY }.size
-            children.get().filterNotNull().filter { it.tokenType == KtNodeTypes.DESTRUCTURING_DECLARATION_ENTRY }.size
         }
         else -> null
     }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/extended/RedundantSingleExpressionStringTemplateChecker.kt

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.fir.analysis.checkers.extended
 
 import com.intellij.lang.LighterASTNode
 import com.intellij.lang.PsiBuilder
-import com.intellij.openapi.util.Ref
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.fir.FirFakeSourceElementKind
@@ -15,6 +14,7 @@ import org.jetbrains.kotlin.fir.FirLightSourceElement
 import org.jetbrains.kotlin.fir.FirPsiSourceElement
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirBasicExpressionChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.getChildren
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.REDUNDANT_SINGLE_EXPRESSION_STRING_TEMPLATE
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
@@ -55,10 +55,9 @@ object RedundantSingleExpressionStringTemplateChecker : FirBasicExpressionChecke
 
     private fun LighterASTNode.stringParentChildrenCount(source: FirLightSourceElement): Int? {
         val parent = source.treeStructure.getParent(this)
-        return if (parent?.tokenType == KtNodeTypes.STRING_TEMPLATE) {
+        return if (parent != null && parent.tokenType == KtNodeTypes.STRING_TEMPLATE) {
-            val childrenOfParent = Ref<Array<LighterASTNode>>()
+            val childrenOfParent = parent.getChildren(source.treeStructure)
-            source.treeStructure.getChildren(parent!!, childrenOfParent)
+            childrenOfParent.filter { it is PsiBuilder.Marker }.size
-            childrenOfParent.get().filter { it is PsiBuilder.Marker }.size
         } else {
             parent?.stringParentChildrenCount(source)
         }