From af17db5f45d01498d9ce53544aa67a0f72712169 Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Mon, 8 Apr 2019 18:45:57 +0300
Subject: [PATCH] New J2K: handle more types of nullability annotations

---
 .../JetbrainsNullableAnnotationsConverter.kt  | 60 +++++++++----------
 1 file changed, 30 insertions(+), 30 deletions(-)

diff --git a/nj2k/src/org/jetbrains/kotlin/nj2k/conversions/JetbrainsNullableAnnotationsConverter.kt b/nj2k/src/org/jetbrains/kotlin/nj2k/conversions/JetbrainsNullableAnnotationsConverter.kt
index f83f3d1364c..f72f41afc47 100644
--- a/nj2k/src/org/jetbrains/kotlin/nj2k/conversions/JetbrainsNullableAnnotationsConverter.kt
+++ b/nj2k/src/org/jetbrains/kotlin/nj2k/conversions/JetbrainsNullableAnnotationsConverter.kt
@@ -6,6 +6,8 @@
 package org.jetbrains.kotlin.nj2k.conversions
 
 import org.jetbrains.kotlin.j2k.ast.Nullability
+import org.jetbrains.kotlin.load.java.NOT_NULL_ANNOTATIONS
+import org.jetbrains.kotlin.load.java.NULLABLE_ANNOTATIONS
 import org.jetbrains.kotlin.nj2k.ConversionContext
 import org.jetbrains.kotlin.nj2k.tree.*
 import org.jetbrains.kotlin.nj2k.tree.impl.JKTypeElementImpl
@@ -14,40 +16,38 @@ import org.jetbrains.kotlin.nj2k.tree.impl.JKTypeElementImpl
 class JetbrainsNullableAnnotationsConverter(private val context: ConversionContext) : RecursiveApplicableConversionBase() {
     override fun applyToElement(element: JKTreeElement): JKTreeElement {
         if (element !is JKAnnotationListOwner) return recurse(element)
-        val nullableAnnotationSymbol =
-            context.symbolProvider.provideByFqName("org.jetbrains.annotations.Nullable")
-        val notNullAnnotationSymbol =
-            context.symbolProvider.provideByFqName("org.jetbrains.annotations.NotNull")
-        val nullableAnnotation =
-            element.annotationList.annotations.firstOrNull { it.classSymbol == nullableAnnotationSymbol }
-        val notNullAnnotation =
-            element.annotationList.annotations.firstOrNull { it.classSymbol == notNullAnnotationSymbol }
-        when (element) {
-            is JKVariable -> {
-                if (nullableAnnotation != null) {
-                    element.annotationList.annotations -= nullableAnnotation
-                    element.type =
-                        JKTypeElementImpl(element.type.type.updateNullability(Nullability.Nullable))
+
+        val annotationsToRemove = mutableListOf<JKAnnotation>()
+        for (annotation in element.annotationList.annotations) {
+            val nullability = annotation.annotationNullability() ?: continue
+            when (element) {
+                is JKVariable -> {
+                    annotationsToRemove += annotation
+                    element.type = JKTypeElementImpl(element.type.type.updateNullability(nullability))
                 }
-                if (notNullAnnotation != null) {
-                    element.annotationList.annotations -= notNullAnnotation
-                    element.type =
-                        JKTypeElementImpl(element.type.type.updateNullability(Nullability.NotNull))
-                }
-            }
-            is JKMethod -> {
-                if (nullableAnnotation != null) {
-                    element.annotationList.annotations -= nullableAnnotation
-                    element.returnType =
-                        JKTypeElementImpl(element.returnType.type.updateNullability(Nullability.Nullable))
-                }
-                if (notNullAnnotation != null) {
-                    element.annotationList.annotations -= notNullAnnotation
-                    element.returnType =
-                        JKTypeElementImpl(element.returnType.type.updateNullability(Nullability.NotNull))
+                is JKMethod -> {
+                    annotationsToRemove += annotation
+                    element.returnType = JKTypeElementImpl(element.returnType.type.updateNullability(nullability))
                 }
             }
         }
+        element.annotationList.annotations -= annotationsToRemove
+
         return recurse(element)
     }
+
+    private fun JKAnnotation.annotationNullability(): Nullability? =
+        when {
+            classSymbol.fqName in nullableAnnotationsFqNames -> Nullability.Nullable
+            classSymbol.fqName in notNullAnnotationsFqNames -> Nullability.NotNull
+            else -> null
+        }
+
+
+    companion object {
+        val nullableAnnotationsFqNames =
+            NULLABLE_ANNOTATIONS.map { it.asString() }.toSet()
+        val notNullAnnotationsFqNames =
+            NOT_NULL_ANNOTATIONS.map { it.asString() }.toSet()
+    }
 }
\ No newline at end of file