Option that loads nullness information of JDK package

nullaway/src/main/java/com/uber/nullaway/Config.java

@@ -285,6 +285,13 @@ public interface Config {
    */
   boolean isJarInferEnabled();
 
+  /**
+   * Checks if <a href="https://github.com/jspecify/jdk">JSpecify JDK models</a> should be enabled.
+   *
+   * @return true if JSpecify JDK models should be enabled, false otherwise
+   */
+  boolean isJSpecifyJDKModels();
+
   /**
    * Gets the URL to show with NullAway error messages.
    *

nullaway/src/main/java/com/uber/nullaway/DummyOptionsConfig.java

@@ -204,6 +204,11 @@ public boolean isJarInferEnabled() {
     throw new IllegalStateException(ERROR_MESSAGE);
   }
 
+  @Override
+  public boolean isJSpecifyJDKModels() {
+    throw new IllegalStateException(ERROR_MESSAGE);
+  }
+
   @Override
   public String getErrorURL() {
     throw new IllegalStateException(ERROR_MESSAGE);

nullaway/src/main/java/com/uber/nullaway/ErrorProneCLIFlagsConfig.java

@@ -87,6 +87,8 @@ final class ErrorProneCLIFlagsConfig implements Config {
   /** --- JarInfer configs --- */
   static final String FL_JI_ENABLED = EP_FL_NAMESPACE + ":JarInferEnabled";
 
+  static final String FL_JSPECIFY_JDK_ENABLED = EP_FL_NAMESPACE + ":JSpecifyJDKModels";
+
   static final String FL_ERROR_URL = EP_FL_NAMESPACE + ":ErrorURL";
 
   /** --- Serialization configs --- */
@@ -247,6 +249,8 @@ final class ErrorProneCLIFlagsConfig implements Config {
   /** --- JarInfer configs --- */
   private final boolean jarInferEnabled;
 
+  private final boolean jspecifyJDKModelsEnabled;
+
   private final String errorURL;
 
   /** --- Fully qualified names of custom nonnull/nullable annotation --- */
@@ -333,6 +337,7 @@ final class ErrorProneCLIFlagsConfig implements Config {
 
     /* --- JarInfer configs --- */
     jarInferEnabled = flags.getBoolean(FL_JI_ENABLED).orElse(false);
+    jspecifyJDKModelsEnabled = flags.getBoolean(FL_JSPECIFY_JDK_ENABLED).orElse(false);
     errorURL = flags.get(FL_ERROR_URL).orElse(DEFAULT_URL);
     if (acknowledgeAndroidRecent && !isAcknowledgeRestrictive) {
       throw new IllegalStateException(
@@ -589,6 +594,11 @@ public boolean isJarInferEnabled() {
     return jarInferEnabled;
   }
 
+  @Override
+  public boolean isJSpecifyJDKModels() {
+    return jspecifyJDKModelsEnabled;
+  }
+
   @Override
   public String getErrorURL() {
     return errorURL;

nullaway/src/main/java/com/uber/nullaway/handlers/LibraryModelsHandler.java

@@ -478,8 +478,9 @@ private static LibraryModels loadLibraryModels(Config config) {
         ServiceLoader.load(LibraryModels.class, LibraryModels.class.getClassLoader());
     ImmutableSet.Builder<LibraryModels> libModelsBuilder = new ImmutableSet.Builder<>();
     libModelsBuilder.add(new DefaultLibraryModels(config)).addAll(externalLibraryModels);
-    if (config.isJarInferEnabled()) {
-      libModelsBuilder.add(new ExternalStubxLibraryModels());
+    if (config.isJarInferEnabled() || config.isJSpecifyJDKModels()) {
+      libModelsBuilder.add(
+          new ExternalStubxLibraryModels(config.isJarInferEnabled(), config.isJSpecifyJDKModels()));
     }
     return new CombinedLibraryModels(libModelsBuilder.build(), config);
   }
@@ -1560,26 +1561,42 @@ private static class ExternalStubxLibraryModels implements LibraryModels {
     private final Multimap<String, Integer> methodTypeParamNullableUpperBoundCache;
     private final Map<String, SetMultimap<Integer, NestedAnnotationInfo>> nestedAnnotationInfo;
 
-    ExternalStubxLibraryModels() {
+    ExternalStubxLibraryModels(boolean isJarInferEnabled, boolean isJSpecifyJDKEnabled) {
       String libraryModelLogName = "LM";
       StubxCacheUtil cacheUtil = new StubxCacheUtil(libraryModelLogName);
       // hardcoded loading of stubx files from android-jarinfer-models-sdkXX artifacts
-      try {
-        InputStream androidStubxIS =
-            Class.forName(ANDROID_MODEL_CLASS)
-                .getClassLoader()
-                .getResourceAsStream(ANDROID_ASTUBX_LOCATION);
-        if (androidStubxIS != null) {
-          cacheUtil.parseStubStream(androidStubxIS, "android.jar: " + ANDROID_ASTUBX_LOCATION);
-          astubxLoadLog("Loaded Android RT models.");
+      if (isJarInferEnabled) {
+        try {
+          InputStream androidStubxIS =
+              Class.forName(ANDROID_MODEL_CLASS)
+                  .getClassLoader()
+                  .getResourceAsStream(ANDROID_ASTUBX_LOCATION);
+          if (androidStubxIS != null) {
+            cacheUtil.parseStubStream(androidStubxIS, "android.jar: " + ANDROID_ASTUBX_LOCATION);
+            astubxLoadLog("Loaded Android RT models.");
+          }
+        } catch (ClassNotFoundException e) {
+          astubxLoadLog(
+              "Cannot find Android RT models locator class."
+                  + " This is expected if not in an Android project, or the Android SDK JarInfer models Jar has not been set up for this build.");
+
+        } catch (Exception e) {
+          astubxLoadLog("Cannot load Android RT models.");
         }
-      } catch (ClassNotFoundException e) {
-        astubxLoadLog(
-            "Cannot find Android RT models locator class."
-                + " This is expected if not in an Android project, or the Android SDK JarInfer models Jar has not been set up for this build.");
+      }
 
-      } catch (Exception e) {
-        astubxLoadLog("Cannot load Android RT models.");
+      // hardcoded loading of stubx files from jdk nullness inferred output.astubx
+      if (isJSpecifyJDKEnabled) {
+        try (InputStream in = getClass().getClassLoader().getResourceAsStream("output.astubx")) {
+          if (in == null) {
+            astubxLoadLog("JDK astubx model not found on classpath: output.astubx");
+          } else {
+            cacheUtil.parseStubStream(in, "output.astubx");
+            astubxLoadLog("Loaded JDK astubx model.");
+          }
+        } catch (Exception e) {
+          astubxLoadLog("Failed to load JDK astubx: " + e);
+        }
       }
 
       argAnnotCache = cacheUtil.getArgAnnotCache();

nullaway/src/test/java/com/uber/nullaway/JSpecifyJDKModelsTest.java

@@ -0,0 +1,55 @@
+package com.uber.nullaway;
+
+import com.google.errorprone.CompilationTestHelper;
+import java.util.List;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class JSpecifyJDKModelsTest extends NullAwayTestsBase {
+
+  @Test
+  public void modelsEnabledLoadsAstubxModel() {
+    CompilationTestHelper compilationTestHelper =
+        makeTestHelperWithArgs(
+                List.of(
+                    "-XepOpt:NullAway:AnnotatedPackages=foo",
+                    "-XepOpt:NullAway:JSpecifyJDKModels=true"))
+            .addSourceLines(
+                "Test.java",
+                """
+                package foo;
+                import javax.naming.directory.Attributes;
+                import org.jspecify.annotations.NullMarked;
+                @NullMarked
+                class Test {
+                  void use(Attributes attrs) {
+                    // BUG: Diagnostic contains: @Nullable
+                    attrs.get("key").toString();
+                  }
+                }
+                """);
+    compilationTestHelper.doTest();
+  }
+
+  @Test
+  public void modelsDisabledDoesNotLoadAstubxModel() {
+    CompilationTestHelper compilationTestHelper =
+        makeTestHelperWithArgs(List.of("-XepOpt:NullAway:AnnotatedPackages=foo"))
+            .addSourceLines(
+                "Test.java",
+                """
+                package foo;
+                import javax.naming.directory.Attributes;
+                import org.jspecify.annotations.NullMarked;
+                @NullMarked
+                class Test {
+                  void use(Attributes attrs) {
+                    attrs.get("key").toString();
+                  }
+                }
+                """);
+    compilationTestHelper.doTest();
+  }
+}