Restore nullability annotations

This is imperfect since arguments to those annotations will be missing, but at least the common case of a plain `@NotNull` or `@Nullable` will be right, and the `@NotNull`s introduced by the Kotlin compiler will be present as expected.

java/kotlin-extractor/src/main/kotlin/KotlinFileExtractor.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.descriptors.java.JavaVisibilities
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.ObsoleteDescriptorBasedAPI
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.backend.js.utils.realOverrideTarget
 import org.jetbrains.kotlin.ir.builders.declarations.*
 import org.jetbrains.kotlin.ir.declarations.*
@@ -44,8 +45,13 @@ import org.jetbrains.kotlin.ir.util.parentClassOrNull
 import org.jetbrains.kotlin.ir.util.primaryConstructor
 import org.jetbrains.kotlin.ir.util.render
 import org.jetbrains.kotlin.ir.util.target
+import org.jetbrains.kotlin.load.java.JvmAnnotationNames
+import org.jetbrains.kotlin.load.java.NOT_NULL_ANNOTATIONS
+import org.jetbrains.kotlin.load.java.NULLABLE_ANNOTATIONS
 import org.jetbrains.kotlin.load.java.sources.JavaSourceElement
+import org.jetbrains.kotlin.load.java.structure.JavaAnnotation
 import org.jetbrains.kotlin.load.java.structure.JavaClass
+import org.jetbrains.kotlin.load.java.structure.JavaConstructor
 import org.jetbrains.kotlin.load.java.structure.JavaMethod
 import org.jetbrains.kotlin.load.java.structure.JavaTypeParameter
 import org.jetbrains.kotlin.load.java.structure.JavaTypeParameterListOwner
@@ -53,6 +59,7 @@ import org.jetbrains.kotlin.load.java.structure.impl.classFiles.BinaryJavaClass
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.util.OperatorNameConventions
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import java.io.Closeable
 import java.util.*
 import kotlin.collections.ArrayList
@@ -887,7 +894,13 @@ open class KotlinFileExtractor(
                 extractTypeAccessRecursive(substitutedType, location, id, -1)
             }
             val syntheticParameterNames = isUnderscoreParameter(vp) || ((vp.parent as? IrFunction)?.let { hasSynthesizedParameterNames(it) } ?: true)
-            extractAnnotations(vp, id, extractTypeAccess)
+            val javaParameter = when(val callable = (vp.parent as? IrFunction)?.let { getJavaCallable(it) }) {
+                is JavaConstructor -> callable.valueParameters.getOrNull(idx)
+                is JavaMethod -> callable.valueParameters.getOrNull(idx)
+                else -> null
+            }
+            val extraAnnotations = listOfNotNull(getNullabilityAnnotation(vp.type, vp.origin, vp.annotations, javaParameter?.annotations))
+            extractAnnotations(vp.annotations + extraAnnotations, id, extractTypeAccess)
             return extractValueParameter(id, substitutedType, vp.name.asString(), location, parent, idx, useValueParameter(vp, parentSourceDeclaration), syntheticParameterNames, vp.isVararg, vp.isNoinline, vp.isCrossinline)
         }
     }
@@ -1337,6 +1350,32 @@ open class KotlinFileExtractor(
                 logger.warn("Needed a signature for a type that doesn't have one")
         }
 
+    private fun getNullabilityAnnotationName(t: IrType, declOrigin: IrDeclarationOrigin, existingAnnotations: List<IrConstructorCall>, javaAnnotations: Collection<JavaAnnotation>?): FqName? {
+        if (t !is IrSimpleType)
+            return null
+
+        return if (declOrigin == IrDeclarationOrigin.IR_EXTERNAL_JAVA_DECLARATION_STUB) {
+            // Java declaration: restore a NotNull or Nullable annotation if the original Java member had one but the Kotlin compiler removed it.
+            javaAnnotations?.mapNotNull { it.classId?.asSingleFqName() }
+            ?.singleOrNull { NOT_NULL_ANNOTATIONS.contains(it) || NULLABLE_ANNOTATIONS.contains(it) }
+            ?.takeUnless { existingAnnotations.any { existing -> existing.type.classFqName == it } }
+        } else {
+            // Kotlin declaration: add a NotNull annotation to a non-nullable non-primitive type.
+            JvmAnnotationNames.JETBRAINS_NOT_NULL_ANNOTATION.takeUnless { t.isNullable() || primitiveTypeMapping.getPrimitiveInfo(t) != null }
+        }
+    }
+
+    private fun getNullabilityAnnotation(t: IrType, declOrigin: IrDeclarationOrigin, existingAnnotations: List<IrConstructorCall>, javaAnnotations: Collection<JavaAnnotation>?) =
+        getNullabilityAnnotationName(t, declOrigin, existingAnnotations, javaAnnotations)?.let {
+            pluginContext.referenceClass(it)?.let { annotationClass ->
+                annotationClass.owner.declarations.firstIsInstanceOrNull<IrConstructor>()?.let { annotationConstructor ->
+                    IrConstructorCallImpl.fromSymbolOwner(
+                        UNDEFINED_OFFSET, UNDEFINED_OFFSET, annotationConstructor.returnType, annotationConstructor.symbol, 0
+                    )
+                }
+            }
+        }
+
     private fun forceExtractFunction(f: IrFunction, parentId: Label<out DbReftype>, extractBody: Boolean, extractMethodAndParameterTypeAccesses: Boolean, extractAnnotations: Boolean, typeSubstitution: TypeSubstitution?, classTypeArgsIncludingOuterClasses: List<IrTypeArgument>?, extractOrigin: Boolean = true, overriddenAttributes: OverriddenFunctionAttributes? = null): Label<out DbCallable>  {
         with("function", f) {
             DeclarationStackAdjuster(f, overriddenAttributes).use {
@@ -1427,8 +1466,10 @@ open class KotlinFileExtractor(
 
                 linesOfCode?.linesOfCodeInDeclaration(f, id)
 
-                if (extractAnnotations)
-                    extractAnnotations(f, id, extractMethodAndParameterTypeAccesses)
+                if (extractAnnotations) {
+                    val extraAnnotations = listOfNotNull(getNullabilityAnnotation(f.returnType, f.origin, f.annotations, getJavaCallable(f)?.annotations))
+                    extractAnnotations(f.annotations + extraAnnotations, id, extractMethodAndParameterTypeAccesses)
+                }
 
                 return id
             }

java/ql/integration-tests/all-platforms/kotlin/nullability-annotations/AnnotatedMethods.java

@@ -0,0 +1,10 @@
+import org.jetbrains.annotations.*;
+import zpkg.A;
+
+public class AnnotatedMethods {
+
+  public @A @NotNull String notNullAnnotated(@A @NotNull String param) { return param; }
+
+  public @A @Nullable String nullableAnnotated(@A @Nullable String param) { return param; }
+
+}

java/ql/integration-tests/all-platforms/kotlin/nullability-annotations/JavaUser.java

@@ -0,0 +1,7 @@
+public class JavaUser {
+
+  public static void test(KotlinAnnotatedMethods km) {
+    km.f(null);
+  }
+
+}

java/ql/integration-tests/all-platforms/kotlin/nullability-annotations/ktUser.kt

@@ -0,0 +1,7 @@
+import zpkg.A
+
+class KotlinAnnotatedMethods {
+
+  @A fun f(@A m: AnnotatedMethods): String = m.notNullAnnotated("hello") + m.nullableAnnotated("world")!!
+
+}

java/ql/integration-tests/all-platforms/kotlin/nullability-annotations/org/jetbrains/annotations/NotNull.java

@@ -0,0 +1,6 @@
+package org.jetbrains.annotations;
+
+public @interface NotNull {
+  String value() default "";
+  Class<? extends Exception> exception() default Exception.class;
+}

java/ql/integration-tests/all-platforms/kotlin/nullability-annotations/org/jetbrains/annotations/Nullable.java

@@ -0,0 +1,5 @@
+package org.jetbrains.annotations;
+
+public @interface Nullable {
+  String value() default "";
+}

java/ql/integration-tests/all-platforms/kotlin/nullability-annotations/test.expected

@@ -0,0 +1,12 @@
+| AnnotatedMethods.java:6:29:6:44 | notNullAnnotated | parameter | AnnotatedMethods.java:6:46:6:47 | A |
+| AnnotatedMethods.java:6:29:6:44 | notNullAnnotated | parameter | AnnotatedMethods.java:6:49:6:56 | NotNull |
+| AnnotatedMethods.java:6:29:6:44 | notNullAnnotated | return value | AnnotatedMethods.java:6:10:6:11 | A |
+| AnnotatedMethods.java:6:29:6:44 | notNullAnnotated | return value | AnnotatedMethods.java:6:13:6:20 | NotNull |
+| AnnotatedMethods.java:8:30:8:46 | nullableAnnotated | parameter | AnnotatedMethods.java:8:48:8:49 | A |
+| AnnotatedMethods.java:8:30:8:46 | nullableAnnotated | parameter | AnnotatedMethods.java:8:51:8:59 | Nullable |
+| AnnotatedMethods.java:8:30:8:46 | nullableAnnotated | return value | AnnotatedMethods.java:8:10:8:11 | A |
+| AnnotatedMethods.java:8:30:8:46 | nullableAnnotated | return value | AnnotatedMethods.java:8:13:8:21 | Nullable |
+| ktUser.kt:5:6:5:105 | f | parameter | ktUser.kt:0:0:0:0 | NotNull |
+| ktUser.kt:5:6:5:105 | f | parameter | ktUser.kt:5:12:5:13 | A |
+| ktUser.kt:5:6:5:105 | f | return value | ktUser.kt:0:0:0:0 | NotNull |
+| ktUser.kt:5:6:5:105 | f | return value | ktUser.kt:5:3:5:4 | A |

java/ql/integration-tests/all-platforms/kotlin/nullability-annotations/test.py

@@ -0,0 +1,6 @@
+from create_database_utils import *
+
+os.mkdir('out')
+os.mkdir('out2')
+os.mkdir('out3')
+run_codeql_database_create(["javac AnnotatedMethods.java zpkg/A.java org/jetbrains/annotations/NotNull.java org/jetbrains/annotations/Nullable.java -d out", "kotlinc ktUser.kt -cp out -d out2", "javac JavaUser.java -cp out:out2 -d out3"], lang="java")

java/ql/integration-tests/all-platforms/kotlin/nullability-annotations/test.ql

@@ -0,0 +1,11 @@
+import java
+
+from Method m, string origin, Annotation a
+where
+  m.fromSource() and
+  (
+    origin = "return value" and a = m.getAnAnnotation()
+    or
+    origin = "parameter" and a = m.getAParameter().getAnAnnotation()
+  )
+select m, origin, a

java/ql/integration-tests/all-platforms/kotlin/nullability-annotations/zpkg/A.java

@@ -0,0 +1,3 @@
+package zpkg;
+
+public @interface A { }