Handle assertThat(...).isNotNull() statements

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

@@ -70,6 +70,8 @@ public abstract class AbstractConfig implements Config {
 
   protected boolean checkOptionalEmptiness;
 
+  protected boolean handleTestAssertionLibraries;
+
   protected Set<String> optionalClassPaths;
 
   protected boolean assertsEnabled;
@@ -191,6 +193,11 @@ public boolean checkOptionalEmptiness() {
     return checkOptionalEmptiness;
   }
 
+  @Override
+  public boolean handleTestAssertionLibraries() {
+    return handleTestAssertionLibraries;
+  }
+
   @Override
   public Set<String> getOptionalClassPaths() {
     return optionalClassPaths;

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

@@ -116,6 +116,15 @@ public interface Config {
    */
   boolean checkOptionalEmptiness();
 
+  /**
+   * @return true if AssertionHandler should be enabled. In the absence of this handler, checks in
+   *     tests using assertion libraries are ignored. So, any deference of an object that follows
+   *     such checks are still considered a potential dereference of null. When this handler is
+   *     enabled, it understands checks using assertion libraries and reasons about the following
+   *     code such that any deference of objects that have been checked will never be null.
+   */
+  boolean handleTestAssertionLibraries();
+
   /**
    * @return the paths for Optional class. The list always contains the path of {@link
    *     java.util.Optional}.

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

@@ -118,6 +118,11 @@ public boolean checkOptionalEmptiness() {
     throw new IllegalStateException(error_msg);
   }
 
+  @Override
+  public boolean handleTestAssertionLibraries() {
+    throw new IllegalStateException(error_msg);
+  }
+
   @Override
   public Set<String> getOptionalClassPaths() {
     throw new IllegalStateException(error_msg);

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

@@ -56,6 +56,8 @@ final class ErrorProneCLIFlagsConfig extends AbstractConfig {
   static final String FL_ACKNOWLEDGE_RESTRICTIVE =
       EP_FL_NAMESPACE + ":AcknowledgeRestrictiveAnnotations";
   static final String FL_CHECK_OPTIONAL_EMPTINESS = EP_FL_NAMESPACE + ":CheckOptionalEmptiness";
+  static final String FL_HANDLE_TEST_ASSERTION_LIBRARIES =
+      EP_FL_NAMESPACE + ":HandleTestAssertionLibraries";
   static final String FL_OPTIONAL_CLASS_PATHS =
       EP_FL_NAMESPACE + ":CheckOptionalEmptinessCustomClasses";
   static final String FL_SUPPRESS_COMMENT = EP_FL_NAMESPACE + ":AutoFixSuppressionComment";
@@ -136,6 +138,8 @@ final class ErrorProneCLIFlagsConfig extends AbstractConfig {
     isSuggestSuppressions = flags.getBoolean(FL_SUGGEST_SUPPRESSIONS).orElse(false);
     isAcknowledgeRestrictive = flags.getBoolean(FL_ACKNOWLEDGE_RESTRICTIVE).orElse(false);
     checkOptionalEmptiness = flags.getBoolean(FL_CHECK_OPTIONAL_EMPTINESS).orElse(false);
+    handleTestAssertionLibraries =
+        flags.getBoolean(FL_HANDLE_TEST_ASSERTION_LIBRARIES).orElse(false);
     treatGeneratedAsUnannotated = flags.getBoolean(FL_GENERATED_UNANNOTATED).orElse(false);
     assertsEnabled = flags.getBoolean(FL_ASSERTS_ENABLED).orElse(false);
     fieldAnnotPattern =

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

@@ -0,0 +1,116 @@
+/*
+ * Copyright (c) 2017 Uber Technologies, Inc.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+package com.uber.nullaway.handlers;
+
+import static com.uber.nullaway.Nullness.NONNULL;
+
+import com.google.errorprone.util.ASTHelpers;
+import com.sun.tools.javac.code.Symbol;
+import com.sun.tools.javac.code.Types;
+import com.sun.tools.javac.util.Context;
+import com.sun.tools.javac.util.Name;
+import com.uber.nullaway.dataflow.AccessPath;
+import com.uber.nullaway.dataflow.AccessPathNullnessPropagation;
+import org.checkerframework.dataflow.cfg.node.MethodInvocationNode;
+import org.checkerframework.dataflow.cfg.node.Node;
+
+/** This Handler deals with assertions which ensure that their arguments cannot be null. */
+public class AssertionHandler extends BaseNoOpHandler {
+
+  // Strings corresponding to the names of the methods (and their owners) used to identify
+  // assertions in this handler.
+  private static final String IS_NOT_NULL_METHOD = "isNotNull";
+  private static final String IS_NOT_NULL_OWNER = "com.google.common.truth.Subject";
+  private static final String ASSERT_THAT_METHOD = "assertThat";
+  private static final String ASSERT_THAT_OWNER = "com.google.common.truth.Truth";
+
+  // Names of the methods (and their owners) used to identify assertions in this handler. Name used
+  // here refers to com.sun.tools.javac.util.Name. Comparing methods using Names is faster than
+  // comparing using strings.
+  private Name isNotNull;
+  private Name isNotNullOwner;
+  private Name assertThat;
+  private Name assertThatOwner;
+
+  @Override
+  public NullnessHint onDataflowVisitMethodInvocation(
+      MethodInvocationNode node,
+      Types types,
+      Context context,
+      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 (!areMethodNamesInitialized()) {
+      initializeMethodNames(callee.name.table);
+    }
+
+    // Look for statements of the form: assertThat(A).isNotNull()
+    // A will not be NULL after this statement.
+    if (isMethodIsNotNull(callee)) {
+      Node receiver = node.getTarget().getReceiver();
+      if (receiver instanceof MethodInvocationNode) {
+        MethodInvocationNode receiver_method = (MethodInvocationNode) receiver;
+        Symbol.MethodSymbol receiver_symbol = ASTHelpers.getSymbol(receiver_method.getTree());
+        if (isMethodAssertThat(receiver_symbol)) {
+          Node arg = receiver_method.getArgument(0);
+          AccessPath ap = AccessPath.getAccessPathForNodeNoMapGet(arg);
+          if (ap != null) {
+            bothUpdates.set(ap, NONNULL);
+          }
+        }
+      }
+    }
+    return NullnessHint.UNKNOWN;
+  }
+
+  private boolean isMethodIsNotNull(Symbol.MethodSymbol methodSymbol) {
+    return matchesMethod(methodSymbol, isNotNull, isNotNullOwner);
+  }
+
+  private boolean isMethodAssertThat(Symbol.MethodSymbol methodSymbol) {
+    return matchesMethod(methodSymbol, assertThat, assertThatOwner);
+  }
+
+  private boolean matchesMethod(
+      Symbol.MethodSymbol methodSymbol, Name toMatchMethodName, Name toMatchOwnerName) {
+    return methodSymbol.name.equals(toMatchMethodName)
+        && methodSymbol.owner.getQualifiedName().equals(toMatchOwnerName);
+  }
+
+  private boolean areMethodNamesInitialized() {
+    return isNotNull != null;
+  }
+
+  private void initializeMethodNames(Name.Table table) {
+    isNotNull = table.fromString(IS_NOT_NULL_METHOD);
+    isNotNullOwner = table.fromString(IS_NOT_NULL_OWNER);
+    assertThat = table.fromString(ASSERT_THAT_METHOD);
+    assertThatOwner = table.fromString(ASSERT_THAT_OWNER);
+  }
+}

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

@@ -46,6 +46,9 @@ public static Handler buildDefault(Config config) {
     if (config.isJarInferEnabled()) {
       handlerListBuilder.add(new InferredJARModelsHandler(config));
     }
+    if (config.handleTestAssertionLibraries()) {
+      handlerListBuilder.add(new AssertionHandler());
+    }
     handlerListBuilder.add(new LibraryModelsHandler());
     handlerListBuilder.add(new RxNullabilityPropagator());
     handlerListBuilder.add(new ContractHandler());

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

@@ -967,6 +967,78 @@ public void supportObjectsIsNull() {
         .doTest();
   }
 
+  @Test
+  public void supportAssertThatIsNotNull_Object() {
+    compilationHelper
+        .setArgs(
+            Arrays.asList(
+                "-d",
+                temporaryFolder.getRoot().getAbsolutePath(),
+                "-XepOpt:NullAway:AnnotatedPackages=com.uber",
+                "-XepOpt:NullAway:HandleTestAssertionLibraries=true"))
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import java.lang.Object;",
+            "import java.util.Objects;",
+            "import javax.annotation.Nullable;",
+            "class Test {",
+            "  private void foo(@Nullable Object o) {",
+            "    com.google.common.truth.Truth.assertThat(o).isNotNull();",
+            "    o.toString();",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void supportAssertThatIsNotNull_String() {
+    compilationHelper
+        .setArgs(
+            Arrays.asList(
+                "-d",
+                temporaryFolder.getRoot().getAbsolutePath(),
+                "-XepOpt:NullAway:AnnotatedPackages=com.uber",
+                "-XepOpt:NullAway:HandleTestAssertionLibraries=true"))
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import java.util.Objects;",
+            "import javax.annotation.Nullable;",
+            "class Test {",
+            "  private void foo(@Nullable String s) {",
+            "    com.google.common.truth.Truth.assertThat(s).isNotNull();",
+            "    s.toString();",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void doNotSupportAssertThatWhenDisabled() {
+    compilationHelper
+        .setArgs(
+            Arrays.asList(
+                "-d",
+                temporaryFolder.getRoot().getAbsolutePath(),
+                "-XepOpt:NullAway:AnnotatedPackages=com.uber",
+                "-XepOpt:NullAway:HandleTestAssertionLibraries=false"))
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import java.lang.Object;",
+            "import java.util.Objects;",
+            "import javax.annotation.Nullable;",
+            "class Test {",
+            "  private void foo(@Nullable Object o) {",
+            "    com.google.common.truth.Truth.assertThat(o).isNotNull();",
+            "    // BUG: Diagnostic contains: dereferenced expression",
+            "    o.toString();",
+            "  }",
+            "}")
+        .doTest();
+  }
+
   @Test
   public void supportSwitchExpression() {
     compilationHelper