Fix bugs in handling of valueOf calls for map keys

nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPath.java

@@ -28,6 +28,8 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.VisitorState;
+import com.google.errorprone.suppliers.Supplier;
+import com.google.errorprone.suppliers.Suppliers;
 import com.google.errorprone.util.ASTHelpers;
 import com.sun.source.tree.LiteralTree;
 import com.sun.source.tree.MethodInvocationTree;
@@ -75,6 +77,12 @@
  */
 public final class AccessPath implements MapKey {
 
+  private static final Supplier<Type> INTEGER_TYPE_SUPPLIER =
+      Suppliers.typeFromString("java.lang.Integer");
+
+  private static final Supplier<Type> LONG_TYPE_SUPPLIER =
+      Suppliers.typeFromString("java.lang.Long");
+
   /**
    * A prefix added for elements appearing in method invocation APs which represent fields that can
    * be proven to be class-initialization time constants (i.e. static final fields of a type known
@@ -278,14 +286,15 @@ private static Node stripCasts(Node node) {
         return new NumericMapKey(((LongLiteralNode) argument).getValue());
       case METHOD_INVOCATION:
         MethodAccessNode target = ((MethodInvocationNode) argument).getTarget();
-        Node receiver = stripCasts(target.getReceiver());
         List<Node> arguments = ((MethodInvocationNode) argument).getArguments();
         // Check for int/long boxing.
         if (target.getMethod().getSimpleName().toString().equals("valueOf")
-            && arguments.size() == 1
-            && castToNonNull(receiver.getTree()).getKind().equals(Tree.Kind.IDENTIFIER)
-            && (receiver.toString().equals("Integer") || receiver.toString().equals("Long"))) {
-          return argumentToMapKeySpecifier(arguments.get(0), state, apContext);
+            && arguments.size() == 1) {
+          Type ownerType = ((Symbol.MethodSymbol) target.getMethod()).owner.type;
+          if (ASTHelpers.isSameType(ownerType, INTEGER_TYPE_SUPPLIER.get(state), state)
+              || ASTHelpers.isSameType(ownerType, LONG_TYPE_SUPPLIER.get(state), state)) {
+            return argumentToMapKeySpecifier(arguments.get(0), state, apContext);
+          }
         }
       // Fine to fallthrough:
       default:

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

@@ -434,4 +434,58 @@ public void testAccessUsingExplicitThis() {
             "}")
         .doTest();
   }
+
+  @Test
+  public void mapKeysFromValueOf() {
+    defaultCompilationHelper
+        .addSourceLines(
+            "Foo.java",
+            "package com.uber;",
+            "import java.util.HashMap;",
+            "import java.util.Map;",
+            "public class Foo {",
+            "  private final Map<Integer, Object> map = new HashMap<>();",
+            "  private final Map<Long, Object> longMap = new HashMap<>();",
+            "  static Integer valueOf(int i) { return 0; }",
+            "  static Integer valueOf(int i, int j) { return i+j; }",
+            "  public void putThenGetIntegerValueOf() {",
+            "    map.put(Integer.valueOf(10), new Object());",
+            "    map.get(Integer.valueOf(10)).toString();",
+            "  }",
+            "  public void putThenGetLongValueOf() {",
+            "    longMap.put(Long.valueOf(10), new Object());",
+            "    longMap.get(Long.valueOf(10)).toString();",
+            "  }",
+            "  public void putThenGetFooValueOf() {",
+            "    map.put(valueOf(10), new Object());",
+            "    // Unknown valueOf method so we report a warning",
+            "    // BUG: Diagnostic contains: dereferenced expression map.get(valueOf(10)) is @Nullable",
+            "    map.get(valueOf(10)).toString();",
+            "    map.put(valueOf(10,20), new Object());",
+            "    // Unknown valueOf method so we report a warning",
+            "    // BUG: Diagnostic contains: dereferenced expression map.get(valueOf(10,20)) is @Nullable",
+            "    map.get(valueOf(10,20)).toString();",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void mapKeyFromIntegerValueOfStaticImport() {
+    defaultCompilationHelper
+        .addSourceLines(
+            "Foo.java",
+            "package com.uber;",
+            "import java.util.HashMap;",
+            "import java.util.Map;",
+            "import static java.lang.Integer.valueOf;",
+            "public class Foo {",
+            "  private final Map<Integer, Object> map = new HashMap<>();",
+            "  public void putThenGet() {",
+            "    map.put(valueOf(10), new Object());",
+            "    map.get(valueOf(10)).toString();",
+            "  }",
+            "}")
+        .doTest();
+  }
 }