From aa53489e0b002854d13d9e14ec046548ecf38791 Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Mon, 8 Apr 2019 18:41:11 +0300
Subject: [PATCH] New J2K: Handle forced nullability correctly in nullability
 analysis

---
 .../kotlin/nj2k/nullabilityAnalysis/constraints.kt    |  2 +-
 .../kotlin/nj2k/nullabilityAnalysis/context.kt        |  6 +++++-
 .../kotlin/nj2k/nullabilityAnalysis/printing.kt       |  4 ++--
 .../testData/nullabilityAnalysis/forcedNullability.kt | 11 +++++++++++
 .../nullabilityAnalysis/forcedNullability.kt.after    | 11 +++++++++++
 .../kotlin/nj2k/NullabilityAnalysisTestGenerated.java |  5 +++++
 6 files changed, 35 insertions(+), 4 deletions(-)
 create mode 100644 nj2k/testData/nullabilityAnalysis/forcedNullability.kt
 create mode 100644 nj2k/testData/nullabilityAnalysis/forcedNullability.kt.after

diff --git a/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/constraints.kt b/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/constraints.kt
index bfbfa480d48..592efb4c1bc 100644
--- a/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/constraints.kt
+++ b/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/constraints.kt
@@ -130,7 +130,7 @@ internal inline val BoundType.bound
         }
 
 internal inline fun <reified T : BoundType> T.withForcedNullability(nullability: Nullability?): T =
-    if (forcedNullabilityTo == nullability) this
+    if (forcedNullabilityTo == nullability || nullability == null) this
     else when (this) {
         is GenericBoundType ->
             GenericBoundType(
diff --git a/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/context.kt b/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/context.kt
index 5223073b77c..213ee7c748a 100644
--- a/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/context.kt
+++ b/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/context.kt
@@ -113,6 +113,10 @@ internal data class AnalysisContext(
     val declarationToTypeVariable: Map<KtCallableDeclaration, TypeVariable>
 )
 
-data class AnalysisScope(val elements: List<KtElement>) : Iterable<KtElement> by elements {
+data class AnalysisScope(val elements: List<PsiElement>) : Iterable<PsiElement> by elements {
     constructor(vararg elements: KtElement) : this(elements.toList())
+    constructor(file: KtFile, rangeMarker: RangeMarker?) :
+            this(rangeMarker?.let { marker ->
+                file.elementsInRange(TextRange(marker.startOffset, marker.endOffset))
+            } ?: listOf(file))
 }
diff --git a/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/printing.kt b/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/printing.kt
index 77490d86db2..15ab3cdddc2 100644
--- a/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/printing.kt
+++ b/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/printing.kt
@@ -5,7 +5,7 @@
 
 package org.jetbrains.kotlin.nj2k.nullabilityAnalysis
 
-import org.jetbrains.kotlin.psi.KtElement
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.jetbrains.kotlin.psi.KtTypeElement
 import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
@@ -55,7 +55,7 @@ internal class Printer(private val analysisContext: AnalysisContext) {
         }
 
 
-    internal fun KtElement.addTypeVariablesNames() {
+    internal fun PsiElement.addTypeVariablesNames() {
         val factory = KtPsiFactory(this)
         for (typeElement in collectDescendantsOfType<KtTypeElement>()) {
             val typeVariableName = this@Printer.analysisContext.typeElementToTypeVariable[typeElement]?.name ?: continue
diff --git a/nj2k/testData/nullabilityAnalysis/forcedNullability.kt b/nj2k/testData/nullabilityAnalysis/forcedNullability.kt
new file mode 100644
index 00000000000..fefd4eae7a2
--- /dev/null
+++ b/nj2k/testData/nullabilityAnalysis/forcedNullability.kt
@@ -0,0 +1,11 @@
+fun a(m: Map<Int, String>): String {
+    return m.get(42)
+}
+
+fun b(m: Map<Int, String>): String {
+    return m.get(42)
+}
+
+fun c() {
+    a(mapOf<Int, String>(q.to("nya")))
+}
\ No newline at end of file
diff --git a/nj2k/testData/nullabilityAnalysis/forcedNullability.kt.after b/nj2k/testData/nullabilityAnalysis/forcedNullability.kt.after
new file mode 100644
index 00000000000..262ee0a29ce
--- /dev/null
+++ b/nj2k/testData/nullabilityAnalysis/forcedNullability.kt.after
@@ -0,0 +1,11 @@
+fun a(m: Map<Int, String>): String? {
+    return m.get(42)
+}
+
+fun b(m: Map<Int, String?>): String? {
+    return m.get(42)
+}
+
+fun c() {
+    a(mapOf<Int, String>(q.to("nya")))
+}
\ No newline at end of file
diff --git a/nj2k/tests/org/jetbrains/kotlin/nj2k/NullabilityAnalysisTestGenerated.java b/nj2k/tests/org/jetbrains/kotlin/nj2k/NullabilityAnalysisTestGenerated.java
index 2fc813a902e..fd84d412ecb 100644
--- a/nj2k/tests/org/jetbrains/kotlin/nj2k/NullabilityAnalysisTestGenerated.java
+++ b/nj2k/tests/org/jetbrains/kotlin/nj2k/NullabilityAnalysisTestGenerated.java
@@ -34,6 +34,11 @@ public class NullabilityAnalysisTestGenerated extends AbstractNullabilityAnalysi
         runTest("nj2k/testData/nullabilityAnalysis/compareWithNull.kt");
     }
 
+    @TestMetadata("forcedNullability.kt")
+    public void testForcedNullability() throws Exception {
+        runTest("nj2k/testData/nullabilityAnalysis/forcedNullability.kt");
+    }
+
     @TestMetadata("functionTypeParameterNullability.kt")
     public void testFunctionTypeParameterNullability() throws Exception {
         runTest("nj2k/testData/nullabilityAnalysis/functionTypeParameterNullability.kt");