(w/lowasser) Streams.forEachPair()

[]

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=151386365

Author: kevinb

guava-tests/test/com/google/common/collect/StreamsTest.java

@@ -31,6 +31,7 @@
 import java.util.OptionalDouble;
 import java.util.OptionalInt;
 import java.util.OptionalLong;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.DoubleStream;
@@ -264,6 +265,54 @@ public void testZipDifferingLengths() {
         .inOrder();
   }
 
+  public void testForEachPair() {
+    List<String> list = new ArrayList<>();
+    Streams.forEachPair(
+        Stream.of("a", "b", "c"), Stream.of(1, 2, 3), (a, b) -> list.add(a + ":" + b));
+    Truth.assertThat(list).containsExactly("a:1", "b:2", "c:3");
+  }
+
+  public void testForEachPair_differingLengths1() {
+    List<String> list = new ArrayList<>();
+    Streams.forEachPair(
+        Stream.of("a", "b", "c", "d"), Stream.of(1, 2, 3), (a, b) -> list.add(a + ":" + b));
+    Truth.assertThat(list).containsExactly("a:1", "b:2", "c:3");
+  }
+
+  public void testForEachPair_differingLengths2() {
+    List<String> list = new ArrayList<>();
+    Streams.forEachPair(
+        Stream.of("a", "b", "c"), Stream.of(1, 2, 3, 4), (a, b) -> list.add(a + ":" + b));
+    Truth.assertThat(list).containsExactly("a:1", "b:2", "c:3");
+  }
+
+  public void testForEachPair_oneEmpty() {
+    Streams.forEachPair(Stream.of("a"), Stream.empty(), (a, b) -> fail());
+  }
+
+  public void testForEachPair_finiteWithInfinite() {
+    List<String> list = new ArrayList<>();
+    Streams.forEachPair(
+        Stream.of("a", "b", "c"), Stream.iterate(1, i -> i + 1), (a, b) -> list.add(a + ":" + b));
+    Truth.assertThat(list).containsExactly("a:1", "b:2", "c:3");
+  }
+
+  public void testForEachPair_parallel() {
+    Stream<String> streamA = IntStream.range(0, 100000).mapToObj(String::valueOf).parallel();
+    Stream<Integer> streamB = IntStream.range(0, 100000).mapToObj(i -> i).parallel();
+
+    AtomicInteger count = new AtomicInteger(0);
+    Streams.forEachPair(
+        streamA,
+        streamB,
+        (a, b) -> {
+          count.incrementAndGet();
+          Truth.assertThat(a.equals(String.valueOf(b))).isTrue();
+        });
+    Truth.assertThat(count.get()).isEqualTo(100000);
+    // of course, this test doesn't prove that anything actually happened in parallel...
+  }
+
   // TODO(kevinb): switch to importing Truth's assertThat(Stream) if we get that added
   private static IterableSubject assertThat(Stream<?> stream) {
     return Truth.assertThat(stream.toArray()).asList();

guava/src/com/google/common/collect/Streams.java

@@ -33,6 +33,7 @@
 import java.util.Spliterator;
 import java.util.Spliterators;
 import java.util.Spliterators.AbstractSpliterator;
+import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
 import java.util.function.DoubleConsumer;
@@ -205,7 +206,7 @@ public static DoubleStream concat(DoubleStream... streams) {
   }
 
   /**
-   * Returns a stream in which each element is the result of passing the corresponding element of
+   * Returns a stream in which each element is the result of passing the corresponding elementY of
    * each of {@code streamA} and {@code streamB} to {@code function}.
    *
    * <p>For example:
@@ -222,7 +223,10 @@ public static DoubleStream concat(DoubleStream... streams) {
    * <p>The resulting stream will only be as long as the shorter of the two input streams; if one
    * stream is longer, its extra elements will be ignored.
    *
-   * <p>The resulting stream is not <a
+   * <p>Note that if you are calling {@link Stream#forEach} on the resulting stream, you might want
+   * to consider using {@link #forEachPair} instead of this method.
+   *
+   * <p><b>Performance note:</b> The resulting stream is not <a
    * href="http://gee.cs.oswego.edu/dl/html/StreamParallelGuidance.html">efficiently splittable</a>.
    * This may harm parallel performance.
    */
@@ -255,6 +259,59 @@ public boolean tryAdvance(Consumer<? super R> action) {
         isParallel);
   }
 
+  /**
+   * Invokes {@code consumer} once for each pair of <i>corresponding</i> elements in {@code streamA}
+   * and {@code streamB}. If one stream is longer than the other, the extra elements are silently
+   * ignored. Elements passed to the consumer are guaranteed to come from the same position in their
+   * respective source streams. For example:
+   *
+   * <pre>{@code
+   * Streams.forEachPair(
+   *   Stream.of("foo1", "foo2", "foo3"),
+   *   Stream.of("bar1", "bar2"),
+   *   (arg1, arg2) -> System.out.println(arg1 + ":" + arg2)
+   * }</pre>
+   *
+   * <p>will print:
+   *
+   * <pre>{@code
+   * foo1:bar1
+   * foo2:bar2
+   * }</pre>
+   *
+   * <p><b>Warning:</b> If either supplied stream is a parallel stream, the same correspondence
+   * between elements will be made, but the order in which those pairs of elements are passed to the
+   * consumer is <i>not</i> defined.
+   *
+   * <p>Note that many usages of this method can be replaced with simpler calls to {@link #zip}.
+   * This method behaves equivalently to {@linkplain #zip zipping} the stream elements into
+   * temporary pair objects and then using {@link Stream#forEach} on that stream.
+   */
+  public static <A, B> void forEachPair(
+      Stream<A> streamA, Stream<B> streamB, BiConsumer<? super A, ? super B> consumer) {
+    checkNotNull(consumer);
+
+    if (streamA.isParallel() || streamB.isParallel()) {
+      zip(streamA, streamB, TemporaryPair::new).forEach(pair -> consumer.accept(pair.a, pair.b));
+    } else {
+      Iterator<A> iterA = streamA.iterator();
+      Iterator<B> iterB = streamB.iterator();
+      while (iterA.hasNext() && iterB.hasNext()) {
+        consumer.accept(iterA.next(), iterB.next());
+      }
+    }
+  }
+
+  private static class TemporaryPair<A, B> {
+    final A a;
+    final B b;
+
+    TemporaryPair(A a, B b) {
+      this.a = a;
+      this.b = b;
+    }
+  }
+
   /**
    * Returns a stream consisting of the results of applying the given function to the elements of
    * {@code stream} and their indices in the stream. For example,