Generally, custom operations will need to deal with the Spliterator
interface. It extends the concept of the Iterator
by adding characteristics and size information and the ability to split off a part of the elements as another spliterator (hence its name). It also simplifies the iteration logic by only needing one method.
public static <T> Stream<T> takeWhile(Stream<T> s, Predicate<? super T> condition) {
boolean parallel = s.isParallel();
Spliterator<T> spliterator = s.spliterator();
return StreamSupport.stream(new Spliterators.AbstractSpliterator<T>(
spliterator.estimateSize(),
spliterator.characteristics()&~(Spliterator.SIZED|Spliterator.SUBSIZED)) {
boolean active = true;
Consumer<? super T> current;
Consumer<T> adapter = t -> {
if((active = condition.test(t))) current.accept(t);
};
@Override
public boolean tryAdvance(Consumer<? super T> action) {
if(!active) return false;
current = action;
try {
return spliterator.tryAdvance(adapter) && active;
}
finally {
current = null;
}
}
}, parallel).onClose(s::close);
}
To keep the stream’s properties, we query the parallel status first, to reestablish it for the new stream. Also, we register a close action that will close the original stream.
The main work is to implement a Spliterator
decorating the previous stream state’s spliterator.
The characteristics are kept, except for the SIZED
and SUBSIZED
, as our operation results in an unpredictable size. The original size is still passed through, it will now be used as an estimate.
This solution stores the Consumer
passed to tryAdvance
for the duration of the operation, to be able to use the same adapter consumer, avoiding to create a new one for each iteration. This works, as it is guaranteed that tryAdvance
is never invoked concurrently.
Parallelism is done via splitting, which is inherited from AbstractSpliterator
. This inherited implementation will buffer some elements, which is reasonable, as implementing a better strategy for an operation like takeWhile
is really complicated.
So you can use it like
takeWhile(Stream.of("foo", "bar", "baz", "hello", "world"), s -> s.length() == 3)
.forEach(System.out::println);
which will print
foo
bar
baz
or
takeWhile(Stream.of("foo", "bar", "baz", "hello", "world")
.peek(s -> System.out.println("before takeWhile: "+s)), s -> s.length() == 3)
.peek(s -> System.out.println("after takeWhile: "+s))
.forEach(System.out::println);
which will print
before takeWhile: foo
after takeWhile: foo
foo
before takeWhile: bar
after takeWhile: bar
bar
before takeWhile: baz
after takeWhile: baz
baz
before takeWhile: hello
which shows that it does not process more than necessary. Before the takeWhile
stage, we have to encounter the first non-matching element, after that, we only encounter the elements up to that.