UAST: Fix getContainingFile, should return KtFile instead of FakeFileForLightClass

#KT-20170 Fixed

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

@@ -29,7 +29,7 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtFile
 
 open class FakeFileForLightClass(
-        protected val ktFile: KtFile,
+        val ktFile: KtFile,
         private val lightClass: () -> KtLightClass,
         private val stub: () -> PsiClassHolderFileStub<*>,
         private val packageFqName: FqName = ktFile.packageFqName

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

@@ -39,15 +39,15 @@ open class KotlinUClass private constructor(
 
     override fun getOriginalElement(): PsiElement? = super.getOriginalElement()
 
-    override fun getNameIdentifier() = UastLightIdentifier(psi, ktClass)
+    override fun getNameIdentifier(): PsiIdentifier? = UastLightIdentifier(psi, ktClass)
 
-    override fun getContainingFile(): PsiFile? = ktClass?.containingFile ?: psi.containingFile
+    override fun getContainingFile(): PsiFile? = unwrapFakeFileForLightClass(psi.containingFile)
 
     override val annotations: List<UAnnotation>
         get() = ktClass?.annotationEntries?.map { KotlinUAnnotation(it, this) } ?: emptyList()
 
     override val uastAnchor: UElement
-        get() = UIdentifier(psi.nameIdentifier, this)
+        get() = UIdentifier(nameIdentifier, this)
 
     override fun getInnerClasses(): Array<UClass> {
         // filter DefaultImpls to avoid processing same methods from original interface multiple times
@@ -123,6 +123,8 @@ class KotlinUAnonymousClass(
     override fun getInitializers(): Array<UClassInitializer> = super<AbstractJavaUClass>.getInitializers()
     override fun getInnerClasses(): Array<UClass> = super<AbstractJavaUClass>.getInnerClasses()
 
+    override fun getContainingFile(): PsiFile = unwrapFakeFileForLightClass(psi.containingFile)
+
     override val uastAnchor: UElement?
         get() {
             val ktClassOrObject = (psi.originalElement as? KtLightClass)?.kotlinOrigin as? KtObjectDeclaration ?: return null 
@@ -134,6 +136,12 @@ class KotlinScriptUClass(
         psi: KtLightClassForScript,
         override val uastParent: UElement?
 ) : AbstractJavaUClass(), PsiClass by psi {
+    override fun getContainingFile(): PsiFile = unwrapFakeFileForLightClass(psi.containingFile)
+
+    override fun getNameIdentifier(): PsiIdentifier? = UastLightIdentifier(psi, psi.kotlinOrigin)
+
+    override val uastAnchor: UElement
+        get() = UIdentifier(nameIdentifier, this)
 
     override val psi = unwrap<UClass, KtLightClassForScript>(psi)
 

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

@@ -42,7 +42,7 @@ open class KotlinUMethod(
 
     private val kotlinOrigin = (psi.originalElement as KtLightElement<*, *>).kotlinOrigin
 
-    override fun getContainingFile(): PsiFile? = kotlinOrigin?.containingFile ?: psi.containingFile
+    override fun getContainingFile(): PsiFile? = unwrapFakeFileForLightClass(psi.containingFile)
 
     override fun getNameIdentifier() = UastLightIdentifier(psi, kotlinOrigin as KtNamedDeclaration?)
 
@@ -57,7 +57,7 @@ open class KotlinUMethod(
     }
 
     override val uastAnchor: UElement
-        get() = UIdentifier((psi.originalElement as? PsiNameIdentifierOwner)?.nameIdentifier ?: psi.nameIdentifier, this)
+        get() = UIdentifier(nameIdentifier, this)
 
 
     override val uastBody by lz {

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

@@ -69,12 +69,14 @@ abstract class AbstractKotlinUVariable : PsiVariable, UVariable, KotlinUElementW
         return UastLightIdentifier(psi, kotlinOrigin as KtNamedDeclaration?)
     }
 
+    override fun getContainingFile(): PsiFile = unwrapFakeFileForLightClass(psi.containingFile)
+
     override val annotations: List<UAnnotation> by lz { psi.annotations.map { JavaUAnnotation(it, this) } }
 
     override val typeReference by lz { getLanguagePlugin().convertOpt<UTypeReferenceExpression>(psi.typeElement, this) }
 
     override val uastAnchor: UElement?
-        get() = UIdentifier(psi.nameIdentifier, this)
+        get() = UIdentifier(nameIdentifier, this)
 
     override fun equals(other: Any?) = other is AbstractKotlinUVariable && psi == other.psi
 
@@ -91,8 +93,6 @@ class KotlinUVariable(
 
     override val typeReference by lz { getLanguagePlugin().convertOpt<UTypeReferenceExpression>(psi.typeElement, this) }
 
-    override fun getContainingFile(): PsiFile? = (psi as? KtLightElement<*, *>)?.kotlinOrigin?.containingFile ?: psi.containingFile
-
     override fun getInitializer(): PsiExpression? {
         return super<AbstractKotlinUVariable>.getInitializer()
     }
@@ -104,6 +104,10 @@ class KotlinUVariable(
     override fun getNameIdentifier(): PsiIdentifier {
         return super.getNameIdentifier()
     }
+
+    override fun getContainingFile(): PsiFile {
+        return super.getContainingFile()
+    }
 }
 
 open class KotlinUParameter(
@@ -124,6 +128,10 @@ open class KotlinUParameter(
     override fun getNameIdentifier(): PsiIdentifier {
         return super.getNameIdentifier()
     }
+
+    override fun getContainingFile(): PsiFile {
+        return super.getContainingFile()
+    }
 }
 
 open class KotlinUField(
@@ -145,12 +153,14 @@ open class KotlinUField(
         return super.getNameIdentifier()
     }
 
+    override fun getContainingFile(): PsiFile {
+        return super.getContainingFile()
+    }
+
     override fun isPhysical(): Boolean {
         return true
     }
 
-    override fun getContainingFile(): PsiFile? = (psi as? KtLightElement<*, *>)?.kotlinOrigin?.containingFile ?: psi.containingFile
-
     override fun accept(visitor: UastVisitor) {
         if (visitor.visitField(this)) return
         annotations.acceptList(visitor)
@@ -179,6 +189,10 @@ open class KotlinULocalVariable(
         return super.getNameIdentifier()
     }
 
+    override fun getContainingFile(): PsiFile {
+        return super.getContainingFile()
+    }
+
     override fun accept(visitor: UastVisitor) {
         if (visitor.visitLocalVariable(this)) return
         annotations.acceptList(visitor)
@@ -201,36 +215,40 @@ open class KotlinUEnumConstant(
         psi: PsiEnumConstant,
         override val uastParent: UElement?
 ) : AbstractKotlinUVariable(), UEnumConstant, PsiEnumConstant by psi {
-    override val initializingClass: UClass? by lz { getLanguagePlugin().convertOpt<UClass>(psi.initializingClass, this) }
 
-    override val psi = unwrap<UEnumConstant, PsiEnumConstant>(psi)
+    override fun getContainingFile(): PsiFile {
-
+        return super.getContainingFile()
-    override fun getInitializer(): PsiExpression? {
-        return super<AbstractKotlinUVariable>.getInitializer()
-    }
-
-    override fun getOriginalElement(): PsiElement? {
-        return super<AbstractKotlinUVariable>.getOriginalElement()
     }
 
     override fun getNameIdentifier(): PsiIdentifier {
         return super.getNameIdentifier()
     }
 
+    override val initializingClass: UClass? by lz { getLanguagePlugin().convertOpt<UClass>(psi.initializingClass, this) }
+
+    override val psi = unwrap<UEnumConstant, PsiEnumConstant>(psi)
+
     override val kind: UastCallKind
         get() = UastCallKind.CONSTRUCTOR_CALL
+
     override val receiver: UExpression?
         get() = null
+
     override val receiverType: PsiType?
         get() = null
+
     override val methodIdentifier: UIdentifier?
         get() = null
+
     override val classReference: UReferenceExpression?
         get() = KotlinEnumConstantClassReference(psi, this)
+
     override val typeArgumentCount: Int
         get() = 0
+
     override val typeArguments: List<PsiType>
         get() = emptyList()
+
     override val valueArgumentCount: Int
         get() = psi.argumentList?.expressions?.size ?: 0
 

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

@@ -20,11 +20,9 @@ import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiNameIdentifierOwner
 import org.jetbrains.kotlin.asJava.elements.KtLightIdentifier
 import org.jetbrains.kotlin.psi.KtNamedDeclaration
+import org.jetbrains.uast.kotlin.unwrapFakeFileForLightClass
 
 class UastLightIdentifier(lightOwner: PsiNameIdentifierOwner, ktDeclaration: KtNamedDeclaration?)
     : KtLightIdentifier(lightOwner, ktDeclaration) {
-
+    override fun getContainingFile(): PsiFile = unwrapFakeFileForLightClass(super.getContainingFile())
-    override fun getContainingFile(): PsiFile {
-        return origin?.containingFile ?: super.getContainingFile()
-    }
 }

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/internal/kotlinInternalUastUtils.kt

@@ -20,6 +20,7 @@ import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.util.Key
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiPrimitiveType
 import com.intellij.psi.PsiType
 import com.intellij.psi.impl.cache.TypeInfo
@@ -28,6 +29,7 @@ import com.intellij.psi.impl.compiled.SignatureParsing
 import com.intellij.psi.impl.compiled.StubBuildingVisitor
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.util.PsiTypesUtil
+import org.jetbrains.kotlin.asJava.elements.FakeFileForLightClass
 import org.jetbrains.kotlin.asJava.toLightClass
 import org.jetbrains.kotlin.asJava.toLightElements
 import org.jetbrains.kotlin.builtins.isBuiltinFunctionalTypeOrSubtype
@@ -36,7 +38,6 @@ import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.PropertyDescriptor
-import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.BindingContext
@@ -198,3 +199,5 @@ internal fun KotlinULambdaExpression.getFunctionalInterfaceType(): PsiType? {
         else -> psi.getExpectedType()?.getFunctionalInterfaceType(this, psi)
     }
 }
+
+internal fun unwrapFakeFileForLightClass(file: PsiFile): PsiFile = (file as? FakeFileForLightClass)?.ktFile ?: file

plugins/uast-kotlin/tests/AbstractKotlinRenderLogTest.kt

@@ -1,9 +1,39 @@
 package org.jetbrains.uast.test.kotlin
 
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.utils.addToStdlib.assertedCast
+import org.jetbrains.uast.UDeclaration
+import org.jetbrains.uast.UElement
+import org.jetbrains.uast.UFile
+import org.jetbrains.uast.UIdentifier
 import org.jetbrains.uast.test.common.RenderLogTestBase
+import org.jetbrains.uast.visitor.UastVisitor
 import java.io.File
 
 abstract class AbstractKotlinRenderLogTest : AbstractKotlinUastTest(), RenderLogTestBase {
     override fun getTestFile(testName: String, ext: String) =
             File(File(TEST_KOTLIN_MODEL_DIR, testName).canonicalPath + '.' + ext)
+
+    override fun check(testName: String, file: UFile) {
+        super.check(testName, file)
+        file.checkContainingFileForAllElements()
+    }
+
+    private fun UFile.checkContainingFileForAllElements() {
+        accept(object : UastVisitor {
+            override fun visitElement(node: UElement): Boolean {
+                if (node is PsiElement) {
+                    node.containingFile.assertedCast<KtFile> { "containingFile should be KtFile for ${node.asLogString()}" }
+                }
+
+                val anchorPsi = (node as? UDeclaration)?.uastAnchor?.psi
+                if (anchorPsi != null) {
+                    anchorPsi.containingFile.assertedCast<KtFile> { "uastAnchor.containingFile should be KtFile for ${node.asLogString()}" }
+                }
+
+                return false
+            }
+        })
+    }
 }