Java Stream forEach 如何实现 break 操作

1. 概述

作为Java开发者，我们经常会有需要遍历数组的场景。利用 Java 8 stream库，以及 forEach 方法能很优雅完成这一任务。

虽然和for循环类似，但其缺少等效的 break 语句来中断遍历。

本文我们将学习如何在 Stream.forEach 中模拟 break 操作。

2. Java 9 中的 Stream.takeWhile()

假设我们有一个 String Stream，我们希望只处理长度为奇数的字符串，当遇到长度为偶数的元素时停止后续处理。

下面代码使用 Java 9 中 Stream.takeWhile 方法来实现这一需求：

Stream.of("cat", "dog", "elephant", "fox", "rabbit", "duck")
  .takeWhile(n -> n.length() % 2 != 0)
  .forEach(System.out::println);

运行后输出结果：

cat
dog

看明白了吗？为了方便理解，我们使用传统for循环 + break 来重新实现下等效代码：

List<String> list = asList("cat", "dog", "elephant", "fox", "rabbit", "duck");
for (int i = 0; i < list.size(); i++) {
    String item = list.get(i);
    if (item.length() % 2 == 0) {
        break;
    }
    System.out.println(item);
}

正如我们所见，takeWhile 方法使我们能够准确地实现我们所需要的。

现在问题来了，如果不用 Java 9，使用Java 8 如何实现类似功能？

3. 自定义 Spliterator

让我们创建一个自定义的Spliterator，它将作为Stream.spliterator的装饰器工作。我们可以使这个Spliterator为我们执行break操作。

首先，我们将从流中获取Spliterator，然后使用我们的CustomSpliterator进行装饰，并提供一个Predicate来控制break操作。最后，我们将使用CustomSpliterator创建一个新的流：

public static <T> Stream<T> takeWhile(Stream<T> stream, Predicate<T> predicate) {
    CustomSpliterator<T> customSpliterator = new CustomSpliterator<>(stream.spliterator(), predicate);
    return StreamSupport.stream(customSpliterator, false);
}

下面我们看如何创建 CustomSpliterator:

public class CustomSpliterator<T> extends Spliterators.AbstractSpliterator<T> {

    private Spliterator<T> splitr;
    private Predicate<T> predicate;
    private boolean isMatched = true;

    public CustomSpliterator(Spliterator<T> splitr, Predicate<T> predicate) {
        super(splitr.estimateSize(), 0);
        this.splitr = splitr;
        this.predicate = predicate;
    }

    @Override
    public synchronized boolean tryAdvance(Consumer<? super T> consumer) {
        boolean hadNext = splitr.tryAdvance(elem -> {
            if (predicate.test(elem) && isMatched) {
                consumer.accept(elem);
            } else {
                isMatched = false;
            }
        });
        return hadNext && isMatched;
    }
}

让我们来看一下tryAdvance方法。在这里我们可以看到，自定义的Spliterator处理装饰的Spliterator的元素。只要我们的断言条件匹配且原始stream还有元素，就会进行处理。当其中任何一个条件变为false时，我们的 Spliterator 会 “break”，Stream 操作结束。

下面是测试代码：

@Test
public void whenCustomTakeWhileIsCalled_ThenCorrectItemsAreReturned() {
    Stream<String> initialStream = 
      Stream.of("cat", "dog", "elephant", "fox", "rabbit", "duck");

    List<String> result = 
      CustomTakeWhile.takeWhile(initialStream, x -> x.length() % 2 != 0)
        .collect(Collectors.toList());

    assertEquals(asList("cat", "dog"), result);
}

如我们所见，stream在满足条件后停止。 出于测试目的，我们已将结果collect到一个list中，但我们也可以使用 forEach 调用或 Stream 的任何其他函数。

4. 自定义 forEach

虽然提供嵌入了 break 机制的 Stream 可能很有用，但只关注 forEach 操作可能更简单

让我们在没有装饰器的情况下直接使用 Stream.spliterator：

public class CustomForEach {

    public static class Breaker {
        private boolean shouldBreak = false;

        public void stop() {
            shouldBreak = true;
        }

        boolean get() {
            return shouldBreak;
        }
    }

    public static <T> void forEach(Stream<T> stream, BiConsumer<T, Breaker> consumer) {
        Spliterator<T> spliterator = stream.spliterator();
        boolean hadNext = true;
        Breaker breaker = new Breaker();

        while (hadNext && !breaker.get()) {
            hadNext = spliterator.tryAdvance(elem -> {
                consumer.accept(elem, breaker);
            });
        }
    }
}

正如我们所见，新的自定义 forEach 方法调用 BiConsumer 为我们的代码提供下一个元素和可用于停止stream的breaker对象。

进行单元测试：

@Test
public void whenCustomForEachIsCalled_ThenCorrectItemsAreReturned() {
    Stream<String> initialStream = Stream.of("cat", "dog", "elephant", "fox", "rabbit", "duck");
    List<String> result = new ArrayList<>();

    CustomForEach.forEach(initialStream, (elem, breaker) -> {
        if (elem.length() % 2 == 0) {
            breaker.stop();
        } else {
            result.add(elem);
        }
    });

    assertEquals(asList("cat", "dog"), result);
}

5. 总结

本文我们学习了如何在 stream 上进行 "break" 操作。我们介绍了了 Java9 中的 takeWhile 方法，以及在 Java 8 中如何自己实现这一功能。

最后本文中的代码可以在 GitHub 上找到。