Model Lombok-generated equals methods as having a @Nullable parameter (#874)

See the new test in `UsesDTO`. Without this change, NullAway reports an
incorrect warning for this code that passes a `@Nullable` parameter to a
Lombok-generated `equals()` method.

The real logic change in this PR is very small. Most of the changes are
required to make a `VisitorState` object available in a `Handler`
method.

13 files changed, 53 insertions, 34 deletions

diff --git a/nullaway/src/main/java/com/uber/nullaway/NullAway.java b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
index 744bddd..f7ee780 100644
--- a/nullaway/src/main/java/com/uber/nullaway/NullAway.java
+++ b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
@@ -760,10 +760,7 @@ public class NullAway extends BugChecker
     // Check handlers for any further/overriding nullness information
     overriddenMethodArgNullnessMap =
         handler.onOverrideMethodInvocationParametersNullability(
-            state.context,
-            overriddenMethod,
-            isOverriddenMethodAnnotated,
-            overriddenMethodArgNullnessMap);
+            state, overriddenMethod, isOverriddenMethodAnnotated, overriddenMethodArgNullnessMap);
 
     // If we have an unbound method reference, the first parameter of the overridden method must be
     // @NonNull, as this parameter will be used as a method receiver inside the generated lambda.
@@ -1714,7 +1711,7 @@ public class NullAway extends BugChecker
     // Allow handlers to override the list of non-null argument positions
     argumentPositionNullness =
         handler.onOverrideMethodInvocationParametersNullability(
-            state.context, methodSymbol, isMethodAnnotated, argumentPositionNullness);
+            state, methodSymbol, isMethodAnnotated, argumentPositionNullness);
 
     // now actually check the arguments
     // NOTE: the case of an invocation on a possibly-null reference
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 b68665a..fb18368 100644
--- a/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessPropagation.java
+++ b/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessPropagation.java
@@ -206,7 +206,7 @@ public class AccessPathNullnessPropagation
   public NullnessStore initialStore(
       UnderlyingAST underlyingAST, List<LocalVariableNode> parameters) {
     return nullnessStoreInitializer.getInitialStore(
-        underlyingAST, parameters, handler, state.context, state.getTypes(), config);
+        underlyingAST, parameters, handler, state, config);
   }
 
   @Override
diff --git a/nullaway/src/main/java/com/uber/nullaway/dataflow/CoreNullnessStoreInitializer.java b/nullaway/src/main/java/com/uber/nullaway/dataflow/CoreNullnessStoreInitializer.java
index 2dfce9f..170ccb1 100644
--- a/nullaway/src/main/java/com/uber/nullaway/dataflow/CoreNullnessStoreInitializer.java
+++ b/nullaway/src/main/java/com/uber/nullaway/dataflow/CoreNullnessStoreInitializer.java
@@ -3,6 +3,7 @@ package com.uber.nullaway.dataflow;
 import static com.uber.nullaway.Nullness.NONNULL;
 import static com.uber.nullaway.Nullness.NULLABLE;
 
+import com.google.errorprone.VisitorState;
 import com.google.errorprone.util.ASTHelpers;
 import com.sun.source.tree.ClassTree;
 import com.sun.source.tree.LambdaExpressionTree;
@@ -30,9 +31,9 @@ class CoreNullnessStoreInitializer extends NullnessStoreInitializer {
       UnderlyingAST underlyingAST,
       List<LocalVariableNode> parameters,
       Handler handler,
-      Context context,
-      Types types,
+      VisitorState state,
       Config config) {
+    Context context = state.context;
     if (underlyingAST.getKind().equals(UnderlyingAST.Kind.ARBITRARY_CODE)) {
       // not a method or a lambda; an initializer expression or block
       UnderlyingAST.CFGStatement ast = (UnderlyingAST.CFGStatement) underlyingAST;
@@ -44,8 +45,7 @@ class CoreNullnessStoreInitializer extends NullnessStoreInitializer {
           (UnderlyingAST.CFGLambda) underlyingAST,
           parameters,
           handler,
-          context,
-          types,
+          state,
           config,
           getCodeAnnotationInfo(context));
     } else {
@@ -77,13 +77,12 @@ class CoreNullnessStoreInitializer extends NullnessStoreInitializer {
       UnderlyingAST.CFGLambda underlyingAST,
       List<LocalVariableNode> parameters,
       Handler handler,
-      Context context,
-      Types types,
+      VisitorState state,
       Config config,
       CodeAnnotationInfo codeAnnotationInfo) {
     // include nullness info for locals from enclosing environment
     EnclosingEnvironmentNullness environmentNullness =
-        EnclosingEnvironmentNullness.instance(context);
+        EnclosingEnvironmentNullness.instance(state.context);
     NullnessStore environmentMapping =
         Objects.requireNonNull(
             environmentNullness.getEnvironmentMapping(underlyingAST.getLambdaTree()),
@@ -91,6 +90,7 @@ class CoreNullnessStoreInitializer extends NullnessStoreInitializer {
     NullnessStore.Builder result = environmentMapping.toBuilder();
     LambdaExpressionTree code = underlyingAST.getLambdaTree();
     // need to check annotation for i'th parameter of functional interface declaration
+    Types types = state.getTypes();
     Symbol.MethodSymbol fiMethodSymbol = NullabilityUtil.getFunctionalInterfaceMethod(code, types);
     com.sun.tools.javac.util.List<Symbol.VarSymbol> fiMethodParameters =
         fiMethodSymbol.getParameters();
@@ -119,7 +119,7 @@ class CoreNullnessStoreInitializer extends NullnessStoreInitializer {
     }
     fiArgumentPositionNullness =
         handler.onOverrideMethodInvocationParametersNullability(
-            context, fiMethodSymbol, isFIAnnotated, fiArgumentPositionNullness);
+            state, fiMethodSymbol, isFIAnnotated, fiArgumentPositionNullness);
 
     for (int i = 0; i < parameters.size(); i++) {
       LocalVariableNode param = parameters.get(i);
diff --git a/nullaway/src/main/java/com/uber/nullaway/dataflow/NullnessStoreInitializer.java b/nullaway/src/main/java/com/uber/nullaway/dataflow/NullnessStoreInitializer.java
index c654621..37c68d3 100644
--- a/nullaway/src/main/java/com/uber/nullaway/dataflow/NullnessStoreInitializer.java
+++ b/nullaway/src/main/java/com/uber/nullaway/dataflow/NullnessStoreInitializer.java
@@ -1,10 +1,10 @@
 package com.uber.nullaway.dataflow;
 
+import com.google.errorprone.VisitorState;
 import com.google.errorprone.util.ASTHelpers;
 import com.sun.source.tree.ClassTree;
 import com.sun.source.util.Trees;
 import com.sun.tools.javac.code.Symbol;
-import com.sun.tools.javac.code.Types;
 import com.sun.tools.javac.processing.JavacProcessingEnvironment;
 import com.sun.tools.javac.util.Context;
 import com.uber.nullaway.Config;
@@ -30,8 +30,7 @@ public abstract class NullnessStoreInitializer {
    * @param underlyingAST The AST node being matched.
    * @param parameters list of local variable nodes.
    * @param handler reference to the handler invoked.
-   * @param context context.
-   * @param types types.
+   * @param state the visitor state.
    * @param config config for analysis.
    * @return Initial Nullness store.
    */
@@ -39,8 +38,7 @@ public abstract class NullnessStoreInitializer {
       UnderlyingAST underlyingAST,
       List<LocalVariableNode> parameters,
       Handler handler,
-      Context context,
-      Types types,
+      VisitorState state,
       Config config);
 
   /**
diff --git a/nullaway/src/main/java/com/uber/nullaway/handlers/BaseNoOpHandler.java b/nullaway/src/main/java/com/uber/nullaway/handlers/BaseNoOpHandler.java
index 242a96a..4fc0c0c 100644
--- a/nullaway/src/main/java/com/uber/nullaway/handlers/BaseNoOpHandler.java
+++ b/nullaway/src/main/java/com/uber/nullaway/handlers/BaseNoOpHandler.java
@@ -120,7 +120,7 @@ public abstract class BaseNoOpHandler implements Handler {
 
   @Override
   public Nullness[] onOverrideMethodInvocationParametersNullability(
-      Context context,
+      VisitorState state,
       Symbol.MethodSymbol methodSymbol,
       boolean isAnnotated,
       Nullness[] argumentPositionNullness) {
diff --git a/nullaway/src/main/java/com/uber/nullaway/handlers/CompositeHandler.java b/nullaway/src/main/java/com/uber/nullaway/handlers/CompositeHandler.java
index 31617da..3602a5f 100644
--- a/nullaway/src/main/java/com/uber/nullaway/handlers/CompositeHandler.java
+++ b/nullaway/src/main/java/com/uber/nullaway/handlers/CompositeHandler.java
@@ -137,14 +137,14 @@ class CompositeHandler implements Handler {
 
   @Override
   public Nullness[] onOverrideMethodInvocationParametersNullability(
-      Context context,
+      VisitorState state,
       Symbol.MethodSymbol methodSymbol,
       boolean isAnnotated,
       Nullness[] argumentPositionNullness) {
     for (Handler h : handlers) {
       argumentPositionNullness =
           h.onOverrideMethodInvocationParametersNullability(
-              context, methodSymbol, isAnnotated, argumentPositionNullness);
+              state, methodSymbol, isAnnotated, argumentPositionNullness);
     }
     return argumentPositionNullness;
   }
diff --git a/nullaway/src/main/java/com/uber/nullaway/handlers/Handler.java b/nullaway/src/main/java/com/uber/nullaway/handlers/Handler.java
index 835c01f..d509d7b 100644
--- a/nullaway/src/main/java/com/uber/nullaway/handlers/Handler.java
+++ b/nullaway/src/main/java/com/uber/nullaway/handlers/Handler.java
@@ -182,7 +182,7 @@ public interface Handler {
    * considered isAnnotated or not. We use a mutable map for performance, but it should not outlive
    * the chain of handler invocations.
    *
-   * @param context The current context.
+   * @param state The current visitor state.
    * @param methodSymbol The method symbol for the method in question.
    * @param isAnnotated A boolean flag indicating whether the called method is considered to be
    *     within isAnnotated or unannotated code, used to avoid querying for this information
@@ -195,7 +195,7 @@ public interface Handler {
    *     handler.
    */
   Nullness[] onOverrideMethodInvocationParametersNullability(
-      Context context,
+      VisitorState state,
       Symbol.MethodSymbol methodSymbol,
       boolean isAnnotated,
       Nullness[] argumentPositionNullness);
diff --git a/nullaway/src/main/java/com/uber/nullaway/handlers/InferredJARModelsHandler.java b/nullaway/src/main/java/com/uber/nullaway/handlers/InferredJARModelsHandler.java
index bd70059..71dbd92 100644
--- a/nullaway/src/main/java/com/uber/nullaway/handlers/InferredJARModelsHandler.java
+++ b/nullaway/src/main/java/com/uber/nullaway/handlers/InferredJARModelsHandler.java
@@ -28,7 +28,6 @@ import com.sun.source.tree.ExpressionTree;
 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.util.Context;
 import com.uber.nullaway.Config;
 import com.uber.nullaway.NullAway;
 import com.uber.nullaway.Nullness;
@@ -122,7 +121,7 @@ public class InferredJARModelsHandler extends BaseNoOpHandler {
 
   @Override
   public Nullness[] onOverrideMethodInvocationParametersNullability(
-      Context context,
+      VisitorState state,
       Symbol.MethodSymbol methodSymbol,
       boolean isAnnotated,
       Nullness[] argumentPositionNullness) {
diff --git a/nullaway/src/main/java/com/uber/nullaway/handlers/LibraryModelsHandler.java b/nullaway/src/main/java/com/uber/nullaway/handlers/LibraryModelsHandler.java
index e000c4f..e5e943c 100644
--- a/nullaway/src/main/java/com/uber/nullaway/handlers/LibraryModelsHandler.java
+++ b/nullaway/src/main/java/com/uber/nullaway/handlers/LibraryModelsHandler.java
@@ -82,11 +82,11 @@ public class LibraryModelsHandler extends BaseNoOpHandler {
 
   @Override
   public Nullness[] onOverrideMethodInvocationParametersNullability(
-      Context context,
+      VisitorState state,
       Symbol.MethodSymbol methodSymbol,
       boolean isAnnotated,
       Nullness[] argumentPositionNullness) {
-    OptimizedLibraryModels optimizedLibraryModels = getOptLibraryModels(context);
+    OptimizedLibraryModels optimizedLibraryModels = getOptLibraryModels(state.context);
     ImmutableSet<Integer> nullableParamsFromModel =
         optimizedLibraryModels.explicitlyNullableParameters(methodSymbol);
     ImmutableSet<Integer> nonNullParamsFromModel =
diff --git a/nullaway/src/main/java/com/uber/nullaway/handlers/LombokHandler.java b/nullaway/src/main/java/com/uber/nullaway/handlers/LombokHandler.java
index 7069497..7d76611 100644
--- a/nullaway/src/main/java/com/uber/nullaway/handlers/LombokHandler.java
+++ b/nullaway/src/main/java/com/uber/nullaway/handlers/LombokHandler.java
@@ -86,4 +86,26 @@ public class LombokHandler extends BaseNoOpHandler {
     }
     return returnNullness;
   }
+
+  /**
+   * Mark the first argument of Lombok-generated {@code equals} methods as {@code @Nullable}, since
+   * Lombok does not generate the annotation.
+   */
+  @Override
+  public Nullness[] onOverrideMethodInvocationParametersNullability(
+      VisitorState state,
+      Symbol.MethodSymbol methodSymbol,
+      boolean isAnnotated,
+      Nullness[] argumentPositionNullness) {
+    if (ASTHelpers.hasAnnotation(methodSymbol, LOMBOK_GENERATED_ANNOTATION_NAME, state)) {
+      // We assume that Lombok-generated equals methods with a single argument override
+      // Object.equals and are not an overload.
+      if (methodSymbol.getSimpleName().contentEquals("equals")
+          && methodSymbol.params().size() == 1) {
+        // The parameter is not annotated with @Nullable, but it should be.
+        argumentPositionNullness[0] = Nullness.NULLABLE;
+      }
+    }
+    return argumentPositionNullness;
+  }
 }
diff --git a/nullaway/src/main/java/com/uber/nullaway/handlers/RestrictiveAnnotationHandler.java b/nullaway/src/main/java/com/uber/nullaway/handlers/RestrictiveAnnotationHandler.java
index 561ac6e..7c5c9ba 100644
--- a/nullaway/src/main/java/com/uber/nullaway/handlers/RestrictiveAnnotationHandler.java
+++ b/nullaway/src/main/java/com/uber/nullaway/handlers/RestrictiveAnnotationHandler.java
@@ -99,7 +99,7 @@ public class RestrictiveAnnotationHandler extends BaseNoOpHandler {
 
   @Override
   public Nullness[] onOverrideMethodInvocationParametersNullability(
-      Context context,
+      VisitorState state,
       Symbol.MethodSymbol methodSymbol,
       boolean isAnnotated,
       Nullness[] argumentPositionNullness) {
diff --git a/nullaway/src/main/java/com/uber/nullaway/handlers/contract/ContractNullnessStoreInitializer.java b/nullaway/src/main/java/com/uber/nullaway/handlers/contract/ContractNullnessStoreInitializer.java
index 41b25d5..155bd90 100644
--- a/nullaway/src/main/java/com/uber/nullaway/handlers/contract/ContractNullnessStoreInitializer.java
+++ b/nullaway/src/main/java/com/uber/nullaway/handlers/contract/ContractNullnessStoreInitializer.java
@@ -3,12 +3,11 @@ package com.uber.nullaway.handlers.contract;
 import static com.uber.nullaway.Nullness.NONNULL;
 import static com.uber.nullaway.Nullness.NULLABLE;
 
+import com.google.errorprone.VisitorState;
 import com.google.errorprone.util.ASTHelpers;
 import com.sun.source.tree.ClassTree;
 import com.sun.source.tree.MethodTree;
 import com.sun.tools.javac.code.Symbol;
-import com.sun.tools.javac.code.Types;
-import com.sun.tools.javac.util.Context;
 import com.uber.nullaway.Config;
 import com.uber.nullaway.Nullness;
 import com.uber.nullaway.dataflow.AccessPath;
@@ -31,8 +30,7 @@ public class ContractNullnessStoreInitializer extends NullnessStoreInitializer {
       UnderlyingAST underlyingAST,
       List<LocalVariableNode> parameters,
       Handler handler,
-      Context context,
-      Types types,
+      VisitorState state,
       Config config) {
     assert underlyingAST.getKind() == UnderlyingAST.Kind.METHOD;
 
@@ -49,7 +47,7 @@ public class ContractNullnessStoreInitializer extends NullnessStoreInitializer {
     String[] parts = clauses[0].split("->");
     String[] antecedent = parts[0].split(",");
 
-    NullnessStore envStore = getEnvNullnessStoreForClass(classTree, context);
+    NullnessStore envStore = getEnvNullnessStoreForClass(classTree, state.context);
     NullnessStore.Builder result = envStore.toBuilder();
 
     for (int i = 0; i < antecedent.length; ++i) {
diff --git a/test-java-lib-lombok/src/main/java/com/uber/lombok/UsesDTO.java b/test-java-lib-lombok/src/main/java/com/uber/lombok/UsesDTO.java
index f236a65..bef6f78 100644
--- a/test-java-lib-lombok/src/main/java/com/uber/lombok/UsesDTO.java
+++ b/test-java-lib-lombok/src/main/java/com/uber/lombok/UsesDTO.java
@@ -38,4 +38,9 @@ class UsesDTO {
     s += (ldto.getNullableField() == null ? "" : ldto.getNullableField().toString());
     return s;
   }
+
+  public static boolean callEquals(LombokDTO ldto, @Nullable Object o) {
+    // No error should be reported since equals() parameter should be treated as @Nullable
+    return ldto.equals(o);
+  }
 }