Fix error inside Lombok generated code for @nullable @Builder.Default (…

…uber#765)

When given code such as:

```
@builder
public class LombokDTO {
  @nullable @Builder.Default private String fieldWithNullDefault = null;
}
```

Lombok internally generates the following method:

```
@java.lang.SuppressWarnings(value = "all")
@lombok.Generated()
private static String $default$fieldWithNullDefault() {
    return null;
}
```

which does not propagate `@Nullable` to the method's return type!

While this method is marked as `@Generated` code and `@SuppressWarnings("all")`, that does not suppress NullAway under all configurations. In fact, we sometimes want to check generated code (setters/getters for auto-annotation, for example), so just counting all Lombok Generated code as unannotated is not always the desired behavior (it's optional behavior, enabled by the `TreatGeneratedAsUnannotated=true` flag).

Instead, we want to internally and implicitly propagate the `@Nullable` annotation from `fieldWithNullDefault` to the generated `$default$fieldWithNullDefault()`.

We do this in two steps:

1. We modify our checking of return statements to allow handlers' existing overriding of method return nullability to be taken into account when deciding if `return [nullable expression];` should result in a NullAway error.
2. We add a new handler for fixing the Lombok nullability propagation, by detecting these `$default$foo()` methods and looking at the nullability of `foo` to determine if the method should also be `@Nullable`.

In addition to this, we can suggest a fix upstream in Lombok to propagate `@Nullable` to these `$default$` methods when present on the field, but this would not obviate the need for this PR, since we are keeping compatibility with multiple older Lombok releases internally.

Edit: Also, note that coverage on `LombokHandler` is bad in terms of unit tests, but that's because coveralls doesn't count `test-java-lib-lombok/` as a test case. We need to build with Lombok to exercise most of this handler, so we can't really exercise it in unit tests (vs integration) without significantly artificial workarounds (i.e. manually replicating the Lombok-generated code).

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

@@ -808,6 +808,21 @@ static Trees getTreesInstance(VisitorState state) {
     return Trees.instance(JavacProcessingEnvironment.instance(state.context));
   }
 
+  private Nullness getMethodReturnNullness(
+      Symbol.MethodSymbol methodSymbol, VisitorState state, Nullness defaultForUnannotated) {
+    final boolean isMethodAnnotated = !codeAnnotationInfo.isSymbolUnannotated(methodSymbol, config);
+    Nullness methodReturnNullness =
+        defaultForUnannotated; // Permissive default for unannotated code.
+    if (isMethodAnnotated) {
+      methodReturnNullness =
+          Nullness.hasNullableAnnotation(methodSymbol, config)
+              ? Nullness.NULLABLE
+              : Nullness.NONNULL;
+    }
+    return handler.onOverrideMethodReturnNullability(
+        methodSymbol, state, isMethodAnnotated, methodReturnNullness);
+  }
+
   private Description checkReturnExpression(
       Tree tree, ExpressionTree retExpr, Symbol.MethodSymbol methodSymbol, VisitorState state) {
     Type returnType = methodSymbol.getReturnType();
@@ -822,8 +837,7 @@ private Description checkReturnExpression(
       // support)
       return Description.NO_MATCH;
     }
-    if (codeAnnotationInfo.isSymbolUnannotated(methodSymbol, config)
-        || Nullness.hasNullableAnnotation(methodSymbol, config)) {
+    if (getMethodReturnNullness(methodSymbol, state, Nullness.NULLABLE).equals(Nullness.NULLABLE)) {
       return Description.NO_MATCH;
     }
     if (mayBeNullExpr(state, retExpr)) {
@@ -908,63 +922,41 @@ private Description checkOverriding(
       Symbol.MethodSymbol overridingMethod,
       @Nullable MemberReferenceTree memberReferenceTree,
       VisitorState state) {
-    final boolean isOverriddenMethodAnnotated =
-        !codeAnnotationInfo.isSymbolUnannotated(overriddenMethod, config);
-    Nullness overriddenMethodReturnNullness =
-        Nullness.NULLABLE; // Permissive default for unannotated code.
-    if (isOverriddenMethodAnnotated && !Nullness.hasNullableAnnotation(overriddenMethod, config)) {
-      overriddenMethodReturnNullness = Nullness.NONNULL;
-    }
-    overriddenMethodReturnNullness =
-        handler.onOverrideMethodInvocationReturnNullability(
-            overriddenMethod, state, isOverriddenMethodAnnotated, overriddenMethodReturnNullness);
-    // if the super method returns nonnull,
-    // overriding method better not return nullable
-    if (overriddenMethodReturnNullness.equals(Nullness.NONNULL)) {
-      final boolean isOverridingMethodAnnotated =
-          !codeAnnotationInfo.isSymbolUnannotated(overridingMethod, config);
-      // Note that, for the overriding method, the permissive default is non-null.
-      Nullness overridingMethodReturnNullness = Nullness.NONNULL;
-      if (isOverridingMethodAnnotated && Nullness.hasNullableAnnotation(overridingMethod, config)) {
-        overridingMethodReturnNullness = Nullness.NULLABLE;
-      }
-      // We must once again check the handler chain, to allow it to update nullability of the
-      // overriding method
-      // (e.g. through AcknowledgeRestrictiveAnnotations=true)
-      overridingMethodReturnNullness =
-          handler.onOverrideMethodInvocationReturnNullability(
-              overridingMethod, state, isOverridingMethodAnnotated, overridingMethodReturnNullness);
-      if (overridingMethodReturnNullness.equals(Nullness.NULLABLE)
-          && (memberReferenceTree == null
-              || getComputedNullness(memberReferenceTree).equals(Nullness.NULLABLE))) {
-        String message;
-        if (memberReferenceTree != null) {
-          message =
-              "referenced method returns @Nullable, but functional interface method "
-                  + ASTHelpers.enclosingClass(overriddenMethod)
-                  + "."
-                  + overriddenMethod.toString()
-                  + " returns @NonNull";
-
-        } else {
-          message =
-              "method returns @Nullable, but superclass method "
-                  + ASTHelpers.enclosingClass(overriddenMethod)
-                  + "."
-                  + overriddenMethod.toString()
-                  + " returns @NonNull";
-        }
+    // if the super method returns nonnull, overriding method better not return nullable
+    // Note that, for the overriding method, the permissive default is non-null,
+    // but it's nullable for the overridden one.
+    if (getMethodReturnNullness(overriddenMethod, state, Nullness.NULLABLE).equals(Nullness.NONNULL)
+        && getMethodReturnNullness(overridingMethod, state, Nullness.NONNULL)
+            .equals(Nullness.NULLABLE)
+        && (memberReferenceTree == null
+            || getComputedNullness(memberReferenceTree).equals(Nullness.NULLABLE))) {
+      String message;
+      if (memberReferenceTree != null) {
+        message =
+            "referenced method returns @Nullable, but functional interface method "
+                + ASTHelpers.enclosingClass(overriddenMethod)
+                + "."
+                + overriddenMethod.toString()
+                + " returns @NonNull";
 
-        Tree errorTree =
-            memberReferenceTree != null
-                ? memberReferenceTree
-                : getTreesInstance(state).getTree(overridingMethod);
-        return errorBuilder.createErrorDescription(
-            new ErrorMessage(MessageTypes.WRONG_OVERRIDE_RETURN, message),
-            buildDescription(errorTree),
-            state,
-            overriddenMethod);
+      } else {
+        message =
+            "method returns @Nullable, but superclass method "
+                + ASTHelpers.enclosingClass(overriddenMethod)
+                + "."
+                + overriddenMethod.toString()
+                + " returns @NonNull";
       }
+
+      Tree errorTree =
+          memberReferenceTree != null
+              ? memberReferenceTree
+              : getTreesInstance(state).getTree(overridingMethod);
+      return errorBuilder.createErrorDescription(
+          new ErrorMessage(MessageTypes.WRONG_OVERRIDE_RETURN, message),
+          buildDescription(errorTree),
+          state,
+          overriddenMethod);
     }
     // if any parameter in the super method is annotated @Nullable,
     // overriding method cannot assume @Nonnull

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

@@ -109,7 +109,7 @@ public void onMatchReturn(NullAway analysis, ReturnTree tree, VisitorState state
   }
 
   @Override
-  public Nullness onOverrideMethodInvocationReturnNullability(
+  public Nullness onOverrideMethodReturnNullability(
       Symbol.MethodSymbol methodSymbol,
       VisitorState state,
       boolean isAnnotated,

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

@@ -123,15 +123,14 @@ public void onMatchReturn(NullAway analysis, ReturnTree tree, VisitorState state
   }
 
   @Override
-  public Nullness onOverrideMethodInvocationReturnNullability(
+  public Nullness onOverrideMethodReturnNullability(
       Symbol.MethodSymbol methodSymbol,
       VisitorState state,
       boolean isAnnotated,
       Nullness returnNullness) {
     for (Handler h : handlers) {
       returnNullness =
-          h.onOverrideMethodInvocationReturnNullability(
-              methodSymbol, state, isAnnotated, returnNullness);
+          h.onOverrideMethodReturnNullability(methodSymbol, state, isAnnotated, returnNullness);
     }
     return returnNullness;
   }

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

@@ -162,12 +162,12 @@ boolean onOverrideMayBeNullExpr(
    * @param methodSymbol The method symbol for the method in question.
    * @param state The current visitor state.
    * @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
-   *     multiple times within the same handler chain.
+   *     within annotated or unannotated code, used to avoid querying for this information multiple
+   *     times within the same handler chain.
    * @param returnNullness return nullness computed by upstream handlers or NullAway core.
    * @return Updated return nullability computed by this handler.
    */
-  Nullness onOverrideMethodInvocationReturnNullability(
+  Nullness onOverrideMethodReturnNullability(
       Symbol.MethodSymbol methodSymbol,
       VisitorState state,
       boolean isAnnotated,

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

@@ -74,6 +74,7 @@ public static Handler buildDefault(Config config) {
     if (config.checkContracts()) {
       handlerListBuilder.add(new ContractCheckHandler(config));
     }
+    handlerListBuilder.add(new LombokHandler(config));
 
     return new CompositeHandler(handlerListBuilder.build());
   }

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

@@ -109,14 +109,16 @@ public Nullness[] onOverrideMethodInvocationParametersNullability(
   }
 
   @Override
-  public Nullness onOverrideMethodInvocationReturnNullability(
+  public Nullness onOverrideMethodReturnNullability(
       Symbol.MethodSymbol methodSymbol,
       VisitorState state,
       boolean isAnnotated,
       Nullness returnNullness) {
     OptimizedLibraryModels optLibraryModels = getOptLibraryModels(state.context);
     if (optLibraryModels.hasNonNullReturn(methodSymbol, state.getTypes(), !isAnnotated)) {
-      return Nullness.NONNULL;
+      return NONNULL;
+    } else if (optLibraryModels.hasNullableReturn(methodSymbol, state.getTypes(), !isAnnotated)) {
+      return NULLABLE;
     }
     return returnNullness;
   }

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

@@ -0,0 +1,95 @@
+package com.uber.nullaway.handlers;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
+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;
+import com.uber.nullaway.Config;
+import com.uber.nullaway.NullAway;
+import com.uber.nullaway.Nullness;
+import java.util.stream.StreamSupport;
+import javax.annotation.Nullable;
+import javax.lang.model.element.ElementKind;
+
+/**
+ * A general handler for Lombok generated code and its internal semantics.
+ *
+ * <p>Currently used to propagate @Nullable in cases where the Lombok annotation processor fails to
+ * do so consistently.
+ */
+public class LombokHandler extends BaseNoOpHandler {
+
+  private static String LOMBOK_GENERATED_ANNOTATION_NAME = "lombok.Generated";
+  private static String LOMBOK_BUILDER_DEFAULT_METHOD_PREFIX = "$default$";
+
+  private final Config config;
+
+  public LombokHandler(Config config) {
+    this.config = config;
+  }
+
+  @SuppressWarnings("ASTHelpersSuggestions") // Suggested API doesn't exist in EP 2.4.0
+  private boolean isLombokMethodWithMissingNullableAnnotation(
+      Symbol.MethodSymbol methodSymbol, VisitorState state) {
+    if (!ASTHelpers.hasAnnotation(methodSymbol, LOMBOK_GENERATED_ANNOTATION_NAME, state)) {
+      return false;
+    }
+    String methodNameString = methodSymbol.name.toString();
+    if (!methodNameString.startsWith(LOMBOK_BUILDER_DEFAULT_METHOD_PREFIX)) {
+      return false;
+    }
+    String originalFieldName =
+        methodNameString.substring(LOMBOK_BUILDER_DEFAULT_METHOD_PREFIX.length());
+    ImmutableList<Symbol> matchingMembers =
+        StreamSupport.stream(
+                methodSymbol
+                    .enclClass()
+                    .members()
+                    .getSymbols(
+                        sym ->
+                            sym.name.contentEquals(originalFieldName)
+                                && sym.getKind().equals(ElementKind.FIELD))
+                    .spliterator(),
+                false)
+            .collect(ImmutableList.toImmutableList());
+    Preconditions.checkArgument(
+        matchingMembers.size() == 1,
+        String.format(
+            "Found %d fields matching Lombok generated builder default method %s",
+            matchingMembers.size(), methodNameString));
+    return Nullness.hasNullableAnnotation(matchingMembers.get(0), config);
+  }
+
+  @Override
+  public boolean onOverrideMayBeNullExpr(
+      NullAway analysis,
+      ExpressionTree expr,
+      @Nullable Symbol exprSymbol,
+      VisitorState state,
+      boolean exprMayBeNull) {
+    if (exprMayBeNull) {
+      return true;
+    }
+    Tree.Kind exprKind = expr.getKind();
+    if (exprSymbol != null && exprKind == Tree.Kind.METHOD_INVOCATION) {
+      Symbol.MethodSymbol methodSymbol = (Symbol.MethodSymbol) exprSymbol;
+      return isLombokMethodWithMissingNullableAnnotation(methodSymbol, state);
+    }
+    return false;
+  }
+
+  @Override
+  public Nullness onOverrideMethodReturnNullability(
+      Symbol.MethodSymbol methodSymbol,
+      VisitorState state,
+      boolean isAnnotated,
+      Nullness returnNullness) {
+    if (isLombokMethodWithMissingNullableAnnotation(methodSymbol, state)) {
+      return Nullness.NULLABLE;
+    }
+    return returnNullness;
+  }
+}

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

@@ -133,7 +133,7 @@ public Nullness[] onOverrideMethodInvocationParametersNullability(
   }
 
   @Override
-  public Nullness onOverrideMethodInvocationReturnNullability(
+  public Nullness onOverrideMethodReturnNullability(
       Symbol.MethodSymbol methodSymbol,
       VisitorState state,
       boolean isAnnotated,

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

@@ -53,20 +53,22 @@ public void allowLibraryModelsOverrideAnnotations() {
                 temporaryFolder.getRoot().getAbsolutePath(),
                 "-XepOpt:NullAway:AnnotatedPackages=com.uber"))
         .addSourceLines(
-            "Foo.java",
+            "AnnotatedWithModels.java",
             "package com.uber;",
-            "public class Foo {",
+            "public class AnnotatedWithModels {",
             "   Object field = new Object();",
-            "   Object bar() {",
-            "      return new Object();",
+            "   // implicitly @Nullable due to library model",
+            "   Object returnsNullFromModel() {",
+            "      // null is valid here only because of the library model",
+            "      return null;",
             "   }",
             "   Object nullableReturn() {",
             "       // BUG: Diagnostic contains: returning @Nullable",
-            "       return bar();",
+            "       return returnsNullFromModel();",
             "   }",
             "   void run() {",
             "       // just to make sure, flow analysis is also impacted by library models information",
-            "      Object temp = bar();",
+            "      Object temp = returnsNullFromModel();",
             "       // BUG: Diagnostic contains: assigning @Nullable",
             "      this.field = temp;",
             "   }",

test-java-lib-lombok/src/main/java/com/uber/lombok/LombokDTO.java

@@ -36,4 +36,5 @@ public class LombokDTO {
   private String field;
   @Builder.Default private String fieldWithDefault = "Default";
   @Nullable private String nullableField;
+  @Nullable @Builder.Default private String fieldWithNullDefault = null;
 }

test-library-models/src/main/java/com/uber/nullaway/testlibrarymodels/TestLibraryModels.java

@@ -66,7 +66,7 @@ public ImmutableSetMultimap<MethodRef, Integer> nullImpliesNullParameters() {
   @Override
   public ImmutableSet<MethodRef> nullableReturns() {
     return ImmutableSet.of(
-        methodRef("com.uber.Foo", "bar()"),
+        methodRef("com.uber.AnnotatedWithModels", "returnsNullFromModel()"),
         methodRef("com.uber.lib.unannotated.UnannotatedWithModels", "returnsNullUnannotated()"),
         methodRef("com.uber.lib.unannotated.UnannotatedWithModels", "returnsNullUnannotated2()"));
   }