Allow library models of the form null param -> null return (#407)

This is equivalent to what `@Contract(!null -> !null)` already
allows for first party code.

In particular, we use this to implement the nullability contract
of `Optional.orElse(...)` and resolve #406.
diff --git a/nullaway/src/main/java/com/uber/nullaway/LibraryModels.java b/nullaway/src/main/java/com/uber/nullaway/LibraryModels.java
index 0073e98..c7ac987 100644
--- a/nullaway/src/main/java/com/uber/nullaway/LibraryModels.java
+++ b/nullaway/src/main/java/com/uber/nullaway/LibraryModels.java
@@ -81,6 +81,19 @@
   ImmutableSetMultimap<MethodRef, Integer> nullImpliesFalseParameters();
 
   /**
+   * Get (method, parameter) pairs that cause the method to return <code>null</code> when passed
+   * <code>null</code> on that parameter.
+   *
+   * <p>This is equivalent to annotating a method with a contract like:
+   *
+   * <pre><code>@Contract("!null -&gt; !null") @Nullable</code></pre>
+   *
+   * @return map from the names of null-in-implies-null out methods to the indexes of the arguments
+   *     that determine nullness of the return.
+   */
+  ImmutableSetMultimap<MethodRef, Integer> nullImpliesNullParameters();
+
+  /**
    * Get the set of library methods that may return null.
    *
    * @return set of library methods that may return null
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 ee8bb25..fc19c8e 100644
--- a/nullaway/src/main/java/com/uber/nullaway/handlers/LibraryModelsHandler.java
+++ b/nullaway/src/main/java/com/uber/nullaway/handlers/LibraryModelsHandler.java
@@ -100,16 +100,18 @@
   @Override
   public boolean onOverrideMayBeNullExpr(
       NullAway analysis, ExpressionTree expr, VisitorState state, boolean exprMayBeNull) {
-    if (expr.getKind() == Tree.Kind.METHOD_INVOCATION
-        && getOptLibraryModels(state.context)
-            .hasNullableReturn(
-                (Symbol.MethodSymbol) ASTHelpers.getSymbol(expr), state.getTypes())) {
-      return analysis.nullnessFromDataflow(state, expr) || exprMayBeNull;
-    }
-    if (expr.getKind() == Tree.Kind.METHOD_INVOCATION
-        && getOptLibraryModels(state.context)
-            .hasNonNullReturn((Symbol.MethodSymbol) ASTHelpers.getSymbol(expr), state.getTypes())) {
-      return false;
+    if (expr.getKind() == Tree.Kind.METHOD_INVOCATION) {
+      OptimizedLibraryModels optLibraryModels = getOptLibraryModels(state.context);
+      Symbol.MethodSymbol methodSymbol = (Symbol.MethodSymbol) ASTHelpers.getSymbol(expr);
+      if (optLibraryModels.hasNullableReturn(methodSymbol, state.getTypes())
+          || !optLibraryModels.nullImpliesNullParameters(methodSymbol).isEmpty()) {
+        // These mean the method might be null, depending on dataflow and arguments. We force
+        // dataflow to run.
+        return analysis.nullnessFromDataflow(state, expr) || exprMayBeNull;
+      } else if (optLibraryModels.hasNonNullReturn(methodSymbol, state.getTypes())) {
+        // This means the method can't be null, so we return false outright.
+        return false;
+      }
     }
     return exprMayBeNull;
   }
@@ -127,6 +129,21 @@
     Preconditions.checkNotNull(callee);
     setUnconditionalArgumentNullness(bothUpdates, node.getArguments(), callee, context);
     setConditionalArgumentNullness(thenUpdates, elseUpdates, node.getArguments(), callee, context);
+    ImmutableSet<Integer> nullImpliesNullIndexes =
+        getOptLibraryModels(context).nullImpliesNullParameters(callee);
+    if (!nullImpliesNullIndexes.isEmpty()) {
+      // If the method is marked as having argument dependent nullability and any of the
+      // corresponding arguments is null, then the return is nullable. If the method is
+      // marked as having argument dependent nullability but NONE of the corresponding
+      // arguments is null, then the return should be non-null.
+      boolean anyNull = false;
+      for (int idx : nullImpliesNullIndexes) {
+        if (!inputs.valueOfSubNode(node.getArgument(idx)).equals(NONNULL)) {
+          anyNull = true;
+        }
+      }
+      return anyNull ? NullnessHint.HINT_NULLABLE : NullnessHint.FORCE_NONNULL;
+    }
     if (getOptLibraryModels(context).hasNonNullReturn(callee, types)) {
       return NullnessHint.FORCE_NONNULL;
     } else if (getOptLibraryModels(context).hasNullableReturn(callee, types)) {
@@ -459,6 +476,11 @@
             .put(methodRef("java.util.Objects", "nonNull(java.lang.Object)"), 0)
             .build();
 
+    private static final ImmutableSetMultimap<MethodRef, Integer> NULL_IMPLIES_NULL_PARAMETERS =
+        new ImmutableSetMultimap.Builder<MethodRef, Integer>()
+            .put(methodRef("java.util.Optional", "orElse(T)"), 0)
+            .build();
+
     private static final ImmutableSet<MethodRef> NULLABLE_RETURNS =
         new ImmutableSet.Builder<MethodRef>()
             .add(methodRef("java.lang.ref.Reference", "get()"))
@@ -554,6 +576,11 @@
     }
 
     @Override
+    public ImmutableSetMultimap<MethodRef, Integer> nullImpliesNullParameters() {
+      return NULL_IMPLIES_NULL_PARAMETERS;
+    }
+
+    @Override
     public ImmutableSet<MethodRef> nullableReturns() {
       return NULLABLE_RETURNS;
     }
@@ -576,6 +603,8 @@
 
     private final ImmutableSetMultimap<MethodRef, Integer> nullImpliesFalseParameters;
 
+    private final ImmutableSetMultimap<MethodRef, Integer> nullImpliesNullParameters;
+
     private final ImmutableSet<MethodRef> nullableReturns;
 
     private final ImmutableSet<MethodRef> nonNullReturns;
@@ -591,6 +620,8 @@
           new ImmutableSetMultimap.Builder<>();
       ImmutableSetMultimap.Builder<MethodRef, Integer> nullImpliesFalseParametersBuilder =
           new ImmutableSetMultimap.Builder<>();
+      ImmutableSetMultimap.Builder<MethodRef, Integer> nullImpliesNullParametersBuilder =
+          new ImmutableSetMultimap.Builder<>();
       ImmutableSet.Builder<MethodRef> nullableReturnsBuilder = new ImmutableSet.Builder<>();
       ImmutableSet.Builder<MethodRef> nonNullReturnsBuilder = new ImmutableSet.Builder<>();
       for (LibraryModels libraryModels : models) {
@@ -612,6 +643,10 @@
             libraryModels.nullImpliesFalseParameters().entries()) {
           nullImpliesFalseParametersBuilder.put(entry);
         }
+        for (Map.Entry<MethodRef, Integer> entry :
+            libraryModels.nullImpliesNullParameters().entries()) {
+          nullImpliesNullParametersBuilder.put(entry);
+        }
         for (MethodRef name : libraryModels.nullableReturns()) {
           nullableReturnsBuilder.add(name);
         }
@@ -624,6 +659,7 @@
       nonNullParameters = nonNullParametersBuilder.build();
       nullImpliesTrueParameters = nullImpliesTrueParametersBuilder.build();
       nullImpliesFalseParameters = nullImpliesFalseParametersBuilder.build();
+      nullImpliesNullParameters = nullImpliesNullParametersBuilder.build();
       nullableReturns = nullableReturnsBuilder.build();
       nonNullReturns = nonNullReturnsBuilder.build();
     }
@@ -654,6 +690,11 @@
     }
 
     @Override
+    public ImmutableSetMultimap<MethodRef, Integer> nullImpliesNullParameters() {
+      return nullImpliesNullParameters;
+    }
+
+    @Override
     public ImmutableSet<MethodRef> nullableReturns() {
       return nullableReturns;
     }
@@ -705,6 +746,7 @@
     private final NameIndexedMap<ImmutableSet<Integer>> nonNullParams;
     private final NameIndexedMap<ImmutableSet<Integer>> nullImpliesTrueParams;
     private final NameIndexedMap<ImmutableSet<Integer>> nullImpliesFalseParams;
+    private final NameIndexedMap<ImmutableSet<Integer>> nullImpliesNullParams;
     private final NameIndexedMap<Boolean> nullableRet;
     private final NameIndexedMap<Boolean> nonNullRet;
 
@@ -717,6 +759,7 @@
       nullImpliesTrueParams = makeOptimizedIntSetLookup(names, models.nullImpliesTrueParameters());
       nullImpliesFalseParams =
           makeOptimizedIntSetLookup(names, models.nullImpliesFalseParameters());
+      nullImpliesNullParams = makeOptimizedIntSetLookup(names, models.nullImpliesNullParameters());
       nullableRet = makeOptimizedBoolLookup(names, models.nullableReturns());
       nonNullRet = makeOptimizedBoolLookup(names, models.nonNullReturns());
     }
@@ -749,6 +792,10 @@
       return lookupImmutableSet(symbol, nullImpliesFalseParams);
     }
 
+    ImmutableSet<Integer> nullImpliesNullParameters(Symbol.MethodSymbol symbol) {
+      return lookupImmutableSet(symbol, nullImpliesNullParams);
+    }
+
     private ImmutableSet<Integer> lookupImmutableSet(
         Symbol.MethodSymbol symbol, NameIndexedMap<ImmutableSet<Integer>> lookup) {
       ImmutableSet<Integer> result = lookup.get(symbol);
diff --git a/nullaway/src/test/java/com/uber/nullaway/NullAwayTest.java b/nullaway/src/test/java/com/uber/nullaway/NullAwayTest.java
index b3a6611..d06e1f4 100644
--- a/nullaway/src/test/java/com/uber/nullaway/NullAwayTest.java
+++ b/nullaway/src/test/java/com/uber/nullaway/NullAwayTest.java
@@ -436,6 +436,10 @@
             "  String test2(Object o2) {",
             "    return NullnessChecker.noOp(o2).toString();",
             "  }",
+            "  Object test3(@Nullable Object o1) {",
+            "    // BUG: Diagnostic contains: returning @Nullable expression",
+            "    return NullnessChecker.noOp(o1);",
+            "  }",
             "}")
         .doTest();
   }
@@ -2524,4 +2528,39 @@
             "}")
         .doTest();
   }
+
+  @Test
+  public void orElseLibraryModelSupport() {
+    // Checks both Optional.orElse(...) support itself and the general nullImpliesNullParameters
+    // Library Models mechanism for encoding @Contract(!null -> !null) as a library model.
+    compilationHelper
+        .addSourceLines(
+            "TestOptionalOrElseNegative.java",
+            "package com.uber;",
+            "import javax.annotation.Nullable;",
+            "import java.util.Optional;",
+            "class TestOptionalOrElseNegative {",
+            "  public Object foo(Optional<Object> o) {",
+            "    return o.orElse(\"Something\");",
+            "  }",
+            "  public @Nullable Object bar(Optional<Object> o) {",
+            "    return o.orElse(null);",
+            "  }",
+            "}")
+        .addSourceLines(
+            "TestOptionalOrElsePositive.java",
+            "package com.uber;",
+            "import java.util.Optional;",
+            "class TestOptionalOrElsePositive {",
+            "  public Object foo(Optional<Object> o) {",
+            "    // BUG: Diagnostic contains: returning @Nullable expression",
+            "    return o.orElse(null);",
+            "  }",
+            "  public void bar(Optional<Object> o) {",
+            "    // BUG: Diagnostic contains: dereferenced expression o.orElse(null) is @Nullable",
+            "    System.out.println(o.orElse(null).toString());",
+            "  }",
+            "}")
+        .doTest();
+  }
 }
diff --git a/sample-library-model/src/main/java/com/uber/modelexample/ExampleLibraryModels.java b/sample-library-model/src/main/java/com/uber/modelexample/ExampleLibraryModels.java
index ff9d707..ff78dac 100644
--- a/sample-library-model/src/main/java/com/uber/modelexample/ExampleLibraryModels.java
+++ b/sample-library-model/src/main/java/com/uber/modelexample/ExampleLibraryModels.java
@@ -53,6 +53,11 @@
   }
 
   @Override
+  public ImmutableSetMultimap<MethodRef, Integer> nullImpliesNullParameters() {
+    return ImmutableSetMultimap.of();
+  }
+
+  @Override
   public ImmutableSet<MethodRef> nullableReturns() {
     return ImmutableSet.of();
   }
diff --git a/test-library-models/src/main/java/com/uber/nullaway/testlibrarymodels/TestLibraryModels.java b/test-library-models/src/main/java/com/uber/nullaway/testlibrarymodels/TestLibraryModels.java
index 1bb60c2..6e9f0a5 100644
--- a/test-library-models/src/main/java/com/uber/nullaway/testlibrarymodels/TestLibraryModels.java
+++ b/test-library-models/src/main/java/com/uber/nullaway/testlibrarymodels/TestLibraryModels.java
@@ -57,6 +57,11 @@
   }
 
   @Override
+  public ImmutableSetMultimap<MethodRef, Integer> nullImpliesNullParameters() {
+    return ImmutableSetMultimap.of();
+  }
+
+  @Override
   public ImmutableSet<MethodRef> nullableReturns() {
     return ImmutableSet.of();
   }