From a5e304f52057b3df9cbc6e424ecf4b6c148b6b6b Mon Sep 17 00:00:00 2001
From: Igor Yakovlev <igor.yakovlev@jetbrains.com>
Date: Tue, 23 Jun 2020 16:25:12 +0300
Subject: [PATCH] [ULC] Better nullability annotations for light classes

---
 .../asJava/elements/KtLightModifierList.kt    | 15 +++++---
 .../asJava/elements/lightAnnotations.kt       | 13 ++++---
 .../kotlin/asJava/lightClassUtils.kt          | 34 +++++++++++++++++--
 3 files changed, 49 insertions(+), 13 deletions(-)

diff --git a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightModifierList.kt b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightModifierList.kt
index 32295dd1ed2..2378ae47b6e 100644
--- a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightModifierList.kt
+++ b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightModifierList.kt
@@ -10,6 +10,7 @@ import com.intellij.util.IncorrectOperationException
 import org.jetbrains.kotlin.asJava.LightClassGenerationSupport
 import org.jetbrains.kotlin.asJava.builder.LightMemberOriginForDeclaration
 import org.jetbrains.kotlin.asJava.classes.*
+import org.jetbrains.kotlin.asJava.fastCheckIsNullabilityApplied
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.PropertyDescriptor
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor
@@ -76,12 +77,16 @@ abstract class KtLightModifierList<out T : KtLightElement<KtModifierListOwner, P
         if ((modifierListOwner is KtLightMember<*> && modifierListOwner !is KtLightFieldImpl.KtLightEnumConstant)
             || modifierListOwner is LightParameter
         ) {
-            val nullabilityAnnotation = when (modifierListOwner) {
-                is KtUltraLightElementWithNullabilityAnnotation<*, *> -> KtUltraLightNullabilityAnnotation(modifierListOwner, this)
-                else -> KtLightNullabilityAnnotation(modifierListOwner as KtLightElement<*, PsiModifierListOwner>, this)
-            }
 
-            return annotationsForEntries + listOf(nullabilityAnnotation)
+            if (fastCheckIsNullabilityApplied(modifierListOwner)) {
+
+                val nullabilityAnnotation = when (modifierListOwner) {
+                    is KtUltraLightElementWithNullabilityAnnotation<*, *> -> KtUltraLightNullabilityAnnotation(modifierListOwner, this)
+                    else -> KtLightNullabilityAnnotation(modifierListOwner as KtLightElement<*, PsiModifierListOwner>, this)
+                }
+
+                return annotationsForEntries + listOf(nullabilityAnnotation)
+            }
         }
         return annotationsForEntries
     }
diff --git a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/lightAnnotations.kt b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/lightAnnotations.kt
index a21127a3efa..582338deec5 100644
--- a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/lightAnnotations.kt
+++ b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/lightAnnotations.kt
@@ -26,6 +26,7 @@ import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.asJava.LightClassGenerationSupport
 import org.jetbrains.kotlin.asJava.classes.cannotModify
 import org.jetbrains.kotlin.asJava.classes.lazyPub
+import org.jetbrains.kotlin.asJava.fastCheckIsNullabilityApplied
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.CallableMemberDescriptor
@@ -37,6 +38,7 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.hasSuspendModifier
+import org.jetbrains.kotlin.psi.psiUtil.isPrivate
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.CompileTimeConstantUtils
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
@@ -308,14 +310,15 @@ open class KtLightNullabilityAnnotation<D : KtLightElement<*, PsiModifierListOwn
             ?: // it is out of our hands
             return getClsNullabilityAnnotation(member)?.qualifiedName
 
+        if (!fastCheckIsNullabilityApplied(member)) return null
+
         // all data-class generated members are not-null
         if (annotatedElement is KtClass && annotatedElement.isData()) return NotNull::class.java.name
 
-        // objects and companion objects are not null
-        if (annotatedElement is KtObjectDeclaration) return NotNull::class.java.name
-
-        if (annotatedElement is KtParameter) {
-            if (annotatedElement.containingClassOrObject?.isAnnotation() == true) return null
+        // objects and companion objects are not null but ctor that does not annotated
+        if (annotatedElement is KtObjectDeclaration) {
+            if ((parent.parent as? PsiMethod)?.isConstructor == true) return null
+            return NotNull::class.java.name
         }
 
         // don't annotate property setters
diff --git a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/lightClassUtils.kt b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/lightClassUtils.kt
index 35e68f33e9e..abd7058deaf 100644
--- a/compiler/light-classes/src/org/jetbrains/kotlin/asJava/lightClassUtils.kt
+++ b/compiler/light-classes/src/org/jetbrains/kotlin/asJava/lightClassUtils.kt
@@ -20,6 +20,7 @@ import com.intellij.psi.*
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.asJava.classes.KtLightClass
 import org.jetbrains.kotlin.asJava.classes.KtLightClassForFacade
+import org.jetbrains.kotlin.asJava.classes.KtUltraLightElementWithNullabilityAnnotation
 import org.jetbrains.kotlin.asJava.elements.PsiElementWithOrigin
 import org.jetbrains.kotlin.asJava.elements.*
 import org.jetbrains.kotlin.builtins.jvm.JavaToKotlinClassMap
@@ -30,9 +31,7 @@ import org.jetbrains.kotlin.load.java.propertyNameBySetMethodName
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
-import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
-import org.jetbrains.kotlin.psi.psiUtil.isExtensionDeclaration
+import org.jetbrains.kotlin.psi.psiUtil.*
 
 /**
  * Can be null in scripts and for elements from non-jvm modules.
@@ -220,4 +219,33 @@ fun KtLightMethod.checkIsMangled(): Boolean {
     val demangledName = KotlinTypeMapper.InternalNameMapper.demangleInternalName(name) ?: return false
     val originalName = propertyNameByAccessor(demangledName, this) ?: demangledName
     return originalName == kotlinOrigin?.name
+}
+
+fun fastCheckIsNullabilityApplied(lightElement: KtLightElement<*, PsiModifierListOwner>): Boolean {
+
+    val annotatedElement = lightElement.kotlinOrigin ?: return true
+
+    // all data-class generated members are not-null
+    if (annotatedElement is KtClass && annotatedElement.isData()) return true
+
+    if (annotatedElement is KtParameter && lightElement !is PsiField) { //KtParameter but not it's backing field
+        val containingClassOrObject = annotatedElement.containingClassOrObject
+        if (containingClassOrObject?.isAnnotation() == true) return false
+        if ((containingClassOrObject as? KtClass)?.isEnum() == true) {
+            if (annotatedElement.parent.parent is KtPrimaryConstructor) return false
+        }
+
+        val parent = annotatedElement.parent.parent
+        if (parent is KtModifierListOwner && parent.isPrivate()) return false
+
+        if (parent is KtPropertyAccessor) {
+            val propertyOfAccessor = parent.parent
+            if (propertyOfAccessor is KtProperty && propertyOfAccessor.isPrivate()) return false
+        }
+    } else {
+        // private properties
+        if ((annotatedElement as? KtModifierListOwner)?.isPrivate() == true) return false
+    }
+
+    return true
 }
\ No newline at end of file