[Refactor] Pass resolved Symbols into Handler methods (#729)

This avoids redundant ASTHelpers.getSymbol calls and is more consistent with the other Handler method signatures

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

@@ -2271,7 +2271,7 @@ private boolean mayBeNullExpr(VisitorState state, ExpressionTree expr) {
           return mayBeNullFieldAccess(state, expr, exprSymbol);
         } else {
           // Check handler.onOverrideMayBeNullExpr before dataflow.
-          exprMayBeNull = handler.onOverrideMayBeNullExpr(this, expr, state, true);
+          exprMayBeNull = handler.onOverrideMayBeNullExpr(this, expr, exprSymbol, state, true);
           return exprMayBeNull ? nullnessFromDataflow(state, expr) : false;
         }
       case METHOD_INVOCATION:
@@ -2290,7 +2290,7 @@ private boolean mayBeNullExpr(VisitorState state, ExpressionTree expr) {
               "whoops, better handle " + expr.getKind() + " " + state.getSourceForNode(expr));
         }
     }
-    exprMayBeNull = handler.onOverrideMayBeNullExpr(this, expr, state, exprMayBeNull);
+    exprMayBeNull = handler.onOverrideMayBeNullExpr(this, expr, exprSymbol, state, exprMayBeNull);
     return exprMayBeNull;
   }
 
@@ -2303,7 +2303,7 @@ private boolean mayBeNullMethodCall(
     if (!Nullness.hasNullableAnnotation(exprSymbol, config)) {
       exprMayBeNull = false;
     }
-    exprMayBeNull = handler.onOverrideMayBeNullExpr(this, expr, state, exprMayBeNull);
+    exprMayBeNull = handler.onOverrideMayBeNullExpr(this, expr, exprSymbol, state, exprMayBeNull);
     return exprMayBeNull ? nullnessFromDataflow(state, expr) : false;
   }
 
@@ -2327,7 +2327,7 @@ private boolean mayBeNullFieldAccess(VisitorState state, ExpressionTree expr, Sy
     if (!NullabilityUtil.mayBeNullFieldFromType(exprSymbol, config, codeAnnotationInfo)) {
       exprMayBeNull = false;
     }
-    exprMayBeNull = handler.onOverrideMayBeNullExpr(this, expr, state, exprMayBeNull);
+    exprMayBeNull = handler.onOverrideMayBeNullExpr(this, expr, exprSymbol, state, exprMayBeNull);
     return exprMayBeNull ? nullnessFromDataflow(state, expr) : false;
   }
 

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

@@ -917,15 +917,16 @@ public TransferResult<Nullness, NullnessStore> visitMethodInvocation(
     ReadableUpdates elseUpdates = new ReadableUpdates();
     ReadableUpdates bothUpdates = new ReadableUpdates();
     Symbol.MethodSymbol callee = ASTHelpers.getSymbol(node.getTree());
-    Preconditions.checkNotNull(callee);
+    Preconditions.checkNotNull(
+        callee); // this could be null before https://github.com/google/error-prone/pull/2902
     setReceiverNonnull(bothUpdates, node.getTarget().getReceiver(), callee);
     setNullnessForMapCalls(
         node, callee, node.getArguments(), values(input), thenUpdates, bothUpdates);
     NullnessHint nullnessHint =
         handler.onDataflowVisitMethodInvocation(
-            node, state, apContext, values(input), thenUpdates, elseUpdates, bothUpdates);
+            node, callee, state, apContext, values(input), thenUpdates, elseUpdates, bothUpdates);
     Nullness nullness = returnValueNullness(node, input, nullnessHint);
-    if (booleanReturnType(node)) {
+    if (booleanReturnType(callee)) {
       ResultingStore thenStore = updateStore(input.getThenStore(), thenUpdates, bothUpdates);
       ResultingStore elseStore = updateStore(input.getElseStore(), elseUpdates, bothUpdates);
       return conditionalResult(
@@ -984,9 +985,8 @@ private void setNullnessForMapCalls(
     }
   }
 
-  private boolean booleanReturnType(MethodInvocationNode node) {
-    Symbol.MethodSymbol methodSymbol = ASTHelpers.getSymbol(node.getTree());
-    return methodSymbol != null && methodSymbol.getReturnType().getTag() == TypeTag.BOOLEAN;
+  private static boolean booleanReturnType(Symbol.MethodSymbol methodSymbol) {
+    return methodSymbol.getReturnType().getTag() == TypeTag.BOOLEAN;
   }
 
   Nullness returnValueNullness(

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

@@ -25,7 +25,6 @@
 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.ClassTree;
 import com.sun.tools.javac.code.Symbol;
 import com.sun.tools.javac.code.Type;
@@ -67,13 +66,13 @@ public void onMatchTopLevelClass(
   @Override
   public NullnessHint onDataflowVisitMethodInvocation(
       MethodInvocationNode node,
+      Symbol.MethodSymbol symbol,
       VisitorState state,
       AccessPath.AccessPathContext apContext,
       AccessPathNullnessPropagation.SubNodeValues inputs,
       AccessPathNullnessPropagation.Updates thenUpdates,
       AccessPathNullnessPropagation.Updates elseUpdates,
       AccessPathNullnessPropagation.Updates bothUpdates) {
-    Symbol.MethodSymbol symbol = ASTHelpers.getSymbol(node.getTree());
     if (thriftIsSetCall(symbol, state.getTypes())) {
       String methodName = symbol.getSimpleName().toString();
       // remove "isSet"

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

@@ -45,17 +45,13 @@ public class AssertionHandler extends BaseNoOpHandler {
   @Override
   public NullnessHint onDataflowVisitMethodInvocation(
       MethodInvocationNode node,
+      Symbol.MethodSymbol callee,
       VisitorState state,
       AccessPath.AccessPathContext apContext,
       AccessPathNullnessPropagation.SubNodeValues inputs,
       AccessPathNullnessPropagation.Updates thenUpdates,
       AccessPathNullnessPropagation.Updates elseUpdates,
       AccessPathNullnessPropagation.Updates bothUpdates) {
-    Symbol.MethodSymbol callee = ASTHelpers.getSymbol(node.getTree());
-    if (callee == null) {
-      return NullnessHint.UNKNOWN;
-    }
-
     if (!methodNameUtil.isUtilInitialized()) {
       methodNameUtil.initializeMethodNames(callee.name.table);
     }

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

@@ -130,7 +130,11 @@ public Nullness[] onOverrideMethodInvocationParametersNullability(
 
   @Override
   public boolean onOverrideMayBeNullExpr(
-      NullAway analysis, ExpressionTree expr, VisitorState state, boolean exprMayBeNull) {
+      NullAway analysis,
+      ExpressionTree expr,
+      @Nullable Symbol exprSymbol,
+      VisitorState state,
+      boolean exprMayBeNull) {
     // NoOp
     return exprMayBeNull;
   }
@@ -146,6 +150,7 @@ public NullnessStore.Builder onDataflowInitialStore(
   @Override
   public NullnessHint onDataflowVisitMethodInvocation(
       MethodInvocationNode node,
+      Symbol.MethodSymbol symbol,
       VisitorState state,
       AccessPath.AccessPathContext apContext,
       AccessPathNullnessPropagation.SubNodeValues inputs,

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

@@ -152,9 +152,13 @@ public Nullness[] onOverrideMethodInvocationParametersNullability(
 
   @Override
   public boolean onOverrideMayBeNullExpr(
-      NullAway analysis, ExpressionTree expr, VisitorState state, boolean exprMayBeNull) {
+      NullAway analysis,
+      ExpressionTree expr,
+      @Nullable Symbol exprSymbol,
+      VisitorState state,
+      boolean exprMayBeNull) {
     for (Handler h : handlers) {
-      exprMayBeNull = h.onOverrideMayBeNullExpr(analysis, expr, state, exprMayBeNull);
+      exprMayBeNull = h.onOverrideMayBeNullExpr(analysis, expr, exprSymbol, state, exprMayBeNull);
     }
     return exprMayBeNull;
   }
@@ -173,6 +177,7 @@ public NullnessStore.Builder onDataflowInitialStore(
   @Override
   public NullnessHint onDataflowVisitMethodInvocation(
       MethodInvocationNode node,
+      Symbol.MethodSymbol symbol,
       VisitorState state,
       AccessPath.AccessPathContext apContext,
       AccessPathNullnessPropagation.SubNodeValues inputs,
@@ -183,7 +188,7 @@ public NullnessHint onDataflowVisitMethodInvocation(
     for (Handler h : handlers) {
       NullnessHint n =
           h.onDataflowVisitMethodInvocation(
-              node, state, apContext, inputs, thenUpdates, elseUpdates, bothUpdates);
+              node, symbol, state, apContext, inputs, thenUpdates, elseUpdates, bothUpdates);
       nullnessHint = nullnessHint.merge(n);
     }
     return nullnessHint;

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

@@ -79,14 +79,14 @@ public void onMatchTopLevelClass(
   @Override
   public NullnessHint onDataflowVisitMethodInvocation(
       MethodInvocationNode node,
+      Symbol.MethodSymbol symbol,
       VisitorState state,
       AccessPath.AccessPathContext apContext,
       AccessPathNullnessPropagation.SubNodeValues inputs,
       AccessPathNullnessPropagation.Updates thenUpdates,
       AccessPathNullnessPropagation.Updates elseUpdates,
       AccessPathNullnessPropagation.Updates bothUpdates) {
     MethodInvocationTree tree = castToNonNull(node.getTree());
-    Symbol.MethodSymbol symbol = ASTHelpers.getSymbol(tree);
     Types types = state.getTypes();
     if (grpcIsMetadataContainsKeyCall(symbol, types)) {
       // On seeing o.containsKey(k), set AP for o.get(k) to @NonNull

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

@@ -141,13 +141,18 @@ void onMatchMethodReference(
    *
    * @param analysis A reference to the running NullAway analysis.
    * @param expr The expression in question.
+   * @param exprSymbol The symbol of the expression, might be null
    * @param state The current visitor state.
    * @param exprMayBeNull Whether or not the expression may be null according to the base analysis
    *     or upstream handlers.
    * @return Whether or not the expression may be null, as updated by this handler.
    */
   boolean onOverrideMayBeNullExpr(
-      NullAway analysis, ExpressionTree expr, VisitorState state, boolean exprMayBeNull);
+      NullAway analysis,
+      ExpressionTree expr,
+      @Nullable Symbol exprSymbol,
+      VisitorState state,
+      boolean exprMayBeNull);
 
   /**
    * Called to potentially override the nullability of an annotated or unannotated method's return,
@@ -217,6 +222,7 @@ NullnessStore.Builder onDataflowInitialStore(
    * Called when the Dataflow analysis visits each method invocation.
    *
    * @param node The AST node for the method callsite.
+   * @param symbol The symbol of the called method
    * @param state The current visitor state.
    * @param apContext the current access path context information (see {@link
    *     AccessPath.AccessPathContext}).
@@ -233,6 +239,7 @@ NullnessStore.Builder onDataflowInitialStore(
    */
   NullnessHint onDataflowVisitMethodInvocation(
       MethodInvocationNode node,
+      Symbol.MethodSymbol symbol,
       VisitorState state,
       AccessPath.AccessPathContext apContext,
       AccessPathNullnessPropagation.SubNodeValues inputs,

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

@@ -24,9 +24,7 @@
 
 import com.google.common.base.Preconditions;
 import com.google.errorprone.VisitorState;
-import com.google.errorprone.util.ASTHelpers;
 import com.sun.source.tree.ExpressionTree;
-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;
@@ -178,24 +176,29 @@ public Nullness[] onOverrideMethodInvocationParametersNullability(
   @Override
   public NullnessHint onDataflowVisitMethodInvocation(
       MethodInvocationNode node,
+      Symbol.MethodSymbol symbol,
       VisitorState state,
       AccessPath.AccessPathContext apContext,
       AccessPathNullnessPropagation.SubNodeValues inputs,
       AccessPathNullnessPropagation.Updates thenUpdates,
       AccessPathNullnessPropagation.Updates elseUpdates,
       AccessPathNullnessPropagation.Updates bothUpdates) {
-    if (isReturnAnnotatedNullable(ASTHelpers.getSymbol(node.getTree()))) {
+    if (isReturnAnnotatedNullable(symbol)) {
       return NullnessHint.HINT_NULLABLE;
     }
     return NullnessHint.UNKNOWN;
   }
 
   @Override
   public boolean onOverrideMayBeNullExpr(
-      NullAway analysis, ExpressionTree expr, VisitorState state, boolean exprMayBeNull) {
-    if (expr.getKind().equals(Tree.Kind.METHOD_INVOCATION)) {
-      return exprMayBeNull
-          || isReturnAnnotatedNullable(ASTHelpers.getSymbol((MethodInvocationTree) expr));
+      NullAway analysis,
+      ExpressionTree expr,
+      @Nullable Symbol exprSymbol,
+      VisitorState state,
+      boolean exprMayBeNull) {
+    if (expr.getKind().equals(Tree.Kind.METHOD_INVOCATION)
+        && exprSymbol instanceof Symbol.MethodSymbol) {
+      return exprMayBeNull || isReturnAnnotatedNullable((Symbol.MethodSymbol) exprSymbol);
     }
     return exprMayBeNull;
   }

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

@@ -123,10 +123,15 @@ public Nullness onOverrideMethodInvocationReturnNullability(
 
   @Override
   public boolean onOverrideMayBeNullExpr(
-      NullAway analysis, ExpressionTree expr, VisitorState state, boolean exprMayBeNull) {
-    if (expr.getKind() == Tree.Kind.METHOD_INVOCATION) {
+      NullAway analysis,
+      ExpressionTree expr,
+      @Nullable Symbol exprSymbol,
+      VisitorState state,
+      boolean exprMayBeNull) {
+    if (expr.getKind().equals(Tree.Kind.METHOD_INVOCATION)
+        && exprSymbol instanceof Symbol.MethodSymbol) {
       OptimizedLibraryModels optLibraryModels = getOptLibraryModels(state.context);
-      Symbol.MethodSymbol methodSymbol = (Symbol.MethodSymbol) ASTHelpers.getSymbol(expr);
+      Symbol.MethodSymbol methodSymbol = (Symbol.MethodSymbol) exprSymbol;
       // When looking up library models of annotated code, we match the exact method signature only;
       // overriding methods in subclasses must be explicitly given their own library model.
       // When dealing with unannotated code, we default to generality: a model applies to a method
@@ -183,14 +188,13 @@ private CodeAnnotationInfo getCodeAnnotationInfo(Context context) {
   @Override
   public NullnessHint onDataflowVisitMethodInvocation(
       MethodInvocationNode node,
+      Symbol.MethodSymbol callee,
       VisitorState state,
       AccessPath.AccessPathContext apContext,
       AccessPathNullnessPropagation.SubNodeValues inputs,
       AccessPathNullnessPropagation.Updates thenUpdates,
       AccessPathNullnessPropagation.Updates elseUpdates,
       AccessPathNullnessPropagation.Updates bothUpdates) {
-    Symbol.MethodSymbol callee = ASTHelpers.getSymbol(node.getTree());
-    Preconditions.checkNotNull(callee);
     boolean isMethodAnnotated =
         !getCodeAnnotationInfo(state.context).isSymbolUnannotated(callee, this.config);
     setUnconditionalArgumentNullness(bothUpdates, node.getArguments(), callee, state, apContext);

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

@@ -82,9 +82,14 @@ public class OptionalEmptinessHandler extends BaseNoOpHandler {
 
   @Override
   public boolean onOverrideMayBeNullExpr(
-      NullAway analysis, ExpressionTree expr, VisitorState state, boolean exprMayBeNull) {
-    if (expr.getKind() == Tree.Kind.METHOD_INVOCATION
-        && optionalIsGetCall((Symbol.MethodSymbol) ASTHelpers.getSymbol(expr), state.getTypes())) {
+      NullAway analysis,
+      ExpressionTree expr,
+      @Nullable Symbol exprSymbol,
+      VisitorState state,
+      boolean exprMayBeNull) {
+    if (expr.getKind().equals(Tree.Kind.METHOD_INVOCATION)
+        && exprSymbol instanceof Symbol.MethodSymbol
+        && optionalIsGetCall((Symbol.MethodSymbol) exprSymbol, state.getTypes())) {
       return true;
     }
     return exprMayBeNull;
@@ -109,14 +114,13 @@ public void onMatchTopLevelClass(
   @Override
   public NullnessHint onDataflowVisitMethodInvocation(
       MethodInvocationNode node,
+      Symbol.MethodSymbol symbol,
       VisitorState state,
       AccessPath.AccessPathContext apContext,
       AccessPathNullnessPropagation.SubNodeValues inputs,
       AccessPathNullnessPropagation.Updates thenUpdates,
       AccessPathNullnessPropagation.Updates elseUpdates,
       AccessPathNullnessPropagation.Updates bothUpdates) {
-    Symbol.MethodSymbol symbol = ASTHelpers.getSymbol(node.getTree());
-
     Types types = state.getTypes();
     if (optionalIsPresentCall(symbol, types)) {
       updateNonNullAPsForOptionalContent(

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

@@ -23,7 +23,6 @@
 package com.uber.nullaway.handlers;
 
 import com.google.errorprone.VisitorState;
-import com.google.errorprone.util.ASTHelpers;
 import com.sun.source.tree.ExpressionTree;
 import com.sun.source.tree.Tree;
 import com.sun.tools.javac.code.Symbol;
@@ -72,11 +71,16 @@ private boolean isSymbolRestrictivelyNullable(Symbol symbol, Context context) {
 
   @Override
   public boolean onOverrideMayBeNullExpr(
-      NullAway analysis, ExpressionTree expr, VisitorState state, boolean exprMayBeNull) {
+      NullAway analysis,
+      ExpressionTree expr,
+      @Nullable Symbol exprSymbol,
+      VisitorState state,
+      boolean exprMayBeNull) {
     Tree.Kind exprKind = expr.getKind();
-    if (exprKind.equals(Tree.Kind.METHOD_INVOCATION) || exprKind.equals(Tree.Kind.IDENTIFIER)) {
-      Symbol symbol = ASTHelpers.getSymbol(expr);
-      return exprMayBeNull || isSymbolRestrictivelyNullable(symbol, state.context);
+    if (exprSymbol != null
+        && (exprKind.equals(Tree.Kind.METHOD_INVOCATION)
+            || exprKind.equals(Tree.Kind.IDENTIFIER))) {
+      return exprMayBeNull || isSymbolRestrictivelyNullable(exprSymbol, state.context);
     }
     return exprMayBeNull;
   }
@@ -148,13 +152,13 @@ public Nullness onOverrideMethodInvocationReturnNullability(
   @Override
   public NullnessHint onDataflowVisitMethodInvocation(
       MethodInvocationNode node,
+      Symbol.MethodSymbol methodSymbol,
       VisitorState state,
       AccessPath.AccessPathContext apContext,
       AccessPathNullnessPropagation.SubNodeValues inputs,
       AccessPathNullnessPropagation.Updates thenUpdates,
       AccessPathNullnessPropagation.Updates elseUpdates,
       AccessPathNullnessPropagation.Updates bothUpdates) {
-    Symbol.MethodSymbol methodSymbol = ASTHelpers.getSymbol(node.getTree());
     return isSymbolRestrictivelyNullable(methodSymbol, state.context)
         ? NullnessHint.HINT_NULLABLE
         : NullnessHint.UNKNOWN;
@@ -169,7 +173,6 @@ public NullnessHint onDataflowVisitFieldAccess(
       AccessPath.AccessPathContext apContext,
       AccessPathNullnessPropagation.SubNodeValues inputs,
       AccessPathNullnessPropagation.Updates updates) {
-    ;
     return isSymbolRestrictivelyNullable(symbol, context)
         ? NullnessHint.HINT_NULLABLE
         : NullnessHint.UNKNOWN;

nullaway/src/main/java/com/uber/nullaway/handlers/contract/ContractHandler.java

@@ -176,15 +176,14 @@ public MethodInvocationNode onCFGBuildPhase1AfterVisitMethodInvocation(
   @Override
   public NullnessHint onDataflowVisitMethodInvocation(
       MethodInvocationNode node,
+      Symbol.MethodSymbol callee,
       VisitorState state,
       AccessPath.AccessPathContext apContext,
       AccessPathNullnessPropagation.SubNodeValues inputs,
       AccessPathNullnessPropagation.Updates thenUpdates,
       AccessPathNullnessPropagation.Updates elseUpdates,
       AccessPathNullnessPropagation.Updates bothUpdates) {
     Preconditions.checkNotNull(analysis);
-    Symbol.MethodSymbol callee = ASTHelpers.getSymbol(node.getTree());
-    Preconditions.checkNotNull(callee);
     MethodInvocationTree tree = castToNonNull(node.getTree());
     for (String clause : ContractUtils.getContractClauses(callee, config)) {