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,25 @@ 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) -> {
+                      // Get type inside the list
+                      Type[] typeArgs = ((ParameterizedType) type).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 ImmutableList.copyOf(standardList);
+                    })
+            .create();
     Type parsedType = new TypeToken<Map<String, List<ClassInfo>>>() {}.getType();
 
     // parse JSON file

jdk-javac-plugin/build.gradle

@@ -57,6 +57,8 @@ dependencies {
 
     implementation libs.gson
     implementation libs.jspecify
+    implementation project(':nullaway')
+    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,153 @@
+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.librarymodel.NestedAnnotationInfo;
+import com.uber.nullaway.librarymodel.NestedAnnotationInfo.Annotation;
+import com.uber.nullaway.librarymodel.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 ArrayDeque<TypePathEntry> path;
+  private Set<NestedAnnotationInfo> nestedAnnotationInfoList;
+
+  public CreateNestedAnnotationInfoVisitor() {
+    path = new ArrayDeque<>();
+    nestedAnnotationInfoList = new HashSet<>();
+  }
+
+  @Override
+  public @Nullable Set<NestedAnnotationInfo> visitClassType(
+      Type.ClassType classType, @Nullable Void unused) {
+    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) {
+          nestedAnnotationInfoList.add(nestedAnnotationInfo);
+        }
+        typeArg.accept(this, null);
+        path.removeLast();
+      }
+    }
+    return nestedAnnotationInfoList;
+  }
+
+  @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) {
+      nestedAnnotationInfoList.add(nestedAnnotationInfo);
+    }
+    t.elemtype.accept(this, null);
+    path.removeLast();
+    return nestedAnnotationInfoList;
+  }
+
+  @Override
+  public @Nullable Set<NestedAnnotationInfo> visitWildcardType(
+      Type.WildcardType t, @Nullable Void unused) {
+    // Upper Bound (? extends T)
+    NestedAnnotationInfo nestedAnnotationInfo = null;
+    if (t.getExtendsBound() != 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) {
+        nestedAnnotationInfoList.add(nestedAnnotationInfo);
+      }
+      t.getExtendsBound().accept(this, null);
+      path.removeLast();
+    }
+
+    // Lower Bound (? super T)
+    if (t.getSuperBound() != 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) {
+        nestedAnnotationInfoList.add(nestedAnnotationInfo);
+      }
+      t.getSuperBound().accept(this, null);
+      path.removeLast();
+    }
+
+    return nestedAnnotationInfoList;
+  }
+
+  @Override
+  public @Nullable Set<NestedAnnotationInfo> visitType(Type t, @Nullable Void unused) {
+    // just return null (no nested types)
+    return null;
+  }
+
+  private boolean hasNullableAnnotation(TypeMirror type) {
+    if (type == null) {
+      return false;
+    }
+    for (AnnotationMirror annotation : type.getAnnotationMirrors()) {
+      String qualifiedName =
+          ((TypeElement) annotation.getAnnotationType().asElement()).getQualifiedName().toString();
+      if (qualifiedName.equals("org.jspecify.annotations.Nullable")) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  private boolean hasNonNullAnnotation(TypeMirror type) {
+    if (type == null) {
+      return false;
+    }
+    for (AnnotationMirror annotation : type.getAnnotationMirrors()) {
+      String qualifiedName =
+          ((TypeElement) annotation.getAnnotationType().asElement()).getQualifiedName().toString();
+      if (qualifiedName.equals("org.jspecify.annotations.NonNull")) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  private ImmutableList<TypePathEntry> getTypePath() {
+    return ImmutableList.copyOf(path);
+  }
+}

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

@@ -14,6 +14,7 @@
 import com.sun.tools.javac.code.Symbol;
 import com.sun.tools.javac.code.Symbol.ClassSymbol;
 import com.sun.tools.javac.code.Symbol.MethodSymbol;
+import com.uber.nullaway.librarymodel.NestedAnnotationInfo;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -25,6 +26,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 +56,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 +156,31 @@ 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);