Properly model AtomicReference.get() in JSpecify mode

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

@@ -386,7 +386,7 @@ private static LibraryModels loadLibraryModels(Config config) {
     Iterable<LibraryModels> externalLibraryModels =
         ServiceLoader.load(LibraryModels.class, LibraryModels.class.getClassLoader());
     ImmutableSet.Builder<LibraryModels> libModelsBuilder = new ImmutableSet.Builder<>();
-    libModelsBuilder.add(new DefaultLibraryModels()).addAll(externalLibraryModels);
+    libModelsBuilder.add(new DefaultLibraryModels(config)).addAll(externalLibraryModels);
     if (config.isJarInferEnabled()) {
       libModelsBuilder.add(new ExternalStubxLibraryModels());
     }
@@ -815,7 +815,7 @@ private static class DefaultLibraryModels implements LibraryModels {
             .put(methodRef("java.util.Objects", "toString(java.lang.Object,java.lang.String)"), 1)
             .build();
 
-    private static final ImmutableSet<MethodRef> NULLABLE_RETURNS =
+    private static final ImmutableSet<MethodRef> ALWAYS_NULLABLE_RETURNS =
         new ImmutableSet.Builder<MethodRef>()
             .add(methodRef("com.sun.source.tree.CompilationUnitTree", "getPackageName()"))
             .add(methodRef("java.lang.Throwable", "getMessage()"))
@@ -826,7 +826,6 @@ private static class DefaultLibraryModels implements LibraryModels {
             .add(methodRef("java.lang.ref.SoftReference", "get()"))
             .add(methodRef("java.lang.ref.WeakReference", "get()"))
             .add(methodRef("java.nio.file.Path", "getParent()"))
-            .add(methodRef("java.util.concurrent.atomic.AtomicReference", "get()"))
             .add(methodRef("java.util.Map", "get(java.lang.Object)"))
             .add(methodRef("javax.lang.model.element.Element", "getEnclosingElement()"))
             .add(methodRef("javax.lang.model.element.ExecutableElement", "getDefaultValue()"))
@@ -840,6 +839,9 @@ private static class DefaultLibraryModels implements LibraryModels {
             .add(methodRef("java.lang.System", "console()"))
             .build();
 
+    private static final ImmutableSet<MethodRef> NULLABLE_RETURNS_JSPECIFY_MODE_DISABLED =
+        ImmutableSet.of(methodRef("java.util.concurrent.atomic.AtomicReference", "get()"));
+
     private static final ImmutableSet<MethodRef> NONNULL_RETURNS =
         new ImmutableSet.Builder<MethodRef>()
             .add(methodRef("com.google.gson", "<T>fromJson(String,Class)"))
@@ -892,13 +894,32 @@ private static class DefaultLibraryModels implements LibraryModels {
         new ImmutableSetMultimap.Builder<String, Integer>()
             .put("java.util.function.Function", 0)
             .put("java.util.function.Function", 1)
+            .put("java.util.concurrent.atomic.AtomicReference", 0)
             .build();
 
     private static final ImmutableSet<String> NULLMARKED_CLASSES =
-        new ImmutableSet.Builder<String>().add("java.util.function.Function").build();
+        new ImmutableSet.Builder<String>()
+            .add("java.util.function.Function")
+            .add("java.util.concurrent.atomic.AtomicReference")
+            .build();
+
     private static final ImmutableSetMultimap<MethodRef, Integer> CAST_TO_NONNULL_METHODS =
         new ImmutableSetMultimap.Builder<MethodRef, Integer>().build();
 
+    /**
+     * Methods with nullable returns for the current configuration (JSpecify mode enabled or not).
+     */
+    private final ImmutableSet<MethodRef> nullableReturnsForConfig;
+
+    DefaultLibraryModels(Config config) {
+      ImmutableSet.Builder<MethodRef> builder = new ImmutableSet.Builder<>();
+      builder.addAll(ALWAYS_NULLABLE_RETURNS);
+      if (!config.isJSpecifyMode()) {
+        builder.addAll(NULLABLE_RETURNS_JSPECIFY_MODE_DISABLED);
+      }
+      nullableReturnsForConfig = builder.build();
+    }
+
     @Override
     public ImmutableSetMultimap<MethodRef, Integer> failIfNullParameters() {
       return FAIL_IF_NULL_PARAMETERS;
@@ -931,7 +952,7 @@ public ImmutableSetMultimap<MethodRef, Integer> nullImpliesNullParameters() {
 
     @Override
     public ImmutableSet<MethodRef> nullableReturns() {
-      return NULLABLE_RETURNS;
+      return nullableReturnsForConfig;
     }
 
     @Override

nullaway/src/test/java/com/uber/nullaway/jspecify/JSpecifyLibraryModelsTests.java

@@ -0,0 +1,37 @@
+package com.uber.nullaway.jspecify;
+
+import com.google.errorprone.CompilationTestHelper;
+import com.uber.nullaway.NullAwayTestsBase;
+import java.util.Arrays;
+import org.junit.Test;
+
+public class JSpecifyLibraryModelsTests extends NullAwayTestsBase {
+
+  @Test
+  public void atomicReferenceGet() {
+    makeHelper()
+        .addSourceLines(
+            "Test.java",
+            "import org.jspecify.annotations.*;",
+            "import java.util.concurrent.atomic.AtomicReference;",
+            "@NullMarked",
+            "class Test {",
+            "  void testNegative() {",
+            "    AtomicReference<Integer> x = new AtomicReference<>(Integer.valueOf(3));",
+            "    x.get().hashCode();",
+            "  }",
+            "  void testPositive() {",
+            "    AtomicReference<@Nullable Integer> x = new AtomicReference<>(Integer.valueOf(3));",
+            "    // BUG: Diagnostic contains: dereferenced expression x.get() is @Nullable",
+            "    x.get().hashCode();",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  private CompilationTestHelper makeHelper() {
+    return makeTestHelperWithArgs(
+        Arrays.asList(
+            "-XepOpt:NullAway:OnlyNullMarked=true", "-XepOpt:NullAway:JSpecifyMode=true"));
+  }
+}