Add support for Preconditions.checkArgument through custom CFG translation

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

@@ -80,7 +80,7 @@ private AccessPathNullnessAnalysis(
             config,
             handler,
             new CoreNullnessStoreInitializer());
-    this.dataFlow = new DataFlow(config.assertsEnabled());
+    this.dataFlow = new DataFlow(config.assertsEnabled(), handler);
 
     if (config.checkContracts()) {
       this.contractNullnessPropagation =

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

@@ -38,6 +38,8 @@
 import com.sun.tools.javac.processing.JavacProcessingEnvironment;
 import com.sun.tools.javac.util.Context;
 import com.uber.nullaway.NullabilityUtil;
+import com.uber.nullaway.dataflow.cfg.NullAwayCFGBuilder;
+import com.uber.nullaway.handlers.Handler;
 import javax.annotation.Nullable;
 import javax.annotation.processing.ProcessingEnvironment;
 import org.checkerframework.nullaway.dataflow.analysis.AbstractValue;
@@ -49,7 +51,6 @@
 import org.checkerframework.nullaway.dataflow.analysis.TransferFunction;
 import org.checkerframework.nullaway.dataflow.cfg.ControlFlowGraph;
 import org.checkerframework.nullaway.dataflow.cfg.UnderlyingAST;
-import org.checkerframework.nullaway.dataflow.cfg.builder.CFGBuilder;
 
 /**
  * Provides a wrapper around {@link org.checkerframework.nullaway.dataflow.analysis.Analysis}.
@@ -70,8 +71,11 @@ public final class DataFlow {
 
   private final boolean assertsEnabled;
 
-  DataFlow(boolean assertsEnabled) {
+  private final Handler handler;
+
+  DataFlow(boolean assertsEnabled, Handler handler) {
     this.assertsEnabled = assertsEnabled;
+    this.handler = handler;
   }
 
   private final LoadingCache<AnalysisParams, Analysis<?, ?, ?>> analysisCache =
@@ -130,7 +134,8 @@ public ControlFlowGraph load(CfgParams key) {
                     bodyPath = codePath;
                   }
 
-                  return CFGBuilder.build(bodyPath, ast, assertsEnabled, !assertsEnabled, env);
+                  return NullAwayCFGBuilder.build(
+                      bodyPath, ast, assertsEnabled, !assertsEnabled, env, handler);
                 }
               });
 

nullaway/src/main/java/com/uber/nullaway/dataflow/cfg/NullAwayCFGBuilder.java

@@ -0,0 +1,185 @@
+package com.uber.nullaway.dataflow.cfg;
+
+import com.google.common.base.Preconditions;
+import com.sun.source.tree.ExpressionTree;
+import com.sun.source.tree.MethodInvocationTree;
+import com.sun.source.tree.ThrowTree;
+import com.sun.source.tree.Tree;
+import com.sun.source.tree.TreeVisitor;
+import com.sun.source.util.TreePath;
+import com.uber.nullaway.handlers.Handler;
+import javax.annotation.processing.ProcessingEnvironment;
+import javax.lang.model.type.TypeMirror;
+import org.checkerframework.nullaway.dataflow.cfg.ControlFlowGraph;
+import org.checkerframework.nullaway.dataflow.cfg.UnderlyingAST;
+import org.checkerframework.nullaway.dataflow.cfg.builder.CFGBuilder;
+import org.checkerframework.nullaway.dataflow.cfg.builder.CFGTranslationPhaseOne;
+import org.checkerframework.nullaway.dataflow.cfg.builder.CFGTranslationPhaseThree;
+import org.checkerframework.nullaway.dataflow.cfg.builder.CFGTranslationPhaseTwo;
+import org.checkerframework.nullaway.dataflow.cfg.builder.ConditionalJump;
+import org.checkerframework.nullaway.dataflow.cfg.builder.ExtendedNode;
+import org.checkerframework.nullaway.dataflow.cfg.builder.Label;
+import org.checkerframework.nullaway.dataflow.cfg.builder.PhaseOneResult;
+import org.checkerframework.nullaway.dataflow.cfg.node.MethodInvocationNode;
+import org.checkerframework.nullaway.dataflow.cfg.node.Node;
+import org.checkerframework.nullaway.dataflow.cfg.node.ThrowNode;
+import org.checkerframework.nullaway.javacutil.AnnotationProvider;
+import org.checkerframework.nullaway.javacutil.BasicAnnotationProvider;
+import org.checkerframework.nullaway.javacutil.trees.TreeBuilder;
+
+/**
+ * A NullAway specific CFGBuilder subclass, which allows to more directly control the AST to CFG
+ * translation performed by the checker framework.
+ *
+ * <p>This holds the static method {@link #build(TreePath, UnderlyingAST, boolean, boolean,
+ * ProcessingEnvironment, Handler)}, called to perform the CFG translation, and the class {@link
+ * NullAwayCFGTranslationPhaseOne}, which extends {@link CFGTranslationPhaseOne} and adds hooks for
+ * the NullAway handlers mechanism and some utility methods.
+ */
+public final class NullAwayCFGBuilder extends CFGBuilder {
+
+  /** This class should never be instantiated. */
+  private NullAwayCFGBuilder() {}
+
+  /**
+   * This static method produces a new CFG representation given a method's (or lambda/initializer)
+   * body.
+   *
+   * <p>It is analogous to {@link CFGBuilder#build(TreePath, UnderlyingAST, boolean, boolean,
+   * ProcessingEnvironment)}, but it also takes a handler to be called at specific extention points
+   * during the CFG translation.
+   *
+   * @param bodyPath the TreePath to the body of the method, lambda, or initializer.
+   * @param underlyingAST the AST that underlies the control frow graph
+   * @param assumeAssertionsEnabled can assertions be assumed to be disabled?
+   * @param assumeAssertionsDisabled can assertions be assumed to be enabled?
+   * @param env annotation processing environment containing type utilities
+   * @param handler a NullAway handler or chain of handlers (through {@link
+   *     com.uber.nullaway.handlers.CompositeHandler}
+   * @return a control flow graph
+   */
+  public static ControlFlowGraph build(
+      TreePath bodyPath,
+      UnderlyingAST underlyingAST,
+      boolean assumeAssertionsEnabled,
+      boolean assumeAssertionsDisabled,
+      ProcessingEnvironment env,
+      Handler handler) {
+    TreeBuilder builder = new TreeBuilder(env);
+    AnnotationProvider annotationProvider = new BasicAnnotationProvider();
+    CFGTranslationPhaseOne phase1translator =
+        new NullAwayCFGTranslationPhaseOne(
+            builder,
+            annotationProvider,
+            assumeAssertionsEnabled,
+            assumeAssertionsDisabled,
+            env,
+            handler);
+    PhaseOneResult phase1result = phase1translator.process(bodyPath, underlyingAST);
+    ControlFlowGraph phase2result = CFGTranslationPhaseTwo.process(phase1result);
+    ControlFlowGraph phase3result = CFGTranslationPhaseThree.process(phase2result);
+    return phase3result;
+  }
+
+  /**
+   * A NullAway specific subclass of the Checker Framework's {@link CFGTranslationPhaseOne},
+   * augmented with handler extension points and some utility methods meant to be called from
+   * handlers to customize the AST to CFG translation.
+   */
+  public static class NullAwayCFGTranslationPhaseOne extends CFGTranslationPhaseOne {
+
+    private final Handler handler;
+
+    /**
+     * Create a new NullAway phase one translation visitor.
+     *
+     * @param builder a TreeBuilder object (used to create synthetic AST nodes to feed to the
+     *     translation process)
+     * @param annotationProvider an {@link AnnotationProvider}.
+     * @param assumeAssertionsEnabled can assertions be assumed to be disabled?
+     * @param assumeAssertionsDisabled can assertions be assumed to be enabled?
+     * @param env annotation processing environment containing type utilities
+     * @param handler a NullAway handler or chain of handlers (through {@link
+     *     com.uber.nullaway.handlers.CompositeHandler}
+     */
+    public NullAwayCFGTranslationPhaseOne(
+        TreeBuilder builder,
+        AnnotationProvider annotationProvider,
+        boolean assumeAssertionsEnabled,
+        boolean assumeAssertionsDisabled,
+        ProcessingEnvironment env,
+        Handler handler) {
+      super(builder, annotationProvider, assumeAssertionsEnabled, assumeAssertionsDisabled, env);
+      this.handler = handler;
+    }
+
+    @SuppressWarnings("NullAway") // (Void)null issue
+    private void scanWithVoid(Tree tree) {
+      this.scan(tree, (Void) null);
+    }
+
+    /**
+     * Obtain the type mirror for a given class, used for exception throwing.
+     *
+     * @param klass a Java class
+     * @return the corresponding type mirror
+     */
+    public TypeMirror classToErrorType(Class<?> klass) {
+      return this.getTypeMirror(klass);
+    }
+
+    /**
+     * Extend the CFG to throw an exception if the passed expression node evaluates to false.
+     *
+     * @param booleanExpressionNode a CFG Node representing a boolean expression.
+     * @param errorType the type of the exception to throw if booleanExpressionNode evaluates to
+     *     false.
+     */
+    public void insertThrowOnFalse(Node booleanExpressionNode, TypeMirror errorType) {
+      Tree tree = booleanExpressionNode.getTree();
+      Preconditions.checkArgument(
+          tree instanceof ExpressionTree,
+          "Argument booleanExpressionNode must represent a boolean expression");
+      ExpressionTree booleanExpressionTree = (ExpressionTree) booleanExpressionNode.getTree();
+      Preconditions.checkNotNull(booleanExpressionTree);
+      Label falsePreconditionEntry = new Label();
+      Label endPrecondition = new Label();
+      this.scanWithVoid(booleanExpressionTree);
+      ConditionalJump cjump = new ConditionalJump(endPrecondition, falsePreconditionEntry);
+      this.extendWithExtendedNode(cjump);
+      this.addLabelForNextNode(falsePreconditionEntry);
+      ExtendedNode exNode =
+          this.extendWithNodeWithException(
+              new ThrowNode(
+                  new ThrowTree() {
+                    // Dummy throw tree, unused. We could use null here, but that violates nullness
+                    // typing.
+                    @Override
+                    public ExpressionTree getExpression() {
+                      return booleanExpressionTree;
+                    }
+
+                    @Override
+                    public Kind getKind() {
+                      return Kind.THROW;
+                    }
+
+                    @Override
+                    public <R, D> R accept(TreeVisitor<R, D> visitor, D data) {
+                      return visitor.visitThrow(this, data);
+                    }
+                  },
+                  booleanExpressionNode,
+                  this.getProcessingEnvironment().getTypeUtils()),
+              errorType);
+      exNode.setTerminatesExecution(true);
+      this.addLabelForNextNode(endPrecondition);
+    }
+
+    @Override
+    public MethodInvocationNode visitMethodInvocation(MethodInvocationTree tree, Void p) {
+      MethodInvocationNode originalNode = super.visitMethodInvocation(tree, p);
+      return handler.onCFGBuildPhase1AfterVisitMethodInvocation(this, tree, originalNode);
+    }
+  }
+}

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

@@ -40,6 +40,7 @@
 import com.uber.nullaway.dataflow.AccessPathNullnessAnalysis;
 import com.uber.nullaway.dataflow.AccessPathNullnessPropagation;
 import com.uber.nullaway.dataflow.NullnessStore;
+import com.uber.nullaway.dataflow.cfg.NullAwayCFGBuilder;
 import java.util.List;
 import java.util.Optional;
 import org.checkerframework.nullaway.dataflow.cfg.UnderlyingAST;
@@ -193,4 +194,12 @@ public void onNonNullFieldAssignment(
       Symbol field, AccessPathNullnessAnalysis analysis, VisitorState state) {
     // NoOp
   }
+
+  @Override
+  public MethodInvocationNode onCFGBuildPhase1AfterVisitMethodInvocation(
+      NullAwayCFGBuilder.NullAwayCFGTranslationPhaseOne phase,
+      MethodInvocationTree tree,
+      MethodInvocationNode originalNode) {
+    return originalNode;
+  }
 }

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

@@ -41,6 +41,7 @@
 import com.uber.nullaway.dataflow.AccessPathNullnessAnalysis;
 import com.uber.nullaway.dataflow.AccessPathNullnessPropagation;
 import com.uber.nullaway.dataflow.NullnessStore;
+import com.uber.nullaway.dataflow.cfg.NullAwayCFGBuilder;
 import java.util.List;
 import java.util.Optional;
 import org.checkerframework.nullaway.dataflow.cfg.UnderlyingAST;
@@ -251,4 +252,16 @@ public void onNonNullFieldAssignment(
       h.onNonNullFieldAssignment(field, analysis, state);
     }
   }
+
+  @Override
+  public MethodInvocationNode onCFGBuildPhase1AfterVisitMethodInvocation(
+      NullAwayCFGBuilder.NullAwayCFGTranslationPhaseOne phase,
+      MethodInvocationTree tree,
+      MethodInvocationNode originalNode) {
+    MethodInvocationNode currentNode = originalNode;
+    for (Handler h : handlers) {
+      currentNode = h.onCFGBuildPhase1AfterVisitMethodInvocation(phase, tree, currentNode);
+    }
+    return currentNode;
+  }
 }

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

@@ -41,6 +41,7 @@
 import com.uber.nullaway.dataflow.AccessPathNullnessAnalysis;
 import com.uber.nullaway.dataflow.AccessPathNullnessPropagation;
 import com.uber.nullaway.dataflow.NullnessStore;
+import com.uber.nullaway.dataflow.cfg.NullAwayCFGBuilder;
 import java.util.List;
 import java.util.Optional;
 import org.checkerframework.nullaway.dataflow.cfg.UnderlyingAST;
@@ -325,6 +326,22 @@ Optional<ErrorMessage> onExpressionDereference(
   void onNonNullFieldAssignment(
       Symbol field, AccessPathNullnessAnalysis analysis, VisitorState state);
 
+  /**
+   * Called during AST to CFG translation (CFGTranslationPhaseOne) immediately after translating a
+   * MethodInvocationTree.
+   *
+   * @param phase a reference to the NullAwayCFGTranslationPhaseOne object and its utility
+   *     functions.
+   * @param tree the MethodInvocationTree being translated.
+   * @param originalNode the resulting MethodInvocationNode right before this handler is called.
+   * @return a MethodInvocationNode which might be originalNode or a modified version, this is
+   *     passed to the next handler in the chain.
+   */
+  MethodInvocationNode onCFGBuildPhase1AfterVisitMethodInvocation(
+      NullAwayCFGBuilder.NullAwayCFGTranslationPhaseOne phase,
+      MethodInvocationTree tree,
+      MethodInvocationNode originalNode);
+
   /**
    * A three value enum for handlers implementing onDataflowVisitMethodInvocation to communicate
    * their knowledge of the method return nullability to the the rest of NullAway.

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

@@ -55,6 +55,7 @@ public static Handler buildDefault(Config config) {
     if (config.handleTestAssertionLibraries()) {
       handlerListBuilder.add(new AssertionHandler(methodNameUtil));
     }
+    handlerListBuilder.add(new PreconditionsHandler());
     handlerListBuilder.add(new LibraryModelsHandler(config));
     handlerListBuilder.add(StreamNullabilityPropagatorFactory.getRxStreamNullabilityPropagator());
     handlerListBuilder.add(StreamNullabilityPropagatorFactory.getJavaStreamNullabilityPropagator());

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

@@ -0,0 +1,44 @@
+package com.uber.nullaway.handlers;
+
+import com.google.common.base.Preconditions;
+import com.google.errorprone.util.ASTHelpers;
+import com.sun.source.tree.MethodInvocationTree;
+import com.sun.tools.javac.code.Symbol;
+import com.uber.nullaway.dataflow.cfg.NullAwayCFGBuilder;
+import javax.annotation.Nullable;
+import javax.lang.model.element.Name;
+import javax.lang.model.type.TypeMirror;
+import org.checkerframework.nullaway.dataflow.cfg.node.MethodInvocationNode;
+
+public class PreconditionsHandler extends BaseNoOpHandler {
+
+  private static final String PRECONDITIONS_CLASS_NAME = "com.google.common.base.Preconditions";
+  private static final String CHECK_ARGUMENT_METHOD_NAME = "checkArgument";
+  private static final String CHECK_STATE_METHOD_NAME = "checkState";
+
+  @Nullable private Name preconditionsClass;
+  @Nullable private Name checkArgumentMethod;
+  @Nullable private Name checkStateMethod;
+  @Nullable TypeMirror preconditionErrorType;
+
+  @Override
+  public MethodInvocationNode onCFGBuildPhase1AfterVisitMethodInvocation(
+      NullAwayCFGBuilder.NullAwayCFGTranslationPhaseOne phase,
+      MethodInvocationTree tree,
+      MethodInvocationNode originalNode) {
+    Symbol.MethodSymbol callee = ASTHelpers.getSymbol(tree);
+    if (preconditionsClass == null) {
+      preconditionsClass = callee.name.table.fromString(PRECONDITIONS_CLASS_NAME);
+      checkArgumentMethod = callee.name.table.fromString(CHECK_ARGUMENT_METHOD_NAME);
+      checkStateMethod = callee.name.table.fromString(CHECK_STATE_METHOD_NAME);
+      preconditionErrorType = phase.classToErrorType(IllegalArgumentException.class);
+    }
+    Preconditions.checkNotNull(preconditionErrorType);
+    if (callee.enclClass().getQualifiedName().equals(preconditionsClass)
+        && (callee.name.equals(checkArgumentMethod) || callee.name.equals(checkStateMethod))
+        && callee.getParameters().size() > 0) {
+      phase.insertThrowOnFalse(originalNode.getArgument(0), preconditionErrorType);
+    }
+    return originalNode;
+  }
+}