How to implement a parallel supporting takeWhile for the Stream API in Java 8?

Question

I've seen some takeWhile implementations for the Java 8 stream API but they all seem to turn the stream into a non-parallel stream. For example this one:

static <T> Spliterator<T> takeWhile(
    Spliterator<T> splitr, Predicate<? super T> predicate) {
  return new Spliterators.AbstractSpliterator<T>(splitr.estimateSize(), 0) {
    boolean stillGoing = true;
    @Override public boolean tryAdvance(Consumer<? super T> consumer) {
      if (stillGoing) {
        boolean hadNext = splitr.tryAdvance(elem -> {
          if (predicate.test(elem)) {
            consumer.accept(elem);
          } else {
            stillGoing = false;
          }
        });
        return hadNext && stillGoing;
      }
      return false;
    }
  };
}

static <T> Stream<T> takeWhile(Stream<T> stream, Predicate<? super T> predicate) {
   return StreamSupport.stream(takeWhile(stream.spliterator(), predicate), false);
}

Here StreamSupport.stream(takeWhile(stream.spliterator(), predicate), false); turns the stream passed to takeWhile into a sequential stream. Is anyone aware of an implementation that supports parallel streams or how can I modify this code to make it maintain/support parallel streams?

Answer 1

If your source is known to be unordered, then the following implementation should work:

static final class UnorderedTakeWhileSpliterator<T> implements Spliterator<T>, Consumer<T>, Cloneable {
    private final Predicate<? super T> predicate;
    private final AtomicBoolean checked = new AtomicBoolean();
    private Spliterator<T> source;
    private T cur;

    UnorderedTakeWhileSpliterator(Spliterator<T> source, Predicate<? super T> predicate) {
        this.predicate = predicate;
        this.source = source;
    }

    @Override
    public void accept(T t) {
        this.cur = t;
    }

    @Override
    public boolean tryAdvance(Consumer<? super T> action) {
        if (!checked.get() && source.tryAdvance(this)) {
            if (predicate.test(cur)) {
                action.accept(cur);
                return true;
            } else {
                checked.set(true);
            }
        }
        return false;
    }

    @Override
    public Spliterator<T> trySplit() {
        Spliterator<T> prefix = source.trySplit();
        if(prefix == null) {
            return null;
        }
        if(checked.get()) {
            return Spliterators.emptySpliterator();
        }
        UnorderedTakeWhileSpliterator<T> clone;
        try {
            clone = (UnorderedTakeWhileSpliterator<T>) clone();
        } catch (CloneNotSupportedException e) {
            throw new InternalError(e);
        }
        clone.source = prefix;
        return clone;
    }

    @Override
    public long estimateSize() {
        return source.estimateSize();
    }

    @Override
    public int characteristics() {
        return source.characteristics() & (DISTINCT | SORTED | NONNULL);
    }

    @Override
    public Comparator<? super T> getComparator() {
        return source.getComparator();
    }
}

Create the stream with the following method:

static <T> Stream<T> takeWhile(Stream<T> stream, Predicate<? super T> predicate) {
   return StreamSupport.stream(UnorderedTakeWhileSpliterator<>(stream.spliterator(), predicate), stream.isParallel());
}

Ordered implementation would be much more tricky as it should buffer non-prefixed items and propagate the cancelling to the suffixes. Something like this is implemented in JDK-9 (not as spliterator, but as normal stream operation), but I doubt that even this tricky implementation wins in many cases over sequential stream.