[FIR] Implement light tree DECLARATION_NAME & SIGNATURE strategies

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

@@ -66,8 +66,8 @@ object FirErrors {
     val SEALED_SUPERTYPE_IN_LOCAL_CLASS by error0<FirSourceElement, PsiElement>()
 
     // Constructor problems
-    val CONSTRUCTOR_IN_OBJECT by existing0<FirSourceElement, KtDeclaration>()
+    val CONSTRUCTOR_IN_OBJECT by error0<FirSourceElement, KtDeclaration>(LightTreePositioningStrategies.DECLARATION_SIGNATURE)
-    val CONSTRUCTOR_IN_INTERFACE by existing0<FirSourceElement, KtDeclaration>()
+    val CONSTRUCTOR_IN_INTERFACE by error0<FirSourceElement, KtDeclaration>(LightTreePositioningStrategies.DECLARATION_SIGNATURE)
     val NON_PRIVATE_CONSTRUCTOR_IN_ENUM by existing0<FirSourceElement, PsiElement>()
     val NON_PRIVATE_CONSTRUCTOR_IN_SEALED by existing0<FirSourceElement, PsiElement>()
     val CYCLIC_CONSTRUCTOR_DELEGATION_CALL by warning0<FirSourceElement, PsiElement>()
@@ -144,11 +144,11 @@ object FirErrors {
     val ANY_METHOD_IMPLEMENTED_IN_INTERFACE by error0<FirSourceElement, PsiElement>()
 
     // Invalid local declarations
-    val LOCAL_OBJECT_NOT_ALLOWED by error1<FirSourceElement, PsiElement, Name>()
+    val LOCAL_OBJECT_NOT_ALLOWED by error1<FirSourceElement, PsiElement, Name>(LightTreePositioningStrategies.DECLARATION_NAME)
-    val LOCAL_INTERFACE_NOT_ALLOWED by error1<FirSourceElement, PsiElement, Name>()
+    val LOCAL_INTERFACE_NOT_ALLOWED by error1<FirSourceElement, PsiElement, Name>(LightTreePositioningStrategies.DECLARATION_NAME)
 
     // Control flow diagnostics
-    val UNINITIALIZED_VARIABLE by error1<FirSourceElement, PsiElement, FirPropertySymbol>()
+    val UNINITIALIZED_VARIABLE by error1<FirSourceElement, PsiElement, FirPropertySymbol>(LightTreePositioningStrategies.DECLARATION_SIGNATURE)
     val WRONG_INVOCATION_KIND by warning3<FirSourceElement, PsiElement, AbstractFirBasedSymbol<*>, EventOccurrencesRange, EventOccurrencesRange>()
     val LEAKED_IN_PLACE_LAMBDA by error1<FirSourceElement, PsiElement, AbstractFirBasedSymbol<*>>()
     val WRONG_IMPLIES_CONDITION by error0<FirSourceElement, PsiElement>()

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

@@ -20,10 +20,10 @@ object LightTreePositioningStrategies {
         override fun mark(node: LighterASTNode, tree: FlyweightCapableTreeStructure<LighterASTNode>): List<TextRange> {
             when (node.tokenType) {
                 KtNodeTypes.OBJECT_DECLARATION -> {
-                    val objectKeyword = tree.findChildByType(node, KtTokens.OBJECT_KEYWORD)!!
+                    val objectKeyword = tree.objectKeyword(node)!!
                     return markRange(
                         from = objectKeyword,
-                        to = tree.findChildByType(node, KtTokens.IDENTIFIER) ?: objectKeyword,
+                        to = tree.nameIdentifier(node) ?: objectKeyword,
                         tree
                     )
                 }
@@ -46,9 +46,9 @@ object LightTreePositioningStrategies {
         override fun mark(node: LighterASTNode, tree: FlyweightCapableTreeStructure<LighterASTNode>): List<TextRange> {
             when (node.tokenType) {
                 KtNodeTypes.SECONDARY_CONSTRUCTOR -> {
-                    val valueParameterList = tree.findChildByType(node, KtNodeTypes.VALUE_PARAMETER_LIST) ?: return markElement(node, tree)
+                    val valueParameterList = tree.valueParameterList(node) ?: return markElement(node, tree)
                     return markRange(
-                        tree.findChildByType(node, KtTokens.CONSTRUCTOR_KEYWORD)!!,
+                        tree.constructorKeyword(node)!!,
                         tree.lastChild(valueParameterList)!!, tree
                     )
                 }
@@ -56,10 +56,10 @@ object LightTreePositioningStrategies {
                     val delegationReference = tree.findChildByType(node, KtNodeTypes.CONSTRUCTOR_DELEGATION_REFERENCE)
                     if (delegationReference != null && tree.firstChild(delegationReference) == null) {
                         val constructor = tree.findParentOfType(node, KtNodeTypes.SECONDARY_CONSTRUCTOR)!!
-                        val valueParameterList = tree.findChildByType(constructor, KtNodeTypes.VALUE_PARAMETER_LIST)
+                        val valueParameterList = tree.valueParameterList(constructor)
                             ?: return markElement(constructor, tree)
                         return markRange(
-                            tree.findChildByType(constructor, KtTokens.CONSTRUCTOR_KEYWORD)!!,
+                            tree.constructorKeyword(constructor)!!,
                             tree.lastChild(valueParameterList)!!, tree
                         )
                     }
@@ -69,6 +69,120 @@ object LightTreePositioningStrategies {
             }
         }
     }
+
+    val DECLARATION_NAME: LightTreePositioningStrategy = object : LightTreePositioningStrategy() {
+        override fun mark(node: LighterASTNode, tree: FlyweightCapableTreeStructure<LighterASTNode>): List<TextRange> {
+            val nameIdentifier = tree.nameIdentifier(node)
+            if (nameIdentifier != null) {
+                if (node.tokenType == KtNodeTypes.CLASS || node.tokenType == KtNodeTypes.OBJECT_DECLARATION) {
+                    val startElement =
+                        tree.modifierList(node)?.let { modifierList -> tree.findChildByType(modifierList, KtTokens.ENUM_KEYWORD) }
+                            ?: tree.findChildByType(node, TokenSet.create(KtTokens.CLASS_KEYWORD, KtTokens.OBJECT_KEYWORD))
+                            ?: node
+
+                    return markRange(startElement, nameIdentifier, tree)
+                }
+                return markElement(nameIdentifier, tree)
+            }
+            if (node.tokenType == KtNodeTypes.FUN) {
+                return DECLARATION_SIGNATURE.mark(node, tree)
+            }
+            return DEFAULT.mark(node, tree)
+        }
+    }
+
+    val DECLARATION_SIGNATURE: LightTreePositioningStrategy = object : LightTreePositioningStrategy() {
+        override fun mark(node: LighterASTNode, tree: FlyweightCapableTreeStructure<LighterASTNode>): List<TextRange> {
+            when (node.tokenType) {
+                KtNodeTypes.PRIMARY_CONSTRUCTOR, KtNodeTypes.SECONDARY_CONSTRUCTOR -> {
+                    val begin = tree.constructorKeyword(node) ?: tree.valueParameterList(node) ?: return markElement(node, tree)
+                    val end = tree.valueParameterList(node) ?: tree.constructorKeyword(node) ?: return markElement(node, tree)
+                    return markRange(begin, end, tree)
+                }
+                KtNodeTypes.FUN, KtNodeTypes.FUNCTION_LITERAL -> {
+                    val endOfSignatureElement =
+                        tree.typeReference(node)
+                            ?: tree.valueParameterList(node)
+                            ?: tree.nameIdentifier(node)
+                            ?: node
+                    val startElement = if (node.tokenType == KtNodeTypes.FUNCTION_LITERAL) {
+                        tree.receiverTypeReference(node)
+                            ?: tree.valueParameterList(node)
+                            ?: node
+                    } else node
+                    return markRange(startElement, endOfSignatureElement, tree)
+                }
+                KtNodeTypes.PROPERTY -> {
+                    val endOfSignatureElement = tree.typeReference(node) ?: tree.nameIdentifier(node) ?: node
+                    return markRange(node, endOfSignatureElement, tree)
+                }
+                KtNodeTypes.PROPERTY_ACCESSOR -> {
+                    val endOfSignatureElement =
+                        tree.typeReference(node)
+                            ?: tree.rightParenthesis(node)
+                            ?: tree.accessorNamePlaceholder(node)
+
+                    return markRange(node, endOfSignatureElement, tree)
+                }
+                KtNodeTypes.CLASS -> {
+                    val nameAsDeclaration = tree.nameIdentifier(node) ?: return markElement(node, tree)
+                    val primaryConstructorParameterList = tree.primaryConstructor(node)?.let { constructor ->
+                        tree.valueParameterList(constructor)
+                    } ?: return markElement(nameAsDeclaration, tree)
+                    return markRange(nameAsDeclaration, primaryConstructorParameterList, tree)
+                }
+                KtNodeTypes.OBJECT_DECLARATION -> {
+                    return DECLARATION_NAME.mark(node, tree)
+                }
+                KtNodeTypes.CLASS_INITIALIZER -> {
+                    return markElement(tree.initKeyword(node)!!, tree)
+                }
+            }
+            return super.mark(node, tree)
+        }
+    }
+}
+
+private fun FlyweightCapableTreeStructure<LighterASTNode>.constructorKeyword(node: LighterASTNode): LighterASTNode? =
+    findChildByType(node, KtTokens.CONSTRUCTOR_KEYWORD)
+
+private fun FlyweightCapableTreeStructure<LighterASTNode>.initKeyword(node: LighterASTNode): LighterASTNode? =
+    findChildByType(node, KtTokens.INIT_KEYWORD)
+
+private fun FlyweightCapableTreeStructure<LighterASTNode>.nameIdentifier(node: LighterASTNode): LighterASTNode? =
+    findChildByType(node, KtTokens.IDENTIFIER)
+
+private fun FlyweightCapableTreeStructure<LighterASTNode>.rightParenthesis(node: LighterASTNode): LighterASTNode? =
+    findChildByType(node, KtTokens.RPAR)
+
+private fun FlyweightCapableTreeStructure<LighterASTNode>.objectKeyword(node: LighterASTNode): LighterASTNode? =
+    findChildByType(node, KtTokens.OBJECT_KEYWORD)
+
+private fun FlyweightCapableTreeStructure<LighterASTNode>.accessorNamePlaceholder(node: LighterASTNode): LighterASTNode =
+    findChildByType(node, KtTokens.GET_KEYWORD) ?: findChildByType(node, KtTokens.SET_KEYWORD)!!
+
+private fun FlyweightCapableTreeStructure<LighterASTNode>.modifierList(node: LighterASTNode): LighterASTNode? =
+    findChildByType(node, KtNodeTypes.MODIFIER_LIST)
+
+private fun FlyweightCapableTreeStructure<LighterASTNode>.primaryConstructor(node: LighterASTNode): LighterASTNode? =
+    findChildByType(node, KtNodeTypes.PRIMARY_CONSTRUCTOR)
+
+private fun FlyweightCapableTreeStructure<LighterASTNode>.valueParameterList(node: LighterASTNode): LighterASTNode? =
+    findChildByType(node, KtNodeTypes.VALUE_PARAMETER_LIST)
+
+private fun FlyweightCapableTreeStructure<LighterASTNode>.typeReference(node: LighterASTNode): LighterASTNode? {
+    val childrenRef = Ref<Array<LighterASTNode>>()
+    getChildren(node, childrenRef)
+    return childrenRef.get()?.dropWhile { it.tokenType != KtTokens.COLON }?.firstOrNull { it.tokenType == KtNodeTypes.TYPE_REFERENCE }
+}
+
+private fun FlyweightCapableTreeStructure<LighterASTNode>.receiverTypeReference(node: LighterASTNode): LighterASTNode? {
+    val childrenRef = Ref<Array<LighterASTNode>>()
+    getChildren(node, childrenRef)
+    return childrenRef.get()?.firstOrNull {
+        if (it.tokenType == KtTokens.COLON || it.tokenType == KtTokens.LPAR) return null
+        it.tokenType == KtNodeTypes.TYPE_REFERENCE
+    }
 }
 
 private fun FlyweightCapableTreeStructure<LighterASTNode>.findChildByType(node: LighterASTNode, type: IElementType): LighterASTNode? {

compiler/testData/diagnostics/tests/funInterface/funInterfaceDeclarationCheck.fir.kt

@@ -87,9 +87,9 @@ class WithNestedFun<K> {
 }
 
 fun <T> local() {
-    <!LOCAL_INTERFACE_NOT_ALLOWED!>fun interface LocalFun {
+    <!LOCAL_INTERFACE_NOT_ALLOWED!>fun interface LocalFun<!> {
         fun invoke(element: T)
-    }<!>
+    }
 }
 
 fun interface WithDefaultValue {

compiler/testData/diagnostics/tests/localInterfaces.fir.kt

@@ -1,14 +1,14 @@
 // !DIAGNOSTICS: -UNUSED_VARIABLE
 
 fun foo() {
-    <!LOCAL_INTERFACE_NOT_ALLOWED!>interface a {}<!>
+    <!LOCAL_INTERFACE_NOT_ALLOWED!>interface a<!> {}
     val b = object {
-        <!LOCAL_INTERFACE_NOT_ALLOWED!>interface c {}<!>
+        <!LOCAL_INTERFACE_NOT_ALLOWED!>interface c<!> {}
     }
     class A {
-        <!LOCAL_INTERFACE_NOT_ALLOWED!>interface d {}<!>
+        <!LOCAL_INTERFACE_NOT_ALLOWED!>interface d<!> {}
     }
     val f = {
-        <!LOCAL_INTERFACE_NOT_ALLOWED!>interface e {}<!>
+        <!LOCAL_INTERFACE_NOT_ALLOWED!>interface e<!> {}
     }
 }

compiler/testData/diagnostics/tests/thisAndSuper/unqualifiedSuper/unqualifiedSuperWithLocalClass.fir.kt

@@ -1,21 +0,0 @@
-interface Interface {
-    fun foo(x: Int): Int
-}
-
-fun withLocalClasses(param: Int): Interface {
-    open class LocalBase {
-        open val param: Int
-            get() = 100
-    }
-
-    <!LOCAL_INTERFACE_NOT_ALLOWED!>interface LocalInterface : Interface {
-        override fun foo(x: Int): Int =
-                x + param
-    }<!>
-
-    return object : LocalBase(), LocalInterface {
-        override fun foo(x: Int): Int =
-                x + super.param
-    }
-
-}

compiler/testData/diagnostics/tests/thisAndSuper/unqualifiedSuper/unqualifiedSuperWithLocalClass.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 interface Interface {
     fun foo(x: Int): Int
 }