Include nested annotation information in astubx files when loading external library models

jar-infer/jar-infer-lib/src/main/java/com/uber/nullaway/jarinfer/DefinitelyDerefedParamsDriver.java

@@ -19,6 +19,7 @@
 import com.google.common.base.Verify;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSetMultimap;
 import com.ibm.wala.cfg.ControlFlowGraph;
 import com.ibm.wala.classLoader.CodeScanner;
 import com.ibm.wala.classLoader.IClass;
@@ -457,14 +458,18 @@ private void writeModel(DataOutputStream out) throws IOException {
           MethodAnnotationsRecord.create(
               nullableReturns.contains(sign) ? ImmutableSet.of("Nullable") : ImmutableSet.of(),
               ImmutableSet.of(),
-              ImmutableMap.copyOf(argAnnotation)));
+              ImmutableMap.copyOf(argAnnotation),
+              ImmutableSetMultimap.of()));
       nullableReturns.remove(sign);
     }
     for (String nullableReturnMethodSign : Iterator2Iterable.make(nullableReturns.iterator())) {
       methodRecords.put(
           nullableReturnMethodSign,
           MethodAnnotationsRecord.create(
-              ImmutableSet.of("Nullable"), ImmutableSet.of(), ImmutableMap.of()));
+              ImmutableSet.of("Nullable"),
+              ImmutableSet.of(),
+              ImmutableMap.of(),
+              ImmutableSetMultimap.of()));
     }
     StubxWriter.write(
         out,

jdk-annotations/astubx-generator/src/main/java/com/uber/nullaway/jdkannotations/AstubxGenerator.java

@@ -3,6 +3,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSetMultimap;
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
 import com.google.gson.JsonDeserializer;
@@ -11,6 +12,7 @@
 import com.uber.nullaway.javacplugin.NullnessAnnotationSerializer.MethodInfo;
 import com.uber.nullaway.javacplugin.NullnessAnnotationSerializer.TypeParamInfo;
 import com.uber.nullaway.libmodel.MethodAnnotationsRecord;
+import com.uber.nullaway.libmodel.NestedAnnotationInfo;
 import com.uber.nullaway.libmodel.StubxWriter;
 import java.io.DataOutputStream;
 import java.io.File;
@@ -272,17 +274,24 @@ private static void getMethodRecords(
           }
         } else {
           // remove any spaces in Array types
-          typeSignature = typeSignature.replace(" ", "");
+          typeSignature = typeSignature.replace(" []", "[]");
         }
         argumentList[i] = typeSignature;
       }
+      ImmutableSetMultimap.Builder<Integer, NestedAnnotationInfo> nestedAnnotations =
+          new ImmutableSetMultimap.Builder<>();
+      for (Map.Entry<Integer, Set<NestedAnnotationInfo>> nestedInfo :
+          method.nestedAnnotationsList().entrySet()) {
+        nestedAnnotations.putAll(nestedInfo.getKey(), nestedInfo.getValue());
+      }
       signatureForMethodRecords += String.join(",", argumentList) + ")";
       methodRecords.put(
           signatureForMethodRecords,
           MethodAnnotationsRecord.create(
               returnTypeNullness,
               nullableTypeParamBuilder.build(),
-              ImmutableMap.copyOf(argAnnotation)));
+              ImmutableMap.copyOf(argAnnotation),
+              nestedAnnotations.build()));
     }
   }
 

jdk-annotations/astubx-generator/src/test/java/com/uber/nullaway/jdkannotations/AstubxTest.java

@@ -4,12 +4,18 @@
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.MatcherAssert.assertThat;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSetMultimap;
 import com.google.errorprone.BugPattern;
 import com.google.errorprone.CompilationTestHelper;
 import com.google.errorprone.bugpatterns.BugChecker;
 import com.uber.nullaway.libmodel.MethodAnnotationsRecord;
+import com.uber.nullaway.libmodel.NestedAnnotationInfo;
+import com.uber.nullaway.libmodel.NestedAnnotationInfo.Annotation;
+import com.uber.nullaway.libmodel.NestedAnnotationInfo.TypePathEntry;
+import com.uber.nullaway.libmodel.NestedAnnotationInfo.TypePathEntry.Kind;
 import java.nio.file.Files;
 import java.nio.file.Paths;
 import java.util.Arrays;
@@ -66,7 +72,10 @@ public void nullableReturn() {
         ImmutableMap.of(
             "AnnotationExample:java.lang.String makeUpperCase(java.lang.String)",
             MethodAnnotationsRecord.create(
-                ImmutableSet.of("Nullable"), ImmutableSet.of(), ImmutableMap.of()));
+                ImmutableSet.of("Nullable"),
+                ImmutableSet.of(),
+                ImmutableMap.of(),
+                ImmutableSetMultimap.of()));
     runTest(expectedMethodRecords, ImmutableMap.of(), ImmutableSet.of("AnnotationExample"));
   }
 
@@ -92,7 +101,10 @@ public void nullableUpperBound() {
         ImmutableMap.of(
             "NullableUpperBound:T withAnnotation()",
             MethodAnnotationsRecord.create(
-                ImmutableSet.of("Nullable"), ImmutableSet.of(), ImmutableMap.of()));
+                ImmutableSet.of("Nullable"),
+                ImmutableSet.of(),
+                ImmutableMap.of(),
+                ImmutableSetMultimap.of()));
     ImmutableMap<String, Set<Integer>> expectedNullableUpperBounds =
         ImmutableMap.of("NullableUpperBound", ImmutableSet.of(0));
     runTest(
@@ -123,7 +135,10 @@ public void nestedNullableUpperBound() {
         ImmutableMap.of(
             "ReturnAnnotation.UpperBoundExample:T withAnnotation()",
             MethodAnnotationsRecord.create(
-                ImmutableSet.of("Nullable"), ImmutableSet.of(), ImmutableMap.of()));
+                ImmutableSet.of("Nullable"),
+                ImmutableSet.of(),
+                ImmutableMap.of(),
+                ImmutableSetMultimap.of()));
     ImmutableMap<String, Set<Integer>> expectedNullableUpperBounds =
         ImmutableMap.of("ReturnAnnotation.UpperBoundExample", ImmutableSet.of(0));
     runTest(
@@ -181,7 +196,8 @@ public void nullableParameters() {
             MethodAnnotationsRecord.create(
                 ImmutableSet.of(),
                 ImmutableSet.of(),
-                ImmutableMap.of(0, ImmutableSet.of("Nullable"))));
+                ImmutableMap.of(0, ImmutableSet.of("Nullable")),
+                ImmutableSetMultimap.of()));
     runTest(expectedMethodRecords, ImmutableMap.of(), ImmutableSet.of("NullableParameters"));
   }
 
@@ -208,7 +224,8 @@ public void nullableParametersInNullUnmarkedClass() {
             MethodAnnotationsRecord.create(
                 ImmutableSet.of(),
                 ImmutableSet.of(),
-                ImmutableMap.of(0, ImmutableSet.of("Nullable"))));
+                ImmutableMap.of(0, ImmutableSet.of("Nullable")),
+                ImmutableSetMultimap.of()));
     runTest(expectedMethodRecords, ImmutableMap.of(), ImmutableSet.of());
   }
 
@@ -237,7 +254,8 @@ public void nullableArrayTypeParameter() {
             MethodAnnotationsRecord.create(
                 ImmutableSet.of(),
                 ImmutableSet.of(),
-                ImmutableMap.of(0, ImmutableSet.of("Nullable"))));
+                ImmutableMap.of(0, ImmutableSet.of("Nullable")),
+                ImmutableSetMultimap.of()));
     runTest(expectedMethodRecords, ImmutableMap.of(), ImmutableSet.of("NullableParameters"));
   }
 
@@ -261,7 +279,8 @@ public void genericParameter() {
             MethodAnnotationsRecord.create(
                 ImmutableSet.of(),
                 ImmutableSet.of(),
-                ImmutableMap.of(0, ImmutableSet.of("Nullable"))));
+                ImmutableMap.of(0, ImmutableSet.of("Nullable")),
+                ImmutableSetMultimap.of()));
     runTest(expectedMethodRecords, ImmutableMap.of(), ImmutableSet.of("Generic"));
   }
 
@@ -282,13 +301,49 @@ public void parameterizedTypeArray() {
             "  public List<@Nullable String>[] nullableIdentity(List<@Nullable String>[] listArray) {",
             "    return listArray;",
             "  }",
+            "  public List<? extends @Nullable String> wildcardIdentity(List<? super @Nullable String> listArray) {",
+            "    return null;",
+            "  }",
             "}")
         .doTest();
     ImmutableMap<String, MethodAnnotationsRecord> expectedMethodRecords =
         ImmutableMap.of(
             "ParameterizedTypeArray:java.util.List<java.lang.String>[] nullableIdentity(java.util.List<java.lang.String>[])",
             MethodAnnotationsRecord.create(
-                ImmutableSet.of(), ImmutableSet.of(), ImmutableMap.of()));
+                ImmutableSet.of(),
+                ImmutableSet.of(),
+                ImmutableMap.of(),
+                ImmutableSetMultimap.of(
+                    -1,
+                    new NestedAnnotationInfo(
+                        Annotation.NULLABLE,
+                        ImmutableList.of(
+                            new TypePathEntry(Kind.ARRAY_ELEMENT, -1),
+                            new TypePathEntry(Kind.TYPE_ARGUMENT, 0))),
+                    0,
+                    new NestedAnnotationInfo(
+                        Annotation.NULLABLE,
+                        ImmutableList.of(
+                            new TypePathEntry(Kind.ARRAY_ELEMENT, -1),
+                            new TypePathEntry(Kind.TYPE_ARGUMENT, 0))))),
+            "ParameterizedTypeArray:java.util.List<? extends java.lang.String> wildcardIdentity(java.util.List<? super java.lang.String>)",
+            MethodAnnotationsRecord.create(
+                ImmutableSet.of(),
+                ImmutableSet.of(),
+                ImmutableMap.of(),
+                ImmutableSetMultimap.of(
+                    -1,
+                    new NestedAnnotationInfo(
+                        Annotation.NULLABLE,
+                        ImmutableList.of(
+                            new TypePathEntry(Kind.TYPE_ARGUMENT, 0),
+                            new TypePathEntry(Kind.WILDCARD_BOUND, 0))),
+                    0,
+                    new NestedAnnotationInfo(
+                        Annotation.NULLABLE,
+                        ImmutableList.of(
+                            new TypePathEntry(Kind.TYPE_ARGUMENT, 0),
+                            new TypePathEntry(Kind.WILDCARD_BOUND, 1))))));
     runTest(expectedMethodRecords, ImmutableMap.of(), ImmutableSet.of("ParameterizedTypeArray"));
   }
 
@@ -315,7 +370,8 @@ public void primitiveTypeReturn() {
             MethodAnnotationsRecord.create(
                 ImmutableSet.of(),
                 ImmutableSet.of(),
-                ImmutableMap.of(0, ImmutableSet.of("Nullable"), 1, ImmutableSet.of("Nullable"))));
+                ImmutableMap.of(0, ImmutableSet.of("Nullable"), 1, ImmutableSet.of("Nullable")),
+                ImmutableSetMultimap.of()));
     runTest(expectedMethodRecords, ImmutableMap.of(), ImmutableSet.of("PrimitiveType"));
   }
 
@@ -341,7 +397,8 @@ public void voidReturn() {
             MethodAnnotationsRecord.create(
                 ImmutableSet.of(),
                 ImmutableSet.of(),
-                ImmutableMap.of(0, ImmutableSet.of("Nullable"), 1, ImmutableSet.of("Nullable"))));
+                ImmutableMap.of(0, ImmutableSet.of("Nullable"), 1, ImmutableSet.of("Nullable")),
+                ImmutableSetMultimap.of()));
     runTest(expectedMethodRecords, ImmutableMap.of(), ImmutableSet.of("VoidReturn"));
   }
 
@@ -361,7 +418,10 @@ public void methodTypeParameter() {
         ImmutableMap.of(
             "Test:void <K,T>nullableTypeVar(K,T)",
             MethodAnnotationsRecord.create(
-                ImmutableSet.of(), ImmutableSet.of(1), ImmutableMap.of()));
+                ImmutableSet.of(),
+                ImmutableSet.of(1),
+                ImmutableMap.of(),
+                ImmutableSetMultimap.of()));
     runTest(expectedMethodRecords, ImmutableMap.of(), ImmutableSet.of("Test"));
   }
 

jdk-annotations/jdk-integration-test/src/test/java/com/uber/nullaway/jdkannotations/JDKIntegrationTest.java

@@ -388,4 +388,95 @@ public void loadLibraryModuleMethodTypeParam() {
             "}")
         .doTest();
   }
+
+  @Test
+  public void loadLibraryModuleMethodReturnTypeNestedAnnotation() {
+    compilationHelper
+        .setArgs(
+            Arrays.asList(
+                "-d",
+                temporaryFolder.getRoot().getAbsolutePath(),
+                "-XepOpt:NullAway:AnnotatedPackages=com.uber",
+                "-XepOpt:NullAway:JarInferEnabled=true",
+                "-XepOpt:NullAway:JSpecifyMode=true",
+                "-XDaddTypeAnnotationsToSymbol=true"))
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import org.jspecify.annotations.Nullable;",
+            "import com.uber.nullaway.jdkannotations.ReturnAnnotation;",
+            "import java.util.*;",
+            "class Test {",
+            "  void testCall() {",
+            "    // -- Type argument nullability",
+            "    // BUG: Diagnostic contains: incompatible types",
+            "    List<String> typeArg1 = ReturnAnnotation.nestedAnnotTypeArg();",
+            "    List<@Nullable String> typeArg2 = ReturnAnnotation.nestedAnnotTypeArg();",
+            "    // -- Array Element nullability",
+            "    // BUG: Diagnostic contains: incompatible types",
+            "    String[] arrayElement1 = ReturnAnnotation.nestedAnnotArrayElement();",
+            "    // BUG: Diagnostic contains: incompatible types",
+            "    String @Nullable [] arrayElement2 = ReturnAnnotation.nestedAnnotArrayElement();",
+            "    @Nullable String [] arrayElement3 = ReturnAnnotation.nestedAnnotArrayElement();",
+            "    // -- mixed type nullability",
+            "    // BUG: Diagnostic contains: incompatible types",
+            "    List<Integer>[] mixed1 = ReturnAnnotation.nestedAnnotMixed();",
+            "    // BUG: Diagnostic contains: incompatible types",
+            "    @Nullable List<Integer>[] mixed2 = ReturnAnnotation.nestedAnnotMixed();",
+            "    // BUG: Diagnostic contains: incompatible types",
+            "    List<@Nullable Integer>[] mixed3 = ReturnAnnotation.nestedAnnotMixed();",
+            "    @Nullable List<@Nullable Integer>[] mixed4 = ReturnAnnotation.nestedAnnotMixed();",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void loadLibraryModuleMethodParameterTypeNestedAnnotation() {
+    compilationHelper
+        .setArgs(
+            Arrays.asList(
+                "-d",
+                temporaryFolder.getRoot().getAbsolutePath(),
+                "-XepOpt:NullAway:AnnotatedPackages=com.uber",
+                "-XepOpt:NullAway:JarInferEnabled=true",
+                "-XepOpt:NullAway:JSpecifyMode=true",
+                "-XDaddTypeAnnotationsToSymbol=true"))
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import org.jspecify.annotations.Nullable;",
+            "import com.uber.nullaway.jdkannotations.ParameterAnnotation;",
+            "import java.util.*;",
+            "class Test {",
+            "  void testCall() {",
+            "    // -- Type argument",
+            "    List<@Nullable String> nullableTypeArg = new ArrayList<>();",
+            "    nullableTypeArg.add(\"string\");",
+            "    nullableTypeArg.add(null);",
+            "    List<String> nonNullTypeArg = new ArrayList<>(2);",
+            "    // -- Array type",
+            "    @Nullable String[] nullableArray = new String[] {\"populated\", \"value\", null};",
+            "    String[] nonNullArray = new String[] {\"populated\", \"value\"};",
+            "    // -- Multiple nested annotations",
+            "    List<@Nullable Integer> innerList = new ArrayList<>();",
+            "    innerList.add(null);",
+            "    innerList.add(4);",
+            "    @Nullable List<@Nullable Integer>[] nullableMixed = (@Nullable List<@Nullable Integer>[]) new List<?>[3];",
+            "    nullableMixed[0] = innerList;",
+            "    nullableMixed[0] = null;",
+            "    List<@Nullable Integer>[] nonNullArrayMixed = (List<@Nullable Integer>[]) new List<?>[3];",
+            "    @Nullable List<Integer>[] nonNullTypeArgMixed = (@Nullable List<Integer>[]) new List<?>[3];",
+            "    // === test calls",
+            "    ParameterAnnotation.nestedAnnotations(nullableTypeArg, nullableArray, nullableMixed);",
+            "    // BUG: Diagnostic contains: incompatible types",
+            "    ParameterAnnotation.nestedAnnotations(nonNullTypeArg, nullableArray, nullableMixed);",
+            "    ParameterAnnotation.nestedAnnotations(nullableTypeArg, nonNullArray, nullableMixed);",
+            "    ParameterAnnotation.nestedAnnotations(nullableTypeArg, nullableArray, nonNullArrayMixed);",
+            "    // BUG: Diagnostic contains: incompatible types",
+            "    ParameterAnnotation.nestedAnnotations(nullableTypeArg, nullableArray, nonNullTypeArgMixed);",
+            "  }",
+            "}")
+        .doTest();
+  }
 }

jdk-annotations/test-annotated/src/main/java/com/uber/nullaway/jdkannotations/ParameterAnnotation.java

@@ -1,5 +1,6 @@
 package com.uber.nullaway.jdkannotations;
 
+import java.util.List;
 import org.jspecify.annotations.NonNull;
 import org.jspecify.annotations.NullMarked;
 import org.jspecify.annotations.Nullable;
@@ -64,4 +65,9 @@ public static void takesNonNullGenericArray(Generic<String>[] objects) {
   }
 
   public static <T> void nonNullTypeParam(T t) {}
+
+  public static void nestedAnnotations(
+      List<@Nullable String> typeArg,
+      @Nullable String[] array,
+      @Nullable List<@Nullable Integer>[] mixed) {}
 }