Update jdk-javac-plugin to contain nested annotations

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

@@ -1,8 +1,11 @@
 package com.uber.nullaway.jdkannotations;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.gson.Gson;
+import com.google.gson.GsonBuilder;
+import com.google.gson.JsonDeserializer;
 import com.google.gson.reflect.TypeToken;
 import com.uber.nullaway.javacplugin.NullnessAnnotationSerializer.ClassInfo;
 import com.uber.nullaway.javacplugin.NullnessAnnotationSerializer.MethodInfo;
@@ -13,6 +16,7 @@
 import java.io.File;
 import java.io.FileOutputStream;
 import java.io.IOException;
+import java.lang.reflect.ParameterizedType;
 import java.lang.reflect.Type;
 import java.nio.file.Files;
 import java.nio.file.Paths;
@@ -164,7 +168,35 @@ private static Map<String, List<ClassInfo>> parseJson(String jsonDirPath) {
       throw new IllegalStateException("No JSON files found in: " + jsonDirPath);
     }
 
-    Gson gson = new Gson();
+    Gson gson =
+        new GsonBuilder()
+            .registerTypeAdapter(
+                ImmutableList.class,
+                (JsonDeserializer<ImmutableList<?>>)
+                    (json, type, context) -> {
+                      if (json.isJsonNull()) {
+                        return ImmutableList.of();
+                      }
+                      if (!(type instanceof ParameterizedType paramType)) {
+                        // Raw ImmutableList, deserialize as List<Object>
+                        List<?> list = context.deserialize(json, List.class);
+                        return list == null ? ImmutableList.of() : ImmutableList.copyOf(list);
+                      }
+                      // Get type inside the list
+                      Type[] typeArgs = paramType.getActualTypeArguments();
+                      Type innerType = typeArgs.length > 0 ? typeArgs[0] : Object.class;
+
+                      // Get as ArrayList
+                      List<?> standardList =
+                          context.deserialize(
+                              json, TypeToken.getParameterized(List.class, innerType).getType());
+
+                      // Convert to Guava ImmutableList
+                      return standardList == null
+                          ? ImmutableList.of()
+                          : ImmutableList.copyOf(standardList);
+                    })
+            .create();
     Type parsedType = new TypeToken<Map<String, List<ClassInfo>>>() {}.getType();
 
     // parse JSON file

jdk-javac-plugin/build.gradle

@@ -57,6 +57,7 @@ dependencies {
 
     implementation libs.gson
     implementation libs.jspecify
+    implementation libs.guava
 
     testImplementation libs.junit4
     testImplementation(libs.error.prone.test.helpers) {

jdk-javac-plugin/src/main/java/com/uber/nullaway/javacplugin/CreateNestedAnnotationInfoVisitor.java

@@ -0,0 +1,150 @@
+package com.uber.nullaway.javacplugin;
+
+import com.google.common.collect.ImmutableList;
+import com.sun.tools.javac.code.Type;
+import com.sun.tools.javac.code.Types;
+import com.uber.nullaway.javacplugin.NestedAnnotationInfo.Annotation;
+import com.uber.nullaway.javacplugin.NestedAnnotationInfo.TypePathEntry;
+import java.util.ArrayDeque;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import javax.lang.model.element.AnnotationMirror;
+import javax.lang.model.element.TypeElement;
+import javax.lang.model.type.TypeMirror;
+import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
+
+@NullMarked
+public class CreateNestedAnnotationInfoVisitor
+    extends Types.DefaultTypeVisitor<Set<NestedAnnotationInfo>, @Nullable Void> {
+
+  private final ArrayDeque<TypePathEntry> path;
+  private final Set<NestedAnnotationInfo> nestedAnnotationInfoSet;
+
+  private static final String NULLABLE_QNAME = "org.jspecify.annotations.Nullable";
+  private static final String NONNULL_QNAME = "org.jspecify.annotations.NonNull";
+
+  public CreateNestedAnnotationInfoVisitor() {
+    path = new ArrayDeque<>();
+    nestedAnnotationInfoSet = new HashSet<>();
+  }
+
+  @Override
+  public @Nullable Set<NestedAnnotationInfo> visitClassType(
+      Type.ClassType classType, @Nullable Void unused) {
+    // only processes type arguments
+    List<Type> typeArguments = classType.getTypeArguments();
+    if (!typeArguments.isEmpty()) {
+      for (int idx = 0; idx < typeArguments.size(); idx++) {
+        path.addLast(new TypePathEntry(TypePathEntry.Kind.TYPE_ARGUMENT, idx));
+        Type typeArg = typeArguments.get(idx);
+        ImmutableList<TypePathEntry> typePath = getTypePath();
+        NestedAnnotationInfo nestedAnnotationInfo = null;
+        if (hasNullableAnnotation(typeArg)) {
+          nestedAnnotationInfo = new NestedAnnotationInfo(Annotation.NULLABLE, typePath);
+        } else if (hasNonNullAnnotation(typeArg)) {
+          nestedAnnotationInfo = new NestedAnnotationInfo(Annotation.NONNULL, typePath);
+        }
+        if (nestedAnnotationInfo != null) {
+          nestedAnnotationInfoSet.add(nestedAnnotationInfo);
+        }
+        typeArg.accept(this, null);
+        path.removeLast();
+      }
+    }
+    return nestedAnnotationInfoSet;
+  }
+
+  @Override
+  public @Nullable Set<NestedAnnotationInfo> visitArrayType(
+      Type.ArrayType t, @Nullable Void unused) {
+    path.addLast(new TypePathEntry(TypePathEntry.Kind.ARRAY_ELEMENT, -1));
+    ImmutableList<TypePathEntry> typePath = getTypePath();
+    NestedAnnotationInfo nestedAnnotationInfo = null;
+    if (hasNullableAnnotation(t)) {
+      nestedAnnotationInfo = new NestedAnnotationInfo(Annotation.NULLABLE, typePath);
+    } else if (hasNonNullAnnotation(t)) {
+      nestedAnnotationInfo = new NestedAnnotationInfo(Annotation.NONNULL, typePath);
+    }
+    if (nestedAnnotationInfo != null) {
+      nestedAnnotationInfoSet.add(nestedAnnotationInfo);
+    }
+    t.elemtype.accept(this, null);
+    path.removeLast();
+    return nestedAnnotationInfoSet;
+  }
+
+  @Override
+  public @Nullable Set<NestedAnnotationInfo> visitWildcardType(
+      Type.WildcardType t, @Nullable Void unused) {
+    // Upper Bound (? extends T)
+    if (t.getExtendsBound() != null) {
+      NestedAnnotationInfo nestedAnnotationInfo = null;
+      path.addLast(new TypePathEntry(TypePathEntry.Kind.WILDCARD_BOUND, 0));
+      ImmutableList<TypePathEntry> typePath = getTypePath();
+      Type upperBound = t.getExtendsBound();
+      if (hasNullableAnnotation(upperBound)) {
+        nestedAnnotationInfo = new NestedAnnotationInfo(Annotation.NULLABLE, typePath);
+      } else if (hasNonNullAnnotation(upperBound)) {
+        nestedAnnotationInfo = new NestedAnnotationInfo(Annotation.NONNULL, typePath);
+      }
+      if (nestedAnnotationInfo != null) {
+        nestedAnnotationInfoSet.add(nestedAnnotationInfo);
+      }
+      t.getExtendsBound().accept(this, null);
+      path.removeLast();
+    }
+
+    // Lower Bound (? super T)
+    if (t.getSuperBound() != null) {
+      NestedAnnotationInfo nestedAnnotationInfo = null;
+      path.addLast(new TypePathEntry(TypePathEntry.Kind.WILDCARD_BOUND, 1));
+      ImmutableList<TypePathEntry> typePath = getTypePath();
+      Type lowerBound = t.getSuperBound();
+      if (hasNullableAnnotation(lowerBound)) {
+        nestedAnnotationInfo = new NestedAnnotationInfo(Annotation.NULLABLE, typePath);
+      } else if (hasNonNullAnnotation(lowerBound)) {
+        nestedAnnotationInfo = new NestedAnnotationInfo(Annotation.NONNULL, typePath);
+      }
+      if (nestedAnnotationInfo != null) {
+        nestedAnnotationInfoSet.add(nestedAnnotationInfo);
+      }
+      t.getSuperBound().accept(this, null);
+      path.removeLast();
+    }
+
+    return nestedAnnotationInfoSet;
+  }
+
+  @Override
+  public @Nullable Set<NestedAnnotationInfo> visitType(Type t, @Nullable Void unused) {
+    return nestedAnnotationInfoSet;
+  }
+
+  private static boolean hasAnnotation(TypeMirror type, String qname) {
+    if (type == null) {
+      return false;
+    }
+    for (AnnotationMirror annotation : type.getAnnotationMirrors()) {
+      String qualifiedName =
+          ((TypeElement) annotation.getAnnotationType().asElement()).getQualifiedName().toString();
+      if (qualifiedName.equals(qname)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  private boolean hasNullableAnnotation(TypeMirror type) {
+    return hasAnnotation(type, NULLABLE_QNAME);
+  }
+
+  private boolean hasNonNullAnnotation(TypeMirror type) {
+    return hasAnnotation(type, NONNULL_QNAME);
+  }
+
+  private ImmutableList<TypePathEntry> getTypePath() {
+    return ImmutableList.copyOf(path);
+  }
+}

jdk-javac-plugin/src/main/java/com/uber/nullaway/javacplugin/NestedAnnotationInfo.java

@@ -0,0 +1,42 @@
+package com.uber.nullaway.javacplugin;
+
+import com.google.common.collect.ImmutableList;
+import org.jspecify.annotations.NullMarked;
+
+/**
+ * Class to hold information about a nested nullability annotation within a type, including the type
+ * of nullability annotation and the type path to reach it.
+ *
+ * @param annotation the nullability annotation
+ * @param typePath the type path to reach the annotation. If empty, the annotation applies to the
+ *     outermost type. Otherwise, each entry indicates one step in how to navigate to the nested
+ *     type.
+ */
+@NullMarked
+public record NestedAnnotationInfo(Annotation annotation, ImmutableList<TypePathEntry> typePath) {
+
+  /**
+   * Class for a single entry in a type path, indicating how to navigate the "next step" in the type
+   * to eventually reach some target type.
+   *
+   * @param kind the kind of this type path entry
+   * @param index the index associated with the kind. For TYPE_ARGUMENT, this is the type argument
+   *     index. For WILDCARD_BOUND, this is 0 for the upper bound ({@code ? extends Foo}) and 1 for
+   *     the lower bound ({@code ? super Foo}). For ARRAY_ELEMENT, this is unused and set to -1.
+   */
+  public record TypePathEntry(Kind kind, int index) {
+
+    /** Possible nested type kinds for an entry */
+    public enum Kind {
+      ARRAY_ELEMENT,
+      TYPE_ARGUMENT,
+      WILDCARD_BOUND
+    }
+  }
+
+  /** Possible annotations for nullability */
+  public enum Annotation {
+    NULLABLE,
+    NONNULL
+  }
+}

jdk-javac-plugin/src/main/java/com/uber/nullaway/javacplugin/NullnessAnnotationSerializer.java

@@ -25,6 +25,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.UUID;
 import javax.lang.model.element.AnnotationMirror;
 import javax.lang.model.element.Modifier;
@@ -54,7 +55,8 @@ public record MethodInfo(
       String name,
       boolean nullMarked,
       boolean nullUnmarked,
-      List<TypeParamInfo> typeParams) {}
+      List<TypeParamInfo> typeParams,
+      Map<Integer, Set<NestedAnnotationInfo>> nestedAnnotationsList) {}
 
   public record ClassInfo(
       String name,
@@ -153,22 +155,32 @@ public void finished(com.sun.source.util.TaskEvent e) {
                     return super.visitMethod(methodTree, null);
                   }
                   boolean methodHasAnnotations = false;
+                  Map<Integer, Set<NestedAnnotationInfo>> nestedAnnotationsMap = new HashMap<>();
                   String returnType = "";
                   if (methodTree.getReturnType() != null) {
                     returnType += mSym.getReturnType().toString();
                     if (hasJSpecifyAnnotationDeep(mSym.getReturnType())) {
                       methodHasAnnotations = true;
+                      Set<NestedAnnotationInfo> nested =
+                          mSym.getReturnType()
+                              .accept(new CreateNestedAnnotationInfoVisitor(), null);
+                      if (nested != null && !nested.isEmpty()) {
+                        nestedAnnotationsMap.put(-1, nested);
+                      }
                     }
                   }
                   boolean hasNullMarked = hasAnnotation(mSym, NULLMARKED_NAME);
                   boolean hasNullUnmarked = hasAnnotation(mSym, NULLUNMARKED_NAME);
                   methodHasAnnotations = methodHasAnnotations || hasNullMarked || hasNullUnmarked;
                   // check each parameter annotations
-                  if (!methodHasAnnotations) {
-                    for (Symbol.VarSymbol vSym : mSym.getParameters()) {
-                      if (hasJSpecifyAnnotationDeep(vSym.asType())) {
-                        methodHasAnnotations = true;
-                        break;
+                  for (int idx = 0; idx < mSym.getParameters().size(); idx++) {
+                    Symbol.VarSymbol vSym = mSym.getParameters().get(idx);
+                    if (hasJSpecifyAnnotationDeep(vSym.asType())) {
+                      methodHasAnnotations = true;
+                      Set<NestedAnnotationInfo> nested =
+                          vSym.asType().accept(new CreateNestedAnnotationInfoVisitor(), null);
+                      if (nested != null && !nested.isEmpty()) {
+                        nestedAnnotationsMap.put(idx, nested);
                       }
                     }
                   }
@@ -185,7 +197,8 @@ public void finished(com.sun.source.util.TaskEvent e) {
                           mSym.toString(),
                           hasNullMarked,
                           hasNullUnmarked,
-                          methodTypeParams);
+                          methodTypeParams,
+                          nestedAnnotationsMap);
                   // only add this method if it uses JSpecify annotations
                   if (currentClass != null && methodHasAnnotations) {
                     currentClass.methods().add(methodInfo);