Allow library models to define custom stream classes. (#807)

This is needed to support both fully custom stream libraries, as well as
backports/shadowed versions of existing stream libraries.

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 {@link com.uber.nullaway.handlers.stream.StreamModelBuilder} for
+   * details on how to construct these {@link com.uber.nullaway.handlers.stream.StreamTypeRecord}
+   * specs. A full example is available at {@link
+   * 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,13 @@ 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(
+        StreamNullabilityPropagatorFactory.fromSpecs(
+            libraryModelsHandler.getStreamNullabilitySpecs()));
     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,25 @@ private void setUnconditionalArgumentNullness(
     }
   }
 
+  /**
+   * Get all the stream specifications loaded from any of our library models.
+   *
+   * <p>This is used in Handlers.java to create a StreamNullabilityPropagator handler, which gets
+   * registered independently of this LibraryModelsHandler itself.
+   *
+   * <p>LibraryModelsHandler is responsible from reading the library models for stream specs, but
+   * beyond that, checking of the specs falls under the responsibility of the generated
+   * StreamNullabilityPropagator handler.
+   *
+   * @return The list of all stream specifications loaded from any of our library models.
+   */
+  public ImmutableList<StreamTypeRecord> getStreamNullabilitySpecs() {
+    // 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 accessed
+    // only once by calling this method during handler setup in Handlers.java.
+    return libraryModels.customStreamNullabilitySpecs();
+  }
+
   private static LibraryModels loadLibraryModels(Config config) {
     Iterable<LibraryModels> externalLibraryModels =
         ServiceLoader.load(LibraryModels.class, LibraryModels.class.getClassLoader());
@@ -827,6 +848,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 +868,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 +929,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 +942,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 +993,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/StreamNullabilityPropagatorFactory.java

@@ -29,6 +29,8 @@
 import com.uber.nullaway.handlers.stream.StreamTypeRecord;
 
 public class StreamNullabilityPropagatorFactory {
+
+  /** Returns a handler for the standard Java 8 stream APIs. */
   public static StreamNullabilityPropagator getJavaStreamNullabilityPropagator() {
     ImmutableList<StreamTypeRecord> streamModels =
         StreamModelBuilder.start()
@@ -78,6 +80,7 @@ public static StreamNullabilityPropagator getJavaStreamNullabilityPropagator() {
     return new StreamNullabilityPropagator(streamModels);
   }
 
+  /** Returns a handler for io.reactivex.* stream APIs */
   public static StreamNullabilityPropagator getRxStreamNullabilityPropagator() {
     ImmutableList<StreamTypeRecord> rxModels =
         StreamModelBuilder.start()
@@ -137,4 +140,19 @@ public static StreamNullabilityPropagator getRxStreamNullabilityPropagator() {
 
     return new StreamNullabilityPropagator(rxModels);
   }
+
+  /**
+   * Create a new StreamNullabilityPropagator from a list of StreamTypeRecord specs.
+   *
+   * <p>This is used to create a new StreamNullabilityPropagator based on stream API specs provided
+   * by library models.
+   *
+   * @param streamNullabilitySpecs the list of StreamTypeRecord objects defining one or more stream
+   *     APIs (from {@link StreamModelBuilder}).
+   * @return A handler corresponding to the stream APIs defined by the given specs.
+   */
+  public static StreamNullabilityPropagator fromSpecs(
+      ImmutableList<StreamTypeRecord> streamNullabilitySpecs) {
+    return new StreamNullabilityPropagator(streamNullabilitySpecs);
+  }
 }

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,16 @@ public StreamModelBuilder addStreamType(TypePredicate tp) {
     return this;
   }
 
+  /**
+   * Add a stream type to our models based on the type's fully qualified name.
+   *
+   * @param fullyQualifiedName the FQN of the class/interface in our stream-based API.
+   * @return This builder (for chaining).
+   */
+  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,26 @@ 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> filterThenMapNullableContainerLambdasCustomStream(
+            CustomStream<NullableContainer<String>> stream) {
+        return stream
+                .filter(c -> c.get() != null)
+                .map(c -> c.get().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/NullAwayStreamSupportPositiveCases.java

@@ -23,6 +23,7 @@
 package com.uber.nullaway.testdata;
 
 import com.google.common.base.Preconditions;
+import com.uber.nullaway.testdata.unannotated.CustomStreamWithoutModel;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.DoubleStream;
@@ -153,6 +154,28 @@ private void forEachOrdered(Stream<NullableContainer<String>> stream) {
     stream.forEachOrdered(s -> System.out.println(s.get().length()));
   }
 
+  // CustomStreamWithoutModel is NOT modeled in TestLibraryModels
+  private CustomStreamWithoutModel<Integer> filterThenMapLambdasCustomStream(CustomStreamWithoutModel<String> stream) {
+    // Safe because generic is String, not @Nullable String
+    return stream.filter(s -> s != null).map(s -> s.length());
+  }
+
+  private CustomStreamWithoutModel<Integer> filterThenMapNullableContainerLambdasCustomStream(
+          CustomStreamWithoutModel<NullableContainer<String>> stream) {
+    return stream
+            .filter(c -> c.get() != null)
+            // BUG: Diagnostic contains: dereferenced expression
+            .map(c -> c.get().length());
+  }
+
+  private CustomStreamWithoutModel<Integer> filterThenMapMethodRefsCustomStream(
+          CustomStreamWithoutModel<NullableContainer<String>> stream) {
+    return stream
+            .filter(c -> c.get() != null && perhaps())
+            .map(NullableContainer::get) // CSWoM<NullableContainer<String>> -> CSWoM<@Nullable String>
+            .map(String::length); // Should be an error with proper generics support!
+  }
+
   private static class CheckNonfinalBeforeStream<T> {
     @Nullable private 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();
+  }
+}

nullaway/src/test/resources/com/uber/nullaway/testdata/unannotated/CustomStreamWithoutModel.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 copy of {@link CustomStream} but for which our test library models have no spec/model, to test that errors still
+ * get reported in the absence of a model.
+ */
+public class CustomStreamWithoutModel<T> {
+
+    private CustomStreamWithoutModel() {}
+
+    public CustomStreamWithoutModel<T> filter(Predicate<? super T> predicate) {
+        throw new UnsupportedOperationException();
+    }
+
+    public <R> CustomStreamWithoutModel<R> map(Function<? super T, ? extends R> mapper) {
+        throw new UnsupportedOperationException();
+    }
+}