From f87c9af24d15d840aa55a7cdc7cf38c00f9adfec Mon Sep 17 00:00:00 2001
From: Jinseong Jeon <jsjeon@google.com>
Date: Thu, 27 Jun 2024 14:59:01 -0700
Subject: [PATCH] K2 UAST: make nullity computation compatible with K1 UAST

In particular, for flexible type due to platform type, it's unknown.

^KTIJ-30464 fixed

GitOrigin-RevId: cb7327a4f6291b94d627b3929c2b10b348b060c5
---
 .../kotlin/UastResolveApiFixtureTestBase.kt   | 35 +++++++++++++++++++
 .../internal/firKotlinInternalUastUtils.kt    | 10 +++---
 .../uast/test/FirUastResolveApiFixtureTest.kt |  4 +++
 .../FE1UastResolveApiFixtureTest.kt           |  4 +++
 4 files changed, 49 insertions(+), 4 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 d9959513814e..8a1c7b20a5a8 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
@@ -616,6 +616,41 @@ interface UastResolveApiFixtureTestBase {
         TestCase.assertEquals("it", resolved.name)
     }
 
+    fun checkNullityOfResolvedLambdaParameter(myFixture: JavaCodeInsightTestFixture) {
+        myFixture.addClass(
+            """
+                class MyLiveData<T> {
+                  public void addSource(MyLiveData<S> source, Observer<? super S> onChanged) {}
+                  public void setValue(T value) {}
+                }
+            """.trimIndent()
+        )
+        myFixture.configureByText(
+            "main.kt", """
+                fun interface Observer<T> {
+                  fun onChanged(value: T)
+                }
+                
+                class Test {
+                  val myData = MyLiveData<List<Boolean>>()
+                  
+                  init {
+                    myData.addSource(getSources()) { data ->
+                      myData.value = da<caret>ta
+                    }
+                  }
+                  
+                  private fun getSources(): MyLiveData<List<Boolean>> = TODO()
+                }
+            """.trimIndent()
+        )
+        val ref = myFixture.file.findElementAt(myFixture.caretOffset).toUElement()?.getParentOfType<UReferenceExpression>()
+            .orFail("cant convert to UReferenceExpression")
+        val resolved = (ref.resolve() as? PsiParameter)
+            .orFail("cant resolve lambda parameter")
+        TestCase.assertFalse(resolved.annotations.any { it.isNullnessAnnotation })
+    }
+
     fun checkResolveSyntheticMethod(myFixture: JavaCodeInsightTestFixture) {
         myFixture.configureByText(
             "MyClass.kt", """
diff --git a/plugins/kotlin/uast/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/internal/firKotlinInternalUastUtils.kt b/plugins/kotlin/uast/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/internal/firKotlinInternalUastUtils.kt
index ca95ca0f95d8..9658f68ed861 100644
--- a/plugins/kotlin/uast/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/internal/firKotlinInternalUastUtils.kt
+++ b/plugins/kotlin/uast/uast-kotlin-fir/src/org/jetbrains/uast/kotlin/internal/firKotlinInternalUastUtils.kt
@@ -367,10 +367,12 @@ context(KaSession)
 internal fun nullability(ktType: KaType?): KaTypeNullability? {
     if (ktType == null) return null
     if (ktType is KaErrorType) return null
-    return if (ktType.fullyExpandedType.canBeNull)
-        KaTypeNullability.NULLABLE
-    else
-        KaTypeNullability.NON_NULLABLE
+    val expanded = ktType.fullyExpandedType
+    return when {
+        expanded.hasFlexibleNullability -> KaTypeNullability.UNKNOWN
+        expanded.canBeNull -> KaTypeNullability.NULLABLE
+        else -> KaTypeNullability.NON_NULLABLE
+    }
 }
 
 context(KaSession)
diff --git a/plugins/kotlin/uast/uast-kotlin-fir/tests/test/org/jetbrains/fir/uast/test/FirUastResolveApiFixtureTest.kt b/plugins/kotlin/uast/uast-kotlin-fir/tests/test/org/jetbrains/fir/uast/test/FirUastResolveApiFixtureTest.kt
index ee8bd575e70e..2819df6c5f12 100644
--- a/plugins/kotlin/uast/uast-kotlin-fir/tests/test/org/jetbrains/fir/uast/test/FirUastResolveApiFixtureTest.kt
+++ b/plugins/kotlin/uast/uast-kotlin-fir/tests/test/org/jetbrains/fir/uast/test/FirUastResolveApiFixtureTest.kt
@@ -142,6 +142,10 @@ class FirUastResolveApiFixtureTest : KotlinLightCodeInsightFixtureTestCase(), Ua
         checkResolveImplicitLambdaParameter_binary(myFixture)
     }
 
+    fun testNullityOfResolvedLambdaParameter() {
+        checkNullityOfResolvedLambdaParameter(myFixture)
+    }
+
     fun testResolveSyntheticMethod() {
         checkResolveSyntheticMethod(myFixture)
     }
diff --git a/plugins/kotlin/uast/uast-kotlin/tests/test/org/jetbrains/uast/test/kotlin/comparison/FE1UastResolveApiFixtureTest.kt b/plugins/kotlin/uast/uast-kotlin/tests/test/org/jetbrains/uast/test/kotlin/comparison/FE1UastResolveApiFixtureTest.kt
index 15ec5512c7a1..248d15e7db8a 100644
--- a/plugins/kotlin/uast/uast-kotlin/tests/test/org/jetbrains/uast/test/kotlin/comparison/FE1UastResolveApiFixtureTest.kt
+++ b/plugins/kotlin/uast/uast-kotlin/tests/test/org/jetbrains/uast/test/kotlin/comparison/FE1UastResolveApiFixtureTest.kt
@@ -98,6 +98,10 @@ class FE1UastResolveApiFixtureTest : KotlinLightCodeInsightFixtureTestCase(), Ua
         checkResolveImplicitLambdaParameter_binary(myFixture)
     }
 
+    fun testNullityOfResolvedLambdaParameter() {
+        checkNullityOfResolvedLambdaParameter(myFixture)
+    }
+
     fun testResolveSyntheticMethod() {
         checkResolveSyntheticMethod(myFixture)
     }