Separate nullability annotation processing for ultra-light classes

KtLightNullabilityAnnotation is specifically designed for old light classes
and includes a bunch of hacks that are applicable for them

So, we have to introduce own hack into it
(see org.jetbrains.kotlin.asJava.classes.KtUltraLightParameter#getTypeForNullability)
that is a bit of hard to support when extending hierarchy

compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightAnnotations.kt

@@ -0,0 +1,43 @@
+package org.jetbrains.kotlin.asJava.classes
+
+import com.intellij.lang.jvm.JvmModifier
+import com.intellij.psi.*
+import org.jetbrains.annotations.NotNull
+import org.jetbrains.annotations.Nullable
+import org.jetbrains.kotlin.asJava.elements.KtLightNullabilityAnnotation
+import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.TypeUtils
+import org.jetbrains.kotlin.types.isError
+import org.jetbrains.kotlin.types.typeUtil.TypeNullability
+import org.jetbrains.kotlin.types.typeUtil.isTypeParameter
+import org.jetbrains.kotlin.types.typeUtil.nullability
+
+class KtUltraLightNullabilityAnnotation(
+    member: KtUltraLightElementWithNullabilityAnnotation<*, *>,
+    parent: PsiElement
+) : KtLightNullabilityAnnotation<KtUltraLightElementWithNullabilityAnnotation<*, *>>(member, parent) {
+    override fun getQualifiedName(): String? {
+        val kotlinType = member.kotlinTypeForNullabilityAnnotation?.takeUnless(KotlinType::isError) ?: return null
+        val psiType = member.psiTypeForNullabilityAnnotation ?: return null
+        if (member.isPrivateOrParameterInPrivateMethod() || psiType is PsiPrimitiveType) return null
+
+        if (kotlinType.isTypeParameter()) {
+            if (!TypeUtils.hasNullableSuperType(kotlinType)) return NotNull::class.java.name
+            if (!kotlinType.isMarkedNullable) return null
+        }
+
+        val nullability = kotlinType.nullability()
+        return when (nullability) {
+            TypeNullability.NOT_NULL -> NotNull::class.java.name
+            TypeNullability.NULLABLE -> Nullable::class.java.name
+            TypeNullability.FLEXIBLE -> null
+        }
+    }
+}
+
+fun PsiModifierListOwner.isPrivateOrParameterInPrivateMethod(): Boolean {
+    if (hasModifier(JvmModifier.PRIVATE)) return true
+    if (this !is PsiParameter) return false
+    val parentMethod = declarationScope as? PsiMethod ?: return false
+    return parentMethod.hasModifier(JvmModifier.PRIVATE)
+}

compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightField.kt

@@ -22,10 +22,7 @@ import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.idea.KotlinLanguage
 import org.jetbrains.kotlin.load.kotlin.TypeMappingMode
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.psi.KtEnumEntry
-import org.jetbrains.kotlin.psi.KtNamedDeclaration
-import org.jetbrains.kotlin.psi.KtObjectDeclaration
-import org.jetbrains.kotlin.psi.KtProperty
+import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.jvm.annotations.TRANSIENT_ANNOTATION_FQ_NAME
 import org.jetbrains.kotlin.resolve.jvm.annotations.VOLATILE_ANNOTATION_FQ_NAME
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOriginKind
@@ -37,7 +34,8 @@ internal open class KtUltraLightField(
     private val containingClass: KtUltraLightClass,
     private val support: UltraLightSupport,
     modifiers: Set<String>
-) : LightFieldBuilder(name, PsiType.NULL, declaration), KtLightField {
+) : LightFieldBuilder(name, PsiType.NULL, declaration), KtLightField,
+    KtUltraLightElementWithNullabilityAnnotation<KtDeclaration, PsiField> {
     private val modList = object : KtLightSimpleModifierList(this, modifiers) {
         override fun hasModifierProperty(name: String): Boolean = when (name) {
             PsiModifier.VOLATILE -> hasFieldAnnotation(VOLATILE_ANNOTATION_FQ_NAME)
@@ -84,6 +82,13 @@ internal open class KtUltraLightField(
         }
     }
 
+    override val kotlinTypeForNullabilityAnnotation: KotlinType?
+        // We don't generate nullability annotations for non-backing fields in backend
+        get() = kotlinType?.takeUnless { declaration is KtEnumEntry || declaration is KtObjectDeclaration }
+
+    override val psiTypeForNullabilityAnnotation: PsiType?
+        get() = type
+
     private val _type: PsiType by lazyPub {
         fun nonExistent() = JavaPsiFacade.getElementFactory(project).createTypeFromText("error.NonExistentClass", declaration)
 

compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightMethod.kt

@@ -13,10 +13,10 @@ import com.intellij.psi.impl.light.LightTypeParameterListBuilder
 import org.jetbrains.annotations.NonNls
 import org.jetbrains.kotlin.asJava.LightClassGenerationSupport
 import org.jetbrains.kotlin.asJava.builder.LightMemberOriginForDeclaration
-import org.jetbrains.kotlin.asJava.elements.KtLightDeclaration
 import org.jetbrains.kotlin.asJava.elements.KtLightMethod
 import org.jetbrains.kotlin.asJava.elements.KtLightMethodImpl
 import org.jetbrains.kotlin.asJava.elements.KtLightSimpleModifierList
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.codegen.FunctionCodegen
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.CallableMemberDescriptor
@@ -33,7 +33,19 @@ internal class KtUltraLightMethod(
     originalElement: KtDeclaration,
     private val support: UltraLightSupport,
     containingClass: KtUltraLightClass
-) : KtLightMethodImpl({ delegate }, LightMemberOriginForDeclaration(originalElement, JvmDeclarationOriginKind.OTHER), containingClass) {
+) : KtLightMethodImpl(
+    { delegate },
+    LightMemberOriginForDeclaration(
+        originalElement, JvmDeclarationOriginKind.OTHER
+    ),
+    containingClass
+), KtUltraLightElementWithNullabilityAnnotation<KtDeclaration, PsiMethod> {
+
+    override val kotlinTypeForNullabilityAnnotation: KotlinType?
+        get() = kotlinOrigin?.getKotlinType()
+
+    override val psiTypeForNullabilityAnnotation: PsiType?
+        get() = returnType
 
     // These two overrides are necessary because ones from KtLightMethodImpl suppose that clsDelegate.returnTypeElement is valid
     // While here we only set return type for LightMethodBuilder (see org.jetbrains.kotlin.asJava.classes.KtUltraLightClass.asJavaMethod)
@@ -80,8 +92,7 @@ internal class KtUltraLightParameter(
     PsiType.NULL,
     method,
     method.language
-),
-    KtLightDeclaration<KtDeclaration, PsiParameter> {
+), KtUltraLightElementWithNullabilityAnnotation<KtDeclaration, PsiParameter> {
 
     override val clsDelegate: PsiParameter
         get() = throw IllegalStateException("Cls delegate shouldn't be loaded for ultra-light PSI!")
@@ -103,6 +114,21 @@ internal class KtUltraLightParameter(
             else -> kotlinOrigin.getKotlinType()
         }
     }
+
+    override val kotlinTypeForNullabilityAnnotation: KotlinType?
+        get() {
+            val type = kotlinType
+            return if (isVarArgs && type != null && KotlinBuiltIns.isArray(type)) {
+                type.arguments[0].type
+            } else {
+                type
+            }
+        }
+
+    override val psiTypeForNullabilityAnnotation: PsiType?
+        get() = type
+
+
     private val _type: PsiType by lazyPub {
         val kotlinType = kotlinType ?: return@lazyPub PsiType.NULL
         val containingDescriptor = containingFunction.resolve() as? CallableDescriptor ?: return@lazyPub PsiType.NULL

compiler/light-classes/src/org/jetbrains/kotlin/asJava/classes/ultraLightPsi.kt

@@ -15,6 +15,7 @@ import com.intellij.psi.impl.light.LightParameterListBuilder
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.asJava.LightClassGenerationSupport
 import org.jetbrains.kotlin.asJava.builder.LightClassData
+import org.jetbrains.kotlin.asJava.elements.KtLightDeclaration
 import org.jetbrains.kotlin.asJava.elements.KtLightField
 import org.jetbrains.kotlin.asJava.elements.KtLightMethod
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
@@ -384,7 +385,7 @@ open class KtUltraLightClass(classOrObject: KtClassOrObject, internal val suppor
             if (numberOfDefaultParametersToAdd >= 0)
                 numberOfDefaultParametersToAdd
             else
-                // Just to avoid computing the actual number of default parameters, we use an upper bound
+            // Just to avoid computing the actual number of default parameters, we use an upper bound
                 ktFunction.valueParameters.size
 
         for (parameter in ktFunction.valueParameters) {
@@ -597,3 +598,9 @@ interface UltraLightSupport {
     fun findAnnotation(owner: KtAnnotated, fqName: FqName): Pair<KtAnnotationEntry, AnnotationDescriptor>?
     fun isTooComplexForUltraLightGeneration(element: KtClassOrObject): Boolean
 }
+
+interface KtUltraLightElementWithNullabilityAnnotation<out T : KtDeclaration, out D : PsiModifierListOwner> : KtLightDeclaration<T, D>,
+    PsiModifierListOwner {
+    val kotlinTypeForNullabilityAnnotation: KotlinType?
+    val psiTypeForNullabilityAnnotation: PsiType?
+}

compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/KtLightModifierList.kt

@@ -22,9 +22,7 @@ import com.intellij.psi.PsiModifierList
 import com.intellij.psi.PsiModifierListOwner
 import org.jetbrains.kotlin.asJava.LightClassGenerationSupport
 import org.jetbrains.kotlin.asJava.builder.LightMemberOriginForDeclaration
-import org.jetbrains.kotlin.asJava.classes.KtLightClassForSourceDeclaration
-import org.jetbrains.kotlin.asJava.classes.KtUltraLightParameter
-import org.jetbrains.kotlin.asJava.classes.lazyPub
+import org.jetbrains.kotlin.asJava.classes.*
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.PropertyDescriptor
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor
@@ -77,8 +75,12 @@ abstract class KtLightModifierList<out T : KtLightElement<KtModifierListOwner, P
         if ((modifierListOwner is KtLightMember<*> && modifierListOwner !is KtLightFieldImpl.KtLightEnumConstant)
             || modifierListOwner is LightParameter
         ) {
-            return annotationsForEntries +
-                    listOf(KtLightNullabilityAnnotation(modifierListOwner as KtLightElement<*, PsiModifierListOwner>, this))
+            val nullabilityAnnotation = when (modifierListOwner) {
+                is KtUltraLightElementWithNullabilityAnnotation<*, *> -> KtUltraLightNullabilityAnnotation(modifierListOwner, this)
+                else -> KtLightNullabilityAnnotation(modifierListOwner as KtLightElement<*, PsiModifierListOwner>, this)
+            }
+
+            return annotationsForEntries + listOf(nullabilityAnnotation)
         }
         return annotationsForEntries
     }

compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/lightAnnotations.kt

@@ -233,7 +233,7 @@ class KtLightEmptyAnnotationParameterList(parent: PsiElement) : KtLightElementBa
     override fun getAttributes(): Array<PsiNameValuePair> = emptyArray()
 }
 
-class KtLightNullabilityAnnotation(val member: KtLightElement<*, PsiModifierListOwner>, parent: PsiElement) :
+open class KtLightNullabilityAnnotation<D : KtLightElement<*, PsiModifierListOwner>>(val member: D, parent: PsiElement) :
     KtLightAbstractAnnotation(parent, {
     // searching for last because nullability annotations are generated after backend generates source annotations
         getClsNullabilityAnnotation(member) ?: KtLightNonExistentAnnotation(member)

compiler/light-classes/src/org/jetbrains/kotlin/asJava/elements/lightAnnotations.kt.173

@@ -234,7 +234,7 @@ class KtLightEmptyAnnotationParameterList(parent: PsiElement) : KtLightElementBa
     override fun getAttributes(): Array<PsiNameValuePair> = emptyArray()
 }
 
-class KtLightNullabilityAnnotation(val member: KtLightElement<*, PsiModifierListOwner>, parent: PsiElement) :
+open class KtLightNullabilityAnnotation<D : KtLightElement<*, PsiModifierListOwner>>(val member: D, parent: PsiElement) :
     KtLightAbstractAnnotation(parent, {
     // searching for last because nullability annotations are generated after backend generates source annotations
         getClsNullabilityAnnotation(member) ?: KtLightNonExistentAnnotation(member)

compiler/testData/asJava/ultraLightClasses/inferringAnonymousObjectTypes.java

@@ -1,5 +1,4 @@
 public final class Prop /* Prop*/ {
-  @org.jetbrains.annotations.NotNull()
   private final java.lang.Object someProp;
 
   @null()
@@ -11,13 +10,11 @@ public final class Fun /* Fun*/ {
   @null()
   public  Fun();
 
-  @org.jetbrains.annotations.NotNull()
   private final java.lang.Object someFun();
 
 }
 
 public final class ArrayOfAnonymous /* ArrayOfAnonymous*/ {
-  @org.jetbrains.annotations.NotNull()
   private final java.lang.Object[] a1;
 
   @null()
@@ -29,11 +26,10 @@ public final class ArrayOfAnonymous /* ArrayOfAnonymous*/ {
 }
 
 final class C /* C*/ {
-  @org.jetbrains.annotations.NotNull()
-  private final kotlin.jvm.functions.Function0<java.lang.Object> initChild;
-
   private final int y;
 
+  private final kotlin.jvm.functions.Function0<java.lang.Object> initChild;
+
   @null()
   public  C(int);
 
@@ -54,7 +50,6 @@ public abstract class Super /* Super*/ {
 }
 
 public final class Sub /* Sub*/ extends Super {
-  @org.jetbrains.annotations.NotNull()
   private final java.lang.Object[] a;
 
   @null()
@@ -66,7 +61,6 @@ public final class Sub /* Sub*/ extends Super {
 }
 
 public final class ValidPublicSupertype /* ValidPublicSupertype*/ {
-  @org.jetbrains.annotations.NotNull()
   private final java.lang.Runnable x;
 
   @null()
@@ -84,7 +78,6 @@ public abstract interface I /* I*/ {
 }
 
 public final class InvalidPublicSupertype /* InvalidPublicSupertype*/ {
-  @org.jetbrains.annotations.NotNull()
   private final java.lang.Runnable x;
 
   @null()

idea/tests/org/jetbrains/kotlin/idea/caches/resolve/AbstractIdeLightClassTest.kt

@@ -272,7 +272,7 @@ object LightClassLazinessChecker {
             }
             clsAnnotations.zip(lightAnnotations).forEach {
                 (clsAnnotation, lightAnnotation) ->
-                if (lightAnnotation !is KtLightNullabilityAnnotation)
+                if (lightAnnotation !is KtLightNullabilityAnnotation<*>)
                     assertNotNull(
                         lightAnnotation!!.nameReferenceElement,
                         "nameReferenceElement should be not null for $lightAnnotation of ${lightAnnotation.javaClass}"