Ensure that UElement.nameIdentifier.containingFile is a real file

This is needed to report inspection problems on UAST elements.

compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightIdentifier.kt

@@ -26,7 +26,7 @@ import org.jetbrains.kotlin.psi.KtPrimaryConstructor
 import org.jetbrains.kotlin.psi.KtSecondaryConstructor
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 
-class KtLightIdentifier(
+open class KtLightIdentifier(
         private val lightOwner: PsiNameIdentifierOwner,
         private val ktDeclaration: KtNamedDeclaration?
 ) : LightIdentifier(lightOwner.manager, ktDeclaration?.name ?: ""), PsiCompiledElement {

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUClass.kt

@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.psi.KtObjectDeclaration
 import org.jetbrains.uast.*
 import org.jetbrains.uast.java.AbstractJavaUClass
 import org.jetbrains.uast.kotlin.declarations.KotlinUMethod
+import org.jetbrains.uast.kotlin.declarations.UastLightIdentifier
 
 class KotlinUClass private constructor(
         psi: KtLightClass,
@@ -40,6 +41,8 @@ class KotlinUClass private constructor(
         return super.getOriginalElement()
     }
 
+    override fun getNameIdentifier() = UastLightIdentifier(psi, ktClass)
+
     override fun getContainingFile(): PsiFile? = ktClass?.containingFile ?: psi.containingFile
 
     override val annotations: List<UAnnotation>

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUMethod.kt

@@ -43,6 +43,8 @@ open class KotlinUMethod(
 
     override fun getContainingFile(): PsiFile? = kotlinOrigin?.containingFile ?: psi.containingFile
 
+    override fun getNameIdentifier() = UastLightIdentifier(psi, kotlinOrigin as KtNamedDeclaration?)
+
     override val annotations by lz {
         (kotlinOrigin as? KtDeclaration)?.annotationEntries?.map { KotlinUAnnotation(it, this) } ?: emptyList()
     }

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUVariable.kt

@@ -18,6 +18,7 @@ package org.jetbrains.uast.kotlin
 
 import com.intellij.psi.*
 import org.jetbrains.kotlin.asJava.elements.KtLightElement
+import org.jetbrains.kotlin.psi.KtNamedDeclaration
 import org.jetbrains.kotlin.psi.KtParameter
 import org.jetbrains.kotlin.psi.KtVariableDeclaration
 import org.jetbrains.uast.*
@@ -25,6 +26,7 @@ import org.jetbrains.uast.java.AbstractJavaUVariable
 import org.jetbrains.uast.java.JavaAbstractUExpression
 import org.jetbrains.uast.java.JavaUAnnotation
 import org.jetbrains.uast.java.annotations
+import org.jetbrains.uast.kotlin.declarations.UastLightIdentifier
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiParameter
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiVariable
 
@@ -47,6 +49,11 @@ abstract class AbstractKotlinUVariable : AbstractJavaUVariable() {
             } ?: return null
             return getLanguagePlugin().convertElement(initializerExpression, this) as? UExpression ?: UastEmptyExpression
         }
+
+    override fun getNameIdentifier(): PsiIdentifier {
+        val kotlinOrigin = (psi as? KtLightElement<*, *>)?.kotlinOrigin
+        return UastLightIdentifier(psi, kotlinOrigin as KtNamedDeclaration?)
+    }
 }
 
 class KotlinUVariable(
@@ -68,6 +75,10 @@ class KotlinUVariable(
         return super<AbstractKotlinUVariable>.getOriginalElement()
     }
 
+    override fun getNameIdentifier(): PsiIdentifier {
+        return super.getNameIdentifier()
+    }
+
     companion object {
         fun create(psi: PsiVariable, containingElement: UElement?): UVariable {
             return when (psi) {
@@ -95,6 +106,10 @@ open class KotlinUParameter(
     override fun getOriginalElement(): PsiElement? {
         return super<AbstractKotlinUVariable>.getOriginalElement()
     }
+
+    override fun getNameIdentifier(): PsiIdentifier {
+        return super.getNameIdentifier()
+    }
 }
 
 open class KotlinUField(
@@ -111,6 +126,10 @@ open class KotlinUField(
     override fun getOriginalElement(): PsiElement? {
         return super<AbstractKotlinUVariable>.getOriginalElement()
     }
+
+    override fun getNameIdentifier(): PsiIdentifier {
+        return super.getNameIdentifier()
+    }
 }
 
 open class KotlinULocalVariable(
@@ -128,6 +147,9 @@ open class KotlinULocalVariable(
         return super<AbstractKotlinUVariable>.getOriginalElement()
     }
 
+    override fun getNameIdentifier(): PsiIdentifier {
+        return super.getNameIdentifier()
+    }
 }
 
 open class KotlinUEnumConstant(
@@ -146,6 +168,10 @@ open class KotlinUEnumConstant(
         return super<AbstractKotlinUVariable>.getOriginalElement()
     }
 
+    override fun getNameIdentifier(): PsiIdentifier {
+        return super.getNameIdentifier()
+    }
+
     override val kind: UastCallKind
         get() = UastCallKind.CONSTRUCTOR_CALL
     override val receiver: UExpression?

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/UastLightIdentifier.kt

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2010-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.uast.kotlin.declarations
+
+import com.intellij.psi.PsiFile
+import com.intellij.psi.PsiNameIdentifierOwner
+import org.jetbrains.kotlin.asJava.elements.KtLightIdentifier
+import org.jetbrains.kotlin.psi.KtNamedDeclaration
+
+class UastLightIdentifier(lightOwner: PsiNameIdentifierOwner, ktDeclaration: KtNamedDeclaration?)
+    : KtLightIdentifier(lightOwner, ktDeclaration) {
+
+    override fun getContainingFile(): PsiFile {
+        return origin?.containingFile ?: super.getContainingFile()
+    }
+}

plugins/uast-kotlin/testData/NameContainingFile.kt

@@ -0,0 +1,6 @@
+class Foo
+
+fun bar() {}
+
+val xyzzy: Int = 0
+

plugins/uast-kotlin/tests/KotlinUastApiTest.kt

@@ -40,4 +40,17 @@ class KotlinUastApiTest : AbstractKotlinUastTest() {
             assertNull(uLiteral)
         }
     }
+
+    @Test fun testNameContainingFile() {
+        doTest("NameContainingFile") { _, file ->
+            val foo = file.findElementByText<UClass>("class Foo")
+            assertEquals(file.psi, foo.nameIdentifier!!.containingFile)
+
+            val bar = file.findElementByText<UMethod>("fun bar() {}")
+            assertEquals(file.psi, bar.nameIdentifier!!.containingFile)
+
+            val xyzzy = file.findElementByText<UVariable>("val xyzzy: Int = 0")
+            assertEquals(file.psi, xyzzy.nameIdentifier!!.containingFile)
+        }
+    }
 }