Java Spliterator for consecutive pairs

PairSpliterator.java

import java.util.Spliterator;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.BiFunction;

/**
 * Class that supports splittable iteration over consecutive pairs of an
 * underlying Spliterator.
 *
 * A simpler implementation of this functionality would be to create an
 * Iterator of pairs from the given Spliterator, and then convert this iterator
 * to a Spliterator using the Spliterators class. However, this implementation
 * would fail to take advantage of the splitting capabilities of the original
 * Spliterator.
 *
 * Instead, this class wraps the original Spliterator and delegates splitting
 * operations to it. We have to do some extra bookkeeping at the split points
 * to ensure that pairs that span two Spliterators are correctly taken care of.
 */
public class PairSpliterator<T, P> implements Spliterator<P> {
    private final Spliterator<? extends T> spliterator;
    private final BiFunction<? super T, ? super T, ? extends P> pairFunction;

    private Link<T> pred;
    private Link<T> succ;
    private boolean topLevel;

    private T prev;
    private boolean started;
    private boolean exhausted;

    /**
     * @param spliterator The underlying spliterator
     * @param pairFunction Mapping function pairs of consecutive elements
     */
    public PairSpliterator(
            Spliterator<? extends T> spliterator,
            BiFunction<? super T, ? super T, ? extends P> pairFunction) {
        this.spliterator = spliterator;
        this.pairFunction = pairFunction;
        this.pred = new Link<>();
        this.succ = new Link<>();
        this.topLevel = true;
    }

    private PairSpliterator(PairSpliterator<T, P> other,
                            Spliterator<? extends T> spliterator) {
        this.spliterator = spliterator;
        this.pairFunction = other.pairFunction;
        this.pred = other.pred;
        this.succ = other.succ;
        this.topLevel = other.topLevel;
        this.prev = other.prev;
        this.started = other.started;
        this.exhausted = other.exhausted;
    }

    @Override
    public int characteristics() {
        // This implementation preserves the ORDERED, IMMUTABLE, and CONCURRENT
        // characteristics of the underlying Spliterator. Before splitting, it
        // also preserves the SIZED characteristic.
        int mask = Spliterator.ORDERED
            | Spliterator.IMMUTABLE
            | Spliterator.CONCURRENT;
        if (topLevel) {
            mask |= Spliterator.SIZED;
        }
        return spliterator.characteristics() & mask;
    }

    @Override
    public long estimateSize() {
        // If the underlying Spliterator produces N elements, this Spliterator
        // will produce between N-1 and N+1 elements, depending on the boundary
        // conditions.
        long estimate = spliterator.estimateSize();
        if (estimate == 0L || estimate == Long.MAX_VALUE) {
            return estimate;
        }
        return estimate - 1;
    }

    @Override
    public PairSpliterator<T, P> trySplit() {
        if (exhausted) {
            return null;
        }

        final Spliterator<? extends T> prefixSpliterator =
            spliterator.trySplit();
        if (prefixSpliterator == null) {
            return null;
        }

        // Create a new link for the boundary between the two Spliterators, and
        // update fields accordingly.
        final Link<T> mid = new Link<>();
        final PairSpliterator<T, P> other = new PairSpliterator<>(
            this, prefixSpliterator);
        this.pred = other.succ = mid;
        this.topLevel = other.topLevel = false;
        this.prev = null;
        this.started = false;
        return other;
    }

    @Override
    public boolean tryAdvance(Consumer<? super P> action) {
        final BiConsumer<? super T, ? super T> biAction = biAction(action);

        if (!started && tryAdvanceInitial(biAction)) {
            return true;
        }

        if (exhausted) {
            return false;
        }

        final Ref<T> tRef = new Ref<>();
        if (spliterator.tryAdvance(tRef)) {
            final T next = tRef.value;
            biAction.accept(prev, next);
            prev = next;
            return true;
        } else {
            exhausted = true;
            // Submit last element to successor link
            return succ.submit(prev, true, biAction);
        }
    }

    @Override
    public void forEachRemaining(Consumer<? super P> action) {
        final BiConsumer<? super T, ? super T> biAction = biAction(action);
        if (!started) {
            tryAdvanceInitial(biAction);
        }
        if (exhausted) {
            return;
        }
        spliterator.forEachRemaining(next -> {
            biAction.accept(prev, next);
            prev = next;
        });
        exhausted = true;
        // Submit last element to successor link
        succ.submit(prev, true, biAction);
    }

    /**
     * Try to advance the underlying Spliterator for the first time.
     *
     * @param biAction Action to perform on the initial pair, if available.
     * @return Whether the action was performed.
     */
    private boolean tryAdvanceInitial(
            BiConsumer<? super T, ? super T> biAction) {
        assert !started;
        started = true;

        final Ref<T> tRef = new Ref<>();
        if (spliterator.tryAdvance(tRef)) {
            // Submit first element to predecessor link
            final T next = tRef.value;
            final boolean didAction = pred.submit(
                next, false, biAction);
            prev = next;
            return didAction;
        } else {
            // Underlying Spliterator was empty. Connect the predecessor link
            // to the successor.
            exhausted = true;
            return pred.connect(succ, biAction);
        }
    }

    private BiConsumer<? super T, ? super T> biAction(
            Consumer<? super P> action) {
        return (t1, t2) -> {
            action.accept(pairFunction.apply(t1, t2));
        };
    }

    /**
     * Class for keeping track of the link between two PairSpliterators.
     */
    private static class Link<T> {
        private T value;
        private boolean hasValue;
        private Link<T> succ;

        /**
         * Submit an element to this link.
         *
         * @param value Value to submit
         * @param left Whether this value comes from the left or right side of
         *     the link
         * @param action Action to perform if the link already contains a value
         * @return Whether the action was performed.
         */
        public boolean submit(T value,
                              boolean left,
                              BiConsumer<? super T, ? super T> action) {
            final Ref<T> ref = new Ref<>();
            boolean result = submit(value, left, ref);
            if (result) {
                if (left) {
                    action.accept(value, ref.value);
                } else {
                    action.accept(ref.value, value);
                }
            }
            return result;
        }

        /**
         * @param value Value to 
         * @param left Whether this value comes from the left or right side of
         *     the link
         * @param ref Ref to store existing value, if one exists.
         * @return True if the Link was already populated.
         */
        private synchronized boolean submit(T value,
                                            boolean left,
                                            Ref<T> ref) {
            if (hasValue) {
                ref.value = this.value;
                return true;
            } else if (succ == null) {
                this.value = value;
                hasValue = true;
                return false;
            } else {
                assert left;
                // Propagate value to successor Link
                boolean result = succ.submit(value, true, ref);
                if (result) {
                    this.value = ref.value;
                } else {
                    this.value = value;
                }
                hasValue = true;
                succ = null;
                return result;
            }
        }

        /**
         * Connect this link to its successor.
         *
         * @param other Successor link
         * @param action Action to perform if both links are already populated
         * @return Whether the action was performed.
         */
        public boolean connect(Link<T> other,
                               BiConsumer<? super T, ? super T> action) {
            Ref<T> thisRef = new Ref<>();
            Ref<T> otherRef = new Ref<>();
            boolean result = connect(other, thisRef, otherRef);
            if (result) {
                action.accept(thisRef.value, otherRef.value);
            }
            return result;
        }

        /**
         * @param other Successor link
         * @param thisRef Ref for this link's value
         * @param otherRef Ref for successor value
         * @return True if both links were already populated
         */
        private synchronized boolean connect(Link<T> other,
                                             Ref<T> thisRef,
                                             Ref<T> otherRef) {
            assert succ == null;
            if (hasValue) {
                thisRef.value = this.value;
                // Eagerly propagate value to successor
                return other.submit(value, true, otherRef);
            } else {
                // We don't currently have a value, so store a reference to the
                // successor. If a value is submitted to this link, we will
                // propagate the information to the successor.
                this.succ = other;
                return false;
            }
        }
    }

    /** Mutable object reference. */
    private static final class Ref<T> implements Consumer<T> {
        public T value;

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