Reason about key set iteration for subtypes of Map (#559)

Fixes #558 

Most of the changes here are to thread around a `VisitorState` instead of a `Types` object, which allows us to use recommended Error Prone APIs to get a reference to a `Type` object and do subtype tests.

7 files changed, 116 insertions, 71 deletions

diff --git a/nullaway/src/main/java/com/uber/nullaway/NullAway.java b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
index 2c29da1..2070ff3 100644
--- a/nullaway/src/main/java/com/uber/nullaway/NullAway.java
+++ b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
@@ -2027,8 +2027,7 @@ public class NullAway extends BugChecker
   }
 
   public AccessPathNullnessAnalysis getNullnessAnalysis(VisitorState state) {
-    return AccessPathNullnessAnalysis.instance(
-        state.context, nonAnnotatedMethod, config, this.handler);
+    return AccessPathNullnessAnalysis.instance(state, nonAnnotatedMethod, config, this.handler);
   }
 
   private boolean mayBeNullFieldAccess(VisitorState state, ExpressionTree expr, Symbol exprSymbol) {
diff --git a/nullaway/src/main/java/com/uber/nullaway/NullabilityUtil.java b/nullaway/src/main/java/com/uber/nullaway/NullabilityUtil.java
index 758b5c5..35f056a 100644
--- a/nullaway/src/main/java/com/uber/nullaway/NullabilityUtil.java
+++ b/nullaway/src/main/java/com/uber/nullaway/NullabilityUtil.java
@@ -24,6 +24,9 @@ package com.uber.nullaway;
 
 import com.google.common.base.Preconditions;
 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.BlockTree;
 import com.sun.source.tree.ClassTree;
@@ -55,6 +58,8 @@ import org.checkerframework.nullaway.javacutil.AnnotationUtils;
 /** Helpful utility methods for nullability analysis. */
 public class NullabilityUtil {
 
+  private static final Supplier<Type> MAP_TYPE_SUPPLIER = Suppliers.typeFromString("java.util.Map");
+
   private NullabilityUtil() {}
 
   /**
@@ -360,4 +365,20 @@ public class NullabilityUtil {
     Symbol.ClassSymbol outermostClassSymbol = getOutermostClassSymbol(symbol);
     return ASTHelpers.hasDirectAnnotationWithSimpleName(outermostClassSymbol, "Generated");
   }
+
+  /**
+   * Checks if {@code symbol} is a method on {@code java.util.Map} (or a subtype) with name {@code
+   * methodName} and {@code numParams} parameters
+   */
+  public static boolean isMapMethod(
+      Symbol.MethodSymbol symbol, VisitorState state, String methodName, int numParams) {
+    if (!symbol.getSimpleName().toString().equals(methodName)) {
+      return false;
+    }
+    if (symbol.getParameters().size() != numParams) {
+      return false;
+    }
+    Symbol owner = symbol.owner;
+    return ASTHelpers.isSubtype(owner.type, MAP_TYPE_SUPPLIER.get(state), state);
+  }
 }
diff --git a/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPath.java b/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPath.java
index b427b05..96f4141 100644
--- a/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPath.java
+++ b/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPath.java
@@ -25,13 +25,14 @@ package com.uber.nullaway.dataflow;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import com.google.errorprone.VisitorState;
 import com.google.errorprone.util.ASTHelpers;
 import com.sun.source.tree.LiteralTree;
 import com.sun.source.tree.MethodInvocationTree;
 import com.sun.source.tree.Tree;
 import com.sun.tools.javac.code.Symbol;
 import com.sun.tools.javac.code.Type;
-import com.sun.tools.javac.code.Types;
+import com.uber.nullaway.NullabilityUtil;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -162,8 +163,8 @@ public final class AccessPath implements MapKey {
    */
   @Nullable
   static AccessPath fromMethodCall(
-      MethodInvocationNode node, @Nullable Types types, AccessPathContext apContext) {
-    if (types != null && isMapGet(ASTHelpers.getSymbol(node.getTree()), types)) {
+      MethodInvocationNode node, @Nullable VisitorState state, AccessPathContext apContext) {
+    if (state != null && isMapGet(ASTHelpers.getSymbol(node.getTree()), state)) {
       return fromMapGetCall(node, apContext);
     }
     return fromVanillaMethodCall(node, apContext);
@@ -325,20 +326,20 @@ public final class AccessPath implements MapKey {
    * </code>
    *
    * @param node AST node
-   * @param types javac {@link Types}
+   * @param state the visitor state
    * @param apContext the current access path context information (see {@link
    *     AccessPath.AccessPathContext}).
    * @return corresponding AccessPath if it exists; <code>null</code> otherwise
    */
   @Nullable
   public static AccessPath getAccessPathForNodeWithMapGet(
-      Node node, @Nullable Types types, AccessPathContext apContext) {
+      Node node, @Nullable VisitorState state, AccessPathContext apContext) {
     if (node instanceof LocalVariableNode) {
       return fromLocal((LocalVariableNode) node);
     } else if (node instanceof FieldAccessNode) {
       return fromFieldAccess((FieldAccessNode) node, apContext);
     } else if (node instanceof MethodInvocationNode) {
-      return fromMethodCall((MethodInvocationNode) node, types, apContext);
+      return fromMethodCall((MethodInvocationNode) node, state, apContext);
     } else {
       return null;
     }
@@ -547,37 +548,16 @@ public final class AccessPath implements MapKey {
     return "AccessPath{" + "root=" + root + ", elements=" + elements + '}';
   }
 
-  private static boolean isMapMethod(
-      Symbol.MethodSymbol symbol, Types types, String methodName, int numParams) {
-    if (!symbol.getSimpleName().toString().equals(methodName)) {
-      return false;
-    }
-    if (symbol.getParameters().size() != numParams) {
-      return false;
-    }
-    Symbol owner = symbol.owner;
-    if (owner.getQualifiedName().toString().equals("java.util.Map")) {
-      return true;
-    }
-    com.sun.tools.javac.util.List<Type> supertypes = types.closure(owner.type);
-    for (Type t : supertypes) {
-      if (t.asElement().getQualifiedName().toString().equals("java.util.Map")) {
-        return true;
-      }
-    }
-    return false;
-  }
-
-  private static boolean isMapGet(Symbol.MethodSymbol symbol, Types types) {
-    return isMapMethod(symbol, types, "get", 1);
+  private static boolean isMapGet(Symbol.MethodSymbol symbol, VisitorState state) {
+    return NullabilityUtil.isMapMethod(symbol, state, "get", 1);
   }
 
-  public static boolean isContainsKey(Symbol.MethodSymbol symbol, Types types) {
-    return isMapMethod(symbol, types, "containsKey", 1);
+  public static boolean isContainsKey(Symbol.MethodSymbol symbol, VisitorState state) {
+    return NullabilityUtil.isMapMethod(symbol, state, "containsKey", 1);
   }
 
-  public static boolean isMapPut(Symbol.MethodSymbol symbol, Types types) {
-    return isMapMethod(symbol, types, "put", 2);
+  public static boolean isMapPut(Symbol.MethodSymbol symbol, VisitorState state) {
+    return NullabilityUtil.isMapMethod(symbol, state, "put", 2);
   }
 
   /**
diff --git a/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessAnalysis.java b/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessAnalysis.java
index 7035a4d..6576472 100644
--- a/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessAnalysis.java
+++ b/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessAnalysis.java
@@ -62,7 +62,7 @@ public final class AccessPathNullnessAnalysis {
   // Use #instance to instantiate
   private AccessPathNullnessAnalysis(
       Predicate<MethodInvocationNode> methodReturnsNonNull,
-      Context context,
+      VisitorState state,
       Config config,
       Handler handler) {
     apContext =
@@ -73,7 +73,7 @@ public final class AccessPathNullnessAnalysis {
         new AccessPathNullnessPropagation(
             Nullness.NONNULL,
             methodReturnsNonNull,
-            context,
+            state,
             apContext,
             config,
             handler,
@@ -85,7 +85,7 @@ public final class AccessPathNullnessAnalysis {
           new AccessPathNullnessPropagation(
               Nullness.NONNULL,
               methodReturnsNonNull,
-              context,
+              state,
               apContext,
               config,
               handler,
@@ -96,20 +96,21 @@ public final class AccessPathNullnessAnalysis {
   /**
    * Get the per-Javac instance of the analysis.
    *
-   * @param context Javac context
+   * @param state visitor state for the compilation
    * @param methodReturnsNonNull predicate determining whether a method is assumed to return NonNull
    *     value
    * @param config analysis config
    * @return instance of the analysis
    */
   public static AccessPathNullnessAnalysis instance(
-      Context context,
+      VisitorState state,
       Predicate<MethodInvocationNode> methodReturnsNonNull,
       Config config,
       Handler handler) {
+    Context context = state.context;
     AccessPathNullnessAnalysis instance = context.get(FIELD_NULLNESS_ANALYSIS_KEY);
     if (instance == null) {
-      instance = new AccessPathNullnessAnalysis(methodReturnsNonNull, context, config, handler);
+      instance = new AccessPathNullnessAnalysis(methodReturnsNonNull, state, config, handler);
       context.put(FIELD_NULLNESS_ANALYSIS_KEY, instance);
     }
     return instance;
diff --git a/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessPropagation.java b/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessPropagation.java
index 9e111fb..3cf5a79 100644
--- a/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessPropagation.java
+++ b/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessPropagation.java
@@ -23,11 +23,13 @@ import static javax.lang.model.element.ElementKind.EXCEPTION_PARAMETER;
 import static org.checkerframework.nullaway.javacutil.TreeUtils.elementFromDeclaration;
 
 import com.google.common.base.Preconditions;
+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.tools.javac.code.Symbol;
+import com.sun.tools.javac.code.Type;
 import com.sun.tools.javac.code.TypeTag;
-import com.sun.tools.javac.code.Types;
-import com.sun.tools.javac.util.Context;
 import com.uber.nullaway.Config;
 import com.uber.nullaway.NullabilityUtil;
 import com.uber.nullaway.Nullness;
@@ -132,16 +134,19 @@ public class AccessPathNullnessPropagation
 
   private static final boolean NO_STORE_CHANGE = false;
 
+  private static final Supplier<Type> SET_TYPE_SUPPLIER = Suppliers.typeFromString("java.util.Set");
+
+  private static final Supplier<Type> ITERATOR_TYPE_SUPPLIER =
+      Suppliers.typeFromString("java.util.Iterator");
+
   private final Nullness defaultAssumption;
 
   private final Predicate<MethodInvocationNode> methodReturnsNonNull;
 
-  private final Context context;
+  private final VisitorState state;
 
   private final AccessPath.AccessPathContext apContext;
 
-  private final Types types;
-
   private final Config config;
 
   private final Handler handler;
@@ -151,16 +156,15 @@ public class AccessPathNullnessPropagation
   public AccessPathNullnessPropagation(
       Nullness defaultAssumption,
       Predicate<MethodInvocationNode> methodReturnsNonNull,
-      Context context,
+      VisitorState state,
       AccessPath.AccessPathContext apContext,
       Config config,
       Handler handler,
       NullnessStoreInitializer nullnessStoreInitializer) {
     this.defaultAssumption = defaultAssumption;
     this.methodReturnsNonNull = methodReturnsNonNull;
-    this.context = context;
+    this.state = state;
     this.apContext = apContext;
-    this.types = Types.instance(context);
     this.config = config;
     this.handler = handler;
     this.nullnessStoreInitializer = nullnessStoreInitializer;
@@ -189,7 +193,7 @@ public class AccessPathNullnessPropagation
   public NullnessStore initialStore(
       UnderlyingAST underlyingAST, List<LocalVariableNode> parameters) {
     return nullnessStoreInitializer.getInitialStore(
-        underlyingAST, parameters, handler, context, types, config);
+        underlyingAST, parameters, handler, state.context, state.getTypes(), config);
   }
 
   @Override
@@ -444,14 +448,14 @@ public class AccessPathNullnessPropagation
     Node realLeftNode = unwrapAssignExpr(leftNode);
     Node realRightNode = unwrapAssignExpr(rightNode);
 
-    AccessPath leftAP = AccessPath.getAccessPathForNodeWithMapGet(realLeftNode, types, apContext);
+    AccessPath leftAP = AccessPath.getAccessPathForNodeWithMapGet(realLeftNode, state, apContext);
     if (leftAP != null) {
       equalBranchUpdates.set(leftAP, equalBranchValue);
       notEqualBranchUpdates.set(
           leftAP, leftVal.greatestLowerBound(rightVal.deducedValueWhenNotEqual()));
     }
 
-    AccessPath rightAP = AccessPath.getAccessPathForNodeWithMapGet(realRightNode, types, apContext);
+    AccessPath rightAP = AccessPath.getAccessPathForNodeWithMapGet(realRightNode, state, apContext);
     if (rightAP != null) {
       equalBranchUpdates.set(rightAP, equalBranchValue);
       notEqualBranchUpdates.set(
@@ -566,9 +570,10 @@ public class AccessPathNullnessPropagation
             AccessPath.mapWithIteratorContentsKey(mapNode, lhs, apContext);
         if (mapWithIteratorContentsKey != null) {
           // put sanity check here to minimize perf impact
-          if (!isCallToMethod(rhsInv, "java.util.Set", "iterator")) {
+          if (!isCallToMethod(rhsInv, SET_TYPE_SUPPLIER, "iterator")) {
             throw new RuntimeException(
-                "expected call to iterator(), instead saw " + rhsInv.getTarget().getMethod());
+                "expected call to iterator(), instead saw "
+                    + state.getSourceForNode(rhsInv.getTree()));
           }
           updates.set(mapWithIteratorContentsKey, NONNULL);
         }
@@ -590,9 +595,10 @@ public class AccessPathNullnessPropagation
                 .getMapGetIteratorContentsAccessPath((LocalVariableNode) receiver);
         if (mapGetPath != null) {
           // put sanity check here to minimize perf impact
-          if (!isCallToMethod(methodInv, "java.util.Iterator", "next")) {
+          if (!isCallToMethod(methodInv, ITERATOR_TYPE_SUPPLIER, "next")) {
             throw new RuntimeException(
-                "expected call to iterator(), instead saw " + methodInv.getTarget().getMethod());
+                "expected call to next(), instead saw "
+                    + state.getSourceForNode(methodInv.getTree()));
           }
           updates.set(AccessPath.replaceMapKey(mapGetPath, AccessPath.fromLocal(lhs)), NONNULL);
         }
@@ -611,7 +617,8 @@ public class AccessPathNullnessPropagation
     if (receiver instanceof MethodInvocationNode) {
       MethodInvocationNode baseInvocation = (MethodInvocationNode) receiver;
       // Check for a call to java.util.Map.keySet()
-      if (isCallToMethod(baseInvocation, "java.util.Map", "keySet")) {
+      if (NullabilityUtil.isMapMethod(
+          ASTHelpers.getSymbol(baseInvocation.getTree()), state, "keySet", 0)) {
         // receiver represents the map
         return baseInvocation.getTarget().getReceiver();
       }
@@ -620,11 +627,13 @@ public class AccessPathNullnessPropagation
   }
 
   private boolean isCallToMethod(
-      MethodInvocationNode invocationNode, String containingClassName, String methodName) {
+      MethodInvocationNode invocationNode,
+      Supplier<Type> containingTypeSupplier,
+      String methodName) {
     Symbol.MethodSymbol symbol = ASTHelpers.getSymbol(invocationNode.getTree());
     return symbol != null
         && symbol.getSimpleName().contentEquals(methodName)
-        && symbol.owner.getQualifiedName().contentEquals(containingClassName);
+        && ASTHelpers.isSubtype(symbol.owner.type, containingTypeSupplier.get(state), state);
   }
 
   /**
@@ -646,7 +655,7 @@ public class AccessPathNullnessPropagation
    * the updates
    */
   private void setNonnullIfAnalyzeable(Updates updates, Node node) {
-    AccessPath ap = AccessPath.getAccessPathForNodeWithMapGet(node, types, apContext);
+    AccessPath ap = AccessPath.getAccessPathForNodeWithMapGet(node, state, apContext);
     if (ap != null) {
       updates.set(ap, NONNULL);
     }
@@ -867,10 +876,17 @@ public class AccessPathNullnessPropagation
     Preconditions.checkNotNull(callee);
     setReceiverNonnull(bothUpdates, node.getTarget().getReceiver(), callee);
     setNullnessForMapCalls(
-        node, callee, node.getArguments(), types, values(input), thenUpdates, bothUpdates);
+        node, callee, node.getArguments(), values(input), thenUpdates, bothUpdates);
     NullnessHint nullnessHint =
         handler.onDataflowVisitMethodInvocation(
-            node, types, context, apContext, values(input), thenUpdates, elseUpdates, bothUpdates);
+            node,
+            state.getTypes(),
+            state.context,
+            apContext,
+            values(input),
+            thenUpdates,
+            elseUpdates,
+            bothUpdates);
     Nullness nullness = returnValueNullness(node, input, nullnessHint);
     if (booleanReturnType(node)) {
       ResultingStore thenStore = updateStore(input.getThenStore(), thenUpdates, bothUpdates);
@@ -885,11 +901,10 @@ public class AccessPathNullnessPropagation
       MethodInvocationNode node,
       Symbol.MethodSymbol callee,
       List<Node> arguments,
-      Types types,
       AccessPathNullnessPropagation.SubNodeValues inputs,
       AccessPathNullnessPropagation.Updates thenUpdates,
       AccessPathNullnessPropagation.Updates bothUpdates) {
-    if (AccessPath.isContainsKey(callee, types)) {
+    if (AccessPath.isContainsKey(callee, state)) {
       // make sure argument is a variable, and get its element
       AccessPath getAccessPath = AccessPath.getForMapInvocation(node, apContext);
       if (getAccessPath != null) {
@@ -898,7 +913,7 @@ public class AccessPathNullnessPropagation
         // as containsKey()
         thenUpdates.set(getAccessPath, NONNULL);
       }
-    } else if (AccessPath.isMapPut(callee, types)) {
+    } else if (AccessPath.isMapPut(callee, state)) {
       AccessPath getAccessPath = AccessPath.getForMapInvocation(node, apContext);
       if (getAccessPath != null) {
         Nullness value = inputs.valueOfSubNode(arguments.get(1));
@@ -928,7 +943,7 @@ public class AccessPathNullnessPropagation
     } else if (node != null && returnValueNullnessHint == NullnessHint.HINT_NULLABLE) {
       // we have a model saying return value is nullable.
       // still, rely on dataflow fact if there is one available
-      nullness = input.getRegularStore().valueOfMethodCall(node, types, NULLABLE, apContext);
+      nullness = input.getRegularStore().valueOfMethodCall(node, state, NULLABLE, apContext);
     } else if (node == null
         || methodReturnsNonNull.test(node)
         || !Nullness.hasNullableAnnotation((Symbol) node.getTarget().getMethod(), config)) {
@@ -936,7 +951,7 @@ public class AccessPathNullnessPropagation
       nullness = NONNULL;
     } else {
       // rely on dataflow, assuming nullable if no fact
-      nullness = input.getRegularStore().valueOfMethodCall(node, types, NULLABLE, apContext);
+      nullness = input.getRegularStore().valueOfMethodCall(node, state, NULLABLE, apContext);
     }
     return nullness;
   }
@@ -1073,7 +1088,7 @@ public class AccessPathNullnessPropagation
 
     @Override
     public void set(MethodInvocationNode node, Nullness value) {
-      AccessPath path = AccessPath.fromMethodCall(node, types, apContext);
+      AccessPath path = AccessPath.fromMethodCall(node, state, apContext);
       values.put(checkNotNull(path), value);
     }
 
diff --git a/nullaway/src/main/java/com/uber/nullaway/dataflow/NullnessStore.java b/nullaway/src/main/java/com/uber/nullaway/dataflow/NullnessStore.java
index 9068f89..3d53556 100644
--- a/nullaway/src/main/java/com/uber/nullaway/dataflow/NullnessStore.java
+++ b/nullaway/src/main/java/com/uber/nullaway/dataflow/NullnessStore.java
@@ -20,7 +20,7 @@ import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.collect.Sets.intersection;
 
 import com.google.common.collect.ImmutableMap;
-import com.sun.tools.javac.code.Types;
+import com.google.errorprone.VisitorState;
 import com.uber.nullaway.Nullness;
 import com.uber.nullaway.dataflow.AccessPath.IteratorContentsKey;
 import java.util.HashMap;
@@ -97,10 +97,10 @@ public class NullnessStore implements Store<NullnessStore> {
    */
   public Nullness valueOfMethodCall(
       MethodInvocationNode node,
-      Types types,
+      VisitorState state,
       Nullness defaultValue,
       AccessPath.AccessPathContext apContext) {
-    AccessPath accessPath = AccessPath.fromMethodCall(node, types, apContext);
+    AccessPath accessPath = AccessPath.fromMethodCall(node, state, apContext);
     if (accessPath == null) {
       return defaultValue;
     }
diff --git a/nullaway/src/test/java/com/uber/nullaway/NullAwayKeySetIteratorTests.java b/nullaway/src/test/java/com/uber/nullaway/NullAwayKeySetIteratorTests.java
index 8ad60e9..5b39012 100644
--- a/nullaway/src/test/java/com/uber/nullaway/NullAwayKeySetIteratorTests.java
+++ b/nullaway/src/test/java/com/uber/nullaway/NullAwayKeySetIteratorTests.java
@@ -131,4 +131,33 @@ public class NullAwayKeySetIteratorTests extends NullAwayTestsBase {
             "}")
         .doTest();
   }
+
+  @Test
+  public void declaredTypeSubtypeOfMap() {
+    defaultCompilationHelper
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import com.google.common.collect.ImmutableMap;",
+            "import java.util.TreeMap;",
+            "import java.util.LinkedHashMap;",
+            "public class Test {",
+            "  public void keySetStuff1(ImmutableMap<Object, Object> m) {",
+            "    for (Object k: m.keySet()) {",
+            "      m.get(k).toString();",
+            "    }",
+            "  }",
+            "  public void keySetStuff2(TreeMap<Object, Object> m) {",
+            "    for (Object k: m.keySet()) {",
+            "      m.get(k).toString();",
+            "    }",
+            "  }",
+            "  public void keySetStuff3(LinkedHashMap<Object, Object> m) {",
+            "    for (Object k: m.keySet()) {",
+            "      m.get(k).toString();",
+            "    }",
+            "  }",
+            "}")
+        .doTest();
+  }
 }