From 7cb0b6c31491efab9299e1f5f0ecc1fa22afca7d Mon Sep 17 00:00:00 2001
From: Ben Yu <yujige@gmail.com>
Date: Sat, 19 Oct 2024 19:10:40 -0700
Subject: [PATCH] Make toBiStream() eagerly call the functions. It's what other
 collectors like toMap() do.

---
 .../java/com/google/mu/util/stream/BiStream.java    | 11 ++++++-----
 .../java/com/google/mu/util/stream/MoreStreams.java |  2 +-
 .../com/google/mu/util/stream/BiStreamTest.java     | 13 ++++++++++++-
 3 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/mug/src/main/java/com/google/mu/util/stream/BiStream.java b/mug/src/main/java/com/google/mu/util/stream/BiStream.java
index 4a86867429..a930350224 100644
--- a/mug/src/main/java/com/google/mu/util/stream/BiStream.java
+++ b/mug/src/main/java/com/google/mu/util/stream/BiStream.java
@@ -14,7 +14,7 @@
  *****************************************************************************/
 package com.google.mu.util.stream;
 
-import static com.google.mu.util.stream.MoreStreams.collectingAndThen;
+import static com.google.mu.util.stream.MoreStreams.toStream;
 import static java.util.Map.Entry.comparingByKey;
 import static java.util.Map.Entry.comparingByValue;
 import static java.util.Objects.requireNonNull;
@@ -333,7 +333,7 @@ public static <K, V> Builder<K, V> builder() {
    */
   public static <T, K, V> Collector<T, ?, BiStream<K, V>> concatenating(
       Function<? super T, ? extends BiStream<? extends K, ? extends V>> toBiStream) {
-    return collectingAndThen(stream -> concat(stream.map(toBiStream)));
+    return collectingAndThen(Collectors.mapping(toBiStream, toStream()), BiStream::concat);
   }
 
   /**
@@ -394,7 +394,9 @@ public static <K, V> Builder<K, V> builder() {
       Function<? super E, ? extends K> toKey, Function<? super E, ? extends V> toValue) {
     requireNonNull(toKey);
     requireNonNull(toValue);
-    return collectingAndThen(stream -> from(stream, toKey, toValue));
+    return collectingAndThen(
+        Collectors.mapping(e -> kv(toKey.apply(e), toValue.apply(e)), toStream()),
+        BiStream::fromEntries);
   }
 
   /**
@@ -409,8 +411,7 @@ public static <K, V> Builder<K, V> builder() {
    */
   public static <E, K, V> Collector<E, ?, BiStream<K, V>> toBiStream(
       Function<? super E, ? extends Both<? extends K, ? extends V>> toPair) {
-    requireNonNull(toPair);;
-    return collectingAndThen(stream -> from(stream.map(toPair)));
+    return collectingAndThen(Collectors.mapping(toPair, toStream()), BiStream::from);
   }
 
   /**
diff --git a/mug/src/main/java/com/google/mu/util/stream/MoreStreams.java b/mug/src/main/java/com/google/mu/util/stream/MoreStreams.java
index e39ddddfcf..cc4737a3b0 100644
--- a/mug/src/main/java/com/google/mu/util/stream/MoreStreams.java
+++ b/mug/src/main/java/com/google/mu/util/stream/MoreStreams.java
@@ -525,7 +525,7 @@ static <F, T> Stream<T> mapBySpliterator(
   }
 
   /** Copying input elements into another stream. */
-  private static <T> Collector<T, ?, Stream<T>> toStream() {
+  static <T> Collector<T, ?, Stream<T>> toStream() {
     return Collector.of(
         Stream::<T>builder,
         Stream.Builder::add,
diff --git a/mug/src/test/java/com/google/mu/util/stream/BiStreamTest.java b/mug/src/test/java/com/google/mu/util/stream/BiStreamTest.java
index 041f1cf769..4aa054817c 100644
--- a/mug/src/test/java/com/google/mu/util/stream/BiStreamTest.java
+++ b/mug/src/test/java/com/google/mu/util/stream/BiStreamTest.java
@@ -20,10 +20,11 @@
 import static com.google.mu.util.Optionals.optional;
 import static com.google.mu.util.stream.BiCollectors.toMap;
 import static com.google.mu.util.stream.BiStream.biStream;
-import static com.google.mu.util.stream.BiStream.crossJoining;
 import static com.google.mu.util.stream.BiStream.concatenating;
+import static com.google.mu.util.stream.BiStream.crossJoining;
 import static com.google.mu.util.stream.BiStream.groupingByEach;
 import static com.google.mu.util.stream.BiStream.toAdjacentPairs;
+import static com.google.mu.util.stream.BiStream.toBiStream;
 import static com.google.mu.util.stream.MoreStreams.indexesFrom;
 import static java.util.Arrays.asList;
 import static java.util.function.Function.identity;
@@ -66,6 +67,7 @@
 import com.google.common.truth.IterableSubject;
 import com.google.common.truth.MultimapSubject;
 import com.google.mu.util.BiOptional;
+import com.google.mu.util.Both;
 import com.google.mu.util.Substring;
 
 @RunWith(JUnit4.class)
@@ -744,6 +746,15 @@ public class BiStreamTest {
         .inOrder();
   }
 
+  @Test public void testToBiStream_exceptionThrownEagerly() {
+    assertThrows(
+        NumberFormatException.class,
+        () -> Stream.of("1", "b").collect(toBiStream(identity(), v -> Integer.parseInt(v))));
+    assertThrows(
+        NumberFormatException.class,
+        () -> Stream.of("1", "b").collect(toBiStream(v -> Both.of(v, Integer.parseInt(v)))));
+  }
+
   @Test public void testGroupingBy() {
     Map<Integer, List<Integer>> groups =
         Stream.of(0, 1, 2).collect(BiStream.groupingBy(n -> n / 2)).toMap();