From 50278058aec5dc84fd2301ad5b7ebcff7592fa8b Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Thu, 17 Jan 2019 14:47:47 +0300
Subject: [PATCH] New J2K: Fix nullability calculation

---
 .../j2k/conversions/TypeMappingConversion.kt  | 23 +++++++++++++++----
 .../org/jetbrains/kotlin/j2k/tree/types.kt    | 17 +++++++++++++-
 2 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/j2k/newSrc/org/jetbrains/kotlin/j2k/conversions/TypeMappingConversion.kt b/j2k/newSrc/org/jetbrains/kotlin/j2k/conversions/TypeMappingConversion.kt
index 95647108307..95cb207d4b8 100644
--- a/j2k/newSrc/org/jetbrains/kotlin/j2k/conversions/TypeMappingConversion.kt
+++ b/j2k/newSrc/org/jetbrains/kotlin/j2k/conversions/TypeMappingConversion.kt
@@ -116,6 +116,7 @@ class TypeMappingConversion(val context: ConversionContext) : RecursiveApplicabl
         }
 
     private fun JKClassSymbol.mapClassSymbol(typeElement: JKTypeElement?): JKClassSymbol {
+        if (this is JKUniverseClassSymbol) return this
         val newFqName = typeElement?.let { kotlinCollectionClassName(it) }
             ?: kotlinStandardType()
             ?: fqName
@@ -131,7 +132,7 @@ class TypeMappingConversion(val context: ConversionContext) : RecursiveApplicabl
         )
 
 
-    private fun JKClassSymbol.kotlinCollectionClassName(typeElement: JKTypeElement): String? {
+    private fun JKClassSymbol.kotlinCollectionClassName(typeElement: JKTypeElement?): String? {
         val isStructureMutable = calculateStructureMutability(typeElement)
         return if (isStructureMutable) toKotlinMutableTypesMap[fqName]
         else toKotlinTypesMap[fqName]
@@ -152,15 +153,27 @@ class TypeMappingConversion(val context: ConversionContext) : RecursiveApplicabl
 
     private fun calculateNullability(typeElement: JKTypeElement?): Nullability {
         val parent = typeElement?.parent ?: return Nullability.Default
-        val psi = parent.psi ?: return Nullability.Default
+        val psi = parent.psi
         return when (parent) {
-            is JKJavaMethod -> typeFlavorCalculator.methodNullability(psi as PsiMethod)
-            is JKJavaField -> typeFlavorCalculator.variableNullability(psi as PsiVariable)
-            is JKVariable -> typeFlavorCalculator.variableNullability(psi as PsiVariable)
+            is JKMethod ->
+                parent.nullabilityBySuperMethod(context.symbolProvider).defaultToNull()
+                    ?: psi?.let { typeFlavorCalculator.methodNullability(it as PsiMethod) }
+                        .nullToDefault()
+
+            is JKVariable -> psi?.let {
+                typeFlavorCalculator.variableNullability(psi as PsiVariable)
+            }.nullToDefault()
+
             else -> Nullability.Default
         }
     }
 
+    private fun Nullability.defaultToNull() =
+        if (this == Nullability.Default) null else this
+
+    private fun Nullability?.nullToDefault() =
+        this ?: Nullability.Default
+
     private fun calculateStructureMutability(typeElement: JKTypeElement?): Boolean {
         val parent = typeElement?.parent ?: return false
         val psi = parent.psi ?: return false
diff --git a/j2k/newSrc/org/jetbrains/kotlin/j2k/tree/types.kt b/j2k/newSrc/org/jetbrains/kotlin/j2k/tree/types.kt
index da1f7b2bd31..7b25b6f9a82 100644
--- a/j2k/newSrc/org/jetbrains/kotlin/j2k/tree/types.kt
+++ b/j2k/newSrc/org/jetbrains/kotlin/j2k/tree/types.kt
@@ -315,4 +315,19 @@ fun JKClassSymbol.isInterface(): Boolean {
 }
 
 fun JKType.isInterface(): Boolean =
-    (this as? JKClassType)?.classReference?.isInterface() ?: false
\ No newline at end of file
+    (this as? JKClassType)?.classReference?.isInterface() ?: false
+
+fun JKMethod.nullabilityBySuperMethod(symbolProvider: JKSymbolProvider): Nullability {
+    if (modality != Modality.OVERRIDE) return Nullability.Default
+    val superMethodSymbol = findSuperMethodSymbol(symbolProvider) ?: return Nullability.Default
+    return superMethodSymbol.returnType?.nullability ?: return Nullability.Default
+}
+
+
+private fun JKMethod.findSuperMethodSymbol(symbolProvider: JKSymbolProvider): JKMethodSymbol? =
+    psi<PsiMethod>()
+        ?.findSuperMethods()
+        ?.firstOrNull()
+        ?.let {
+            symbolProvider.provideDirectSymbol(it) as? JKMethodSymbol
+        }
\ No newline at end of file