Allow library models to define custom stream classes.

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

@@ -22,9 +22,11 @@
 
 package com.uber.nullaway;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSetMultimap;
 import com.sun.tools.javac.code.Symbol;
+import com.uber.nullaway.handlers.stream.StreamTypeRecord;
 import java.util.Objects;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -130,7 +132,23 @@ public interface LibraryModels {
   ImmutableSetMultimap<MethodRef, Integer> castToNonNullMethods();
 
   /**
-   * representation of a method as a qualified class name + a signature for the method
+   * Get a list of custom stream library specifications.
+   *
+   * <p>This allows users to define filter/map/other methods for APIs which behave similarly to Java
+   * 8 streams or ReactiveX streams, so that NullAway is able to understand nullability invariants
+   * across stream API calls. See {@see com.uber.nullaway.handlers.stream.StreamModelBuilder} for
+   * details on how to construct these {@see com.uber.nullaway.handlers.stream.StreamTypeRecord}
+   * specs. A full example is available at {@see
+   * com.uber.nullaway.testlibrarymodels.TestLibraryModels}.
+   *
+   * @return A list of StreamTypeRecord specs (usually generated using StreamModelBuilder).
+   */
+  default ImmutableList<StreamTypeRecord> customStreamNullabilitySpecs() {
+    return ImmutableList.of();
+  }
+
+  /**
+   * Representation of a method as a qualified class name + a signature for the method
    *
    * <p>The formatting of a method signature should match the result of calling {@link
    * Symbol.MethodSymbol#toString()} on the corresponding symbol. See {@link

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

@@ -57,9 +57,11 @@ public static Handler buildDefault(Config config) {
       handlerListBuilder.add(new AssertionHandler(methodNameUtil));
     }
     handlerListBuilder.add(new GuavaAssertionsHandler());
-    handlerListBuilder.add(new LibraryModelsHandler(config));
+    LibraryModelsHandler libraryModelsHandler = new LibraryModelsHandler(config);
+    handlerListBuilder.add(libraryModelsHandler);
     handlerListBuilder.add(StreamNullabilityPropagatorFactory.getRxStreamNullabilityPropagator());
     handlerListBuilder.add(StreamNullabilityPropagatorFactory.getJavaStreamNullabilityPropagator());
+    handlerListBuilder.add(libraryModelsHandler.getStreamNullabilityPropagatorFromModels());
     handlerListBuilder.add(new ContractHandler(config));
     handlerListBuilder.add(new ApacheThriftIsSetHandler());
     handlerListBuilder.add(new GrpcHandler());

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

@@ -27,6 +27,7 @@
 import static com.uber.nullaway.Nullness.NULLABLE;
 
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSetMultimap;
 import com.google.errorprone.VisitorState;
@@ -46,6 +47,7 @@
 import com.uber.nullaway.Nullness;
 import com.uber.nullaway.dataflow.AccessPath;
 import com.uber.nullaway.dataflow.AccessPathNullnessPropagation;
+import com.uber.nullaway.handlers.stream.StreamTypeRecord;
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
@@ -291,6 +293,21 @@ private void setUnconditionalArgumentNullness(
     }
   }
 
+  /**
+   * Generate a StreamNullabilityPropagator handler based on all the stream specifications loaded
+   * from any of our library models.
+   *
+   * <p>This handler must be registered in Handlers.java separatedly from the LibraryModelsHandler
+   * itself.
+   */
+  public StreamNullabilityPropagator getStreamNullabilityPropagatorFromModels() {
+    // Note: Currently, OptimizedLibraryModels doesn't carry the information about stream type
+    // records,
+    // it is not clear what it means to "optimize" lookup for those and they get access only once:
+    // here.
+    return new StreamNullabilityPropagator(libraryModels.customStreamNullabilitySpecs());
+  }
+
   private static LibraryModels loadLibraryModels(Config config) {
     Iterable<LibraryModels> externalLibraryModels =
         ServiceLoader.load(LibraryModels.class, LibraryModels.class.getClassLoader());
@@ -827,6 +844,8 @@ private static class CombinedLibraryModels implements LibraryModels {
 
     private final ImmutableSetMultimap<MethodRef, Integer> castToNonNullMethods;
 
+    private final ImmutableList<StreamTypeRecord> customStreamNullabilitySpecs;
+
     public CombinedLibraryModels(Iterable<LibraryModels> models, Config config) {
       this.config = config;
       ImmutableSetMultimap.Builder<MethodRef, Integer> failIfNullParametersBuilder =
@@ -845,6 +864,8 @@ public CombinedLibraryModels(Iterable<LibraryModels> models, Config config) {
       ImmutableSet.Builder<MethodRef> nonNullReturnsBuilder = new ImmutableSet.Builder<>();
       ImmutableSetMultimap.Builder<MethodRef, Integer> castToNonNullMethodsBuilder =
           new ImmutableSetMultimap.Builder<>();
+      ImmutableList.Builder<StreamTypeRecord> customStreamNullabilitySpecsBuilder =
+          new ImmutableList.Builder<>();
       for (LibraryModels libraryModels : models) {
         for (Map.Entry<MethodRef, Integer> entry : libraryModels.failIfNullParameters().entries()) {
           if (shouldSkipModel(entry.getKey())) {
@@ -904,6 +925,9 @@ public CombinedLibraryModels(Iterable<LibraryModels> models, Config config) {
           }
           castToNonNullMethodsBuilder.put(entry);
         }
+        for (StreamTypeRecord streamTypeRecord : libraryModels.customStreamNullabilitySpecs()) {
+          customStreamNullabilitySpecsBuilder.add(streamTypeRecord);
+        }
       }
       failIfNullParameters = failIfNullParametersBuilder.build();
       explicitlyNullableParameters = explicitlyNullableParametersBuilder.build();
@@ -914,6 +938,7 @@ public CombinedLibraryModels(Iterable<LibraryModels> models, Config config) {
       nullableReturns = nullableReturnsBuilder.build();
       nonNullReturns = nonNullReturnsBuilder.build();
       castToNonNullMethods = castToNonNullMethodsBuilder.build();
+      customStreamNullabilitySpecs = customStreamNullabilitySpecsBuilder.build();
     }
 
     private boolean shouldSkipModel(MethodRef key) {
@@ -964,6 +989,11 @@ public ImmutableSet<MethodRef> nonNullReturns() {
     public ImmutableSetMultimap<MethodRef, Integer> castToNonNullMethods() {
       return castToNonNullMethods;
     }
+
+    @Override
+    public ImmutableList<StreamTypeRecord> customStreamNullabilitySpecs() {
+      return customStreamNullabilitySpecs;
+    }
   }
 
   /**

nullaway/src/main/java/com/uber/nullaway/handlers/stream/StreamModelBuilder.java

@@ -25,6 +25,8 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.predicates.TypePredicate;
+import com.google.errorprone.predicates.type.DescendantOf;
+import com.google.errorprone.suppliers.Suppliers;
 import java.util.ArrayList;
 import java.util.List;
 import javax.annotation.Nullable;
@@ -90,6 +92,10 @@ public StreamModelBuilder addStreamType(TypePredicate tp) {
     return this;
   }
 
+  public StreamModelBuilder addStreamTypeFromName(String fullyQualifiedName) {
+    return this.addStreamType(new DescendantOf(Suppliers.typeFromString(fullyQualifiedName)));
+  }
+
   private void initializeBuilders() {
     this.filterMethodSigs = ImmutableSet.builder();
     this.filterMethodSimpleNames = ImmutableSet.builder();

nullaway/src/test/resources/com/uber/nullaway/testdata/NullAwayStreamSupportNegativeCases.java

@@ -24,6 +24,7 @@
 
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
+import com.uber.nullaway.testdata.unannotated.CustomStream;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.DoubleStream;
@@ -288,6 +289,19 @@ private void filterThenForEachOrdered(Stream<NullableContainer<String>> stream)
     stream.filter(s -> s.get() != null).forEachOrdered(s -> System.out.println(s.get().length()));
   }
 
+  // CustomStream is modeled in TestLibraryModels
+  private CustomStream<Integer> filterThenMapLambdasCustomStream(CustomStream<String> stream) {
+    return stream.filter(s -> s != null).map(s -> s.length());
+  }
+
+  private CustomStream<Integer> filterThenMapMethodRefsCustomStream(
+      CustomStream<NullableContainer<String>> stream) {
+    return stream
+        .filter(c -> c.get() != null && perhaps())
+        .map(NullableContainer::get)
+        .map(String::length);
+  }
+
   private static class CheckFinalBeforeStream<T> {
     @Nullable private final T ref;
 

nullaway/src/test/resources/com/uber/nullaway/testdata/unannotated/CustomStream.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright (c) 2023 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.testdata.unannotated;
+
+import java.util.function.Function;
+import java.util.function.Predicate;
+
+/**
+ * A class representing a custom implementation of java.util.stream.Stream to test the ability to
+ * define stream handler specs through library models.
+ */
+public class CustomStream<T> {
+
+  private CustomStream() {}
+
+  public CustomStream<T> filter(Predicate<? super T> predicate) {
+    throw new UnsupportedOperationException();
+  }
+
+  public <R> CustomStream<R> map(Function<? super T, ? extends R> mapper) {
+    throw new UnsupportedOperationException();
+  }
+}

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

@@ -18,9 +18,12 @@
 import static com.uber.nullaway.LibraryModels.MethodRef.methodRef;
 
 import com.google.auto.service.AutoService;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSetMultimap;
 import com.uber.nullaway.LibraryModels;
+import com.uber.nullaway.handlers.stream.StreamModelBuilder;
+import com.uber.nullaway.handlers.stream.StreamTypeRecord;
 
 @AutoService(LibraryModels.class)
 public class TestLibraryModels implements LibraryModels {
@@ -87,4 +90,36 @@ public ImmutableSetMultimap<MethodRef, Integer> castToNonNullMethods() {
             1)
         .build();
   }
+
+  @Override
+  public ImmutableList<StreamTypeRecord> customStreamNullabilitySpecs() {
+    // Identical to the default model for java.util.stream.Stream, but with the original type
+    // renamed
+    return StreamModelBuilder.start()
+        .addStreamTypeFromName("com.uber.unannotated.CustomStreamA")
+        .withFilterMethodFromSignature("filter(java.util.function.Predicate<? super T>)")
+        .withMapMethodFromSignature(
+            "<R>map(java.util.function.Function<? super T,? extends R>)",
+            "apply",
+            ImmutableSet.of(0))
+        .withMapMethodFromSignature(
+            "mapToInt(java.util.function.ToIntFunction<? super T>)",
+            "applyAsInt",
+            ImmutableSet.of(0))
+        .withMapMethodFromSignature(
+            "mapToLong(java.util.function.ToLongFunction<? super T>)",
+            "applyAsLong",
+            ImmutableSet.of(0))
+        .withMapMethodFromSignature(
+            "mapToDouble(java.util.function.ToDoubleFunction<? super T>)",
+            "applyAsDouble",
+            ImmutableSet.of(0))
+        .withMapMethodFromSignature(
+            "forEach(java.util.function.Consumer<? super T>)", "accept", ImmutableSet.of(0))
+        .withMapMethodFromSignature(
+            "forEachOrdered(java.util.function.Consumer<? super T>)", "accept", ImmutableSet.of(0))
+        .withMapMethodAllFromName("flatMap", "apply", ImmutableSet.of(0))
+        .withPassthroughMethodFromSignature("distinct()")
+        .end();
+  }
 }