From 443d18fa69f8752d335f7b85ec4e37de594579df Mon Sep 17 00:00:00 2001
From: Jinseong Jeon <jsjeon@google.com>
Date: Fri, 2 May 2025 12:37:11 -0700
Subject: [PATCH] K2 UAST: avoid module lookup for fake PSI

^KTIJ-34040 fixed

GitOrigin-RevId: 6429cb31ea559b6ddd7ea3b292fcd7be80bab4a7
---
 .../kotlin/UastResolveApiFixtureTestBase.kt   | 20 +++++++++++++------
 .../FirKotlinUastResolveProviderService.kt    |  7 +++++++
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/plugins/kotlin/uast/uast-kotlin-base/tests/test/org/jetbrains/uast/test/common/kotlin/UastResolveApiFixtureTestBase.kt b/plugins/kotlin/uast/uast-kotlin-base/tests/test/org/jetbrains/uast/test/common/kotlin/UastResolveApiFixtureTestBase.kt
index 0cd1557ce914..f6e3dc075ccc 100644
--- a/plugins/kotlin/uast/uast-kotlin-base/tests/test/org/jetbrains/uast/test/common/kotlin/UastResolveApiFixtureTestBase.kt
+++ b/plugins/kotlin/uast/uast-kotlin-base/tests/test/org/jetbrains/uast/test/common/kotlin/UastResolveApiFixtureTestBase.kt
@@ -2623,16 +2623,16 @@ interface UastResolveApiFixtureTestBase {
             "MyStringJVM.kt", """
                 @file:kotlin.jvm.JvmMultifileClass
                 @file:kotlin.jvm.JvmName("MyStringsKt")
-                
+
                 package test.pkg
-                
+
                 annotation class MyAnnotation(
                   val myAttr: String = "defaultValue",
                 )
-                
+
                 @MyAnnotation
                 inline fun <T> T.belongsToClassPart(): String = TODO()
-                
+
                 @MyAnnotation("myAttrValue")
                 inline fun <reified T : Any> T.needFake(): String = TODO()
             """.trimIndent()
@@ -2640,8 +2640,9 @@ interface UastResolveApiFixtureTestBase {
         myFixture.configureByText(
             "main.kt", """
                 import java.util.function.Consumer
-                import test.pkg.*
-                
+                import test.pkg.belongsToClassPart
+                import test.pkg.needFake
+
                 fun test() {
                   Any().belongsToClassPart()
                   Any().needFake()
@@ -2653,6 +2654,13 @@ interface UastResolveApiFixtureTestBase {
         try {
             val uFile = myFixture.file.toUElementOfType<UFile>()!!
             uFile.accept(object : AbstractUastVisitor() {
+                override fun visitImportStatement(node: UImportStatement): Boolean {
+                    val txt = node.sourcePsi?.text
+                    val resolved = node.resolve()
+                    TestCase.assertNotNull(txt, resolved)
+                    return super.visitImportStatement(node)
+                }
+
                 override fun visitCallExpression(node: UCallExpression): Boolean {
                     if (node.isConstructorCall()) {
                         // Like Any()
diff --git a/plugins/kotlin/uast/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/FirKotlinUastResolveProviderService.kt b/plugins/kotlin/uast/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/FirKotlinUastResolveProviderService.kt
index d5a65d150a7f..4d3df6ac4f5c 100644
--- a/plugins/kotlin/uast/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/FirKotlinUastResolveProviderService.kt
+++ b/plugins/kotlin/uast/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/FirKotlinUastResolveProviderService.kt
@@ -35,6 +35,7 @@ import org.jetbrains.uast.*
 import org.jetbrains.uast.analysis.KotlinExtensionConstants.LAMBDA_THIS_PARAMETER_NAME
 import org.jetbrains.uast.kotlin.internal.*
 import org.jetbrains.uast.kotlin.psi.UastFakeDeserializedSymbolAnnotation
+import org.jetbrains.uast.kotlin.psi.UastFakeLightMethodBase
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiParameterBase
 
 interface FirKotlinUastResolveProviderService : BaseKotlinUastResolveProviderService {
@@ -493,6 +494,12 @@ interface FirKotlinUastResolveProviderService : BaseKotlinUastResolveProviderSer
                 return resolvedTargetElement
             }
 
+            // If the resolution result is "fake" PSI, it doesn't belong to any module.
+            // Before falling to the following module lookup, bail out here.
+            if (resolvedTargetElement is UastFakeLightMethodBase) {
+                return resolvedTargetElement
+            }
+
             if (resolvedTargetElement != null) {
                 when (KaModuleProvider.getModule(project, resolvedTargetElement, useSiteModule = null)) {
                     is KaSourceModule -> {