Index.java

package com.codepoetics.octarine.joins;

import com.codepoetics.octarine.functional.tuples.T2;

import java.util.*;
import java.util.function.BiFunction;
import java.util.function.Function;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

public final class Index<K, L> {
    private final Comparator<? super K> comparator;
    private final SortedMap<K, Set<L>> indexed;

    private Index(Comparator<? super K> comparator, SortedMap<K, Set<L>> indexed) {
        this.comparator = comparator;
        this.indexed = indexed;
    }

    public static <K extends Comparable<K>, L> Index<K, L> on(Stream<? extends L> stream, Function<? super L, ? extends K> key) {
        return on(stream, key, Comparator.naturalOrder());
    }

    public static <K, L> Index<K, L> on(Stream<? extends L> stream, Function<? super L, ? extends K> key, Comparator<? super K> comparator) {
        return new Index<K, L>(comparator, stream.collect(IndexCollector.<K, L>on(key, comparator)));
    }

    public Spliterator<Map.Entry<K, Set<L>>> entrySpliterator() {
        return indexed.entrySet().spliterator();
    }

    public Set<K> keys() {
        return indexed.keySet();
    }

    public <R> Stream<T2<L, Set<R>>> oneToMany(Index<K, R> other) {
        return matchAndMerge(other, oneToManyMerger());
    }

    public <R> Stream<T2<L, Set<R>>> strictOneToMany(Index<K, R> other) {
        return matchAndMerge(other, strictOneToManyMerger());
    }

    private <R, T> Stream<T> matchAndMerge(Index<K, R> other, BiFunction<Set<L>, Set<R>, Stream<T>> merger) {
        return matchedSublists(other).flatMap(t -> t.pack(merger));
    }

    private <R> BiFunction<Set<L>, Set<R>, Stream<T2<L, Set<R>>>> oneToManyMerger() {
        return (lefts, rights) ->
                lefts.stream().map(left -> T2.of(left, rights));
    }

    private <R> BiFunction<Set<L>, Set<R>, Stream<T2<L, Set<R>>>> strictOneToManyMerger() {
        return (lefts, rights) -> {
            if (lefts.isEmpty()) {
                throw new IllegalArgumentException("Unmatched right values found");
            }
            if (lefts.size() > 1) {
                throw new IllegalArgumentException("Duplicate left values found");
            }
            return Stream.of(T2.of(lefts.iterator().next(), rights));
        };
    }

    public <R> Stream<T2<L, R>> manyToOne(Index<K, R> other) {
        return innerJoin(other);
    }

    public <R> Stream<T2<L, R>> strictManyToOne(Index<K, R> other) {
        return matchAndMerge(other, strictManyToOneMerger());
    }

    private <R> BiFunction<Set<L>, Set<R>, Stream<T2<L, R>>> strictManyToOneMerger() {
        return (lefts, rights) -> {
            if (rights.isEmpty()) {
                throw new IllegalArgumentException("Unmatched left values found");
            }
            if (rights.size() > 1) {
                throw new IllegalArgumentException("Duplicate right values found");
            }
            R right = rights.iterator().next();
            return lefts.stream().map(left -> T2.of(left, right));
        };
    }

    public <R> Stream<T2<L, R>> strictOneToOne(Index<K, R> other) {
        return matchAndMerge(other, strictOneToOneMerger());
    }

    private <R> BiFunction<Set<L>, Set<R>, Stream<T2<L, R>>> strictOneToOneMerger() {
        return (lefts, rights) -> {
            if (lefts.isEmpty()) {
                throw new IllegalArgumentException("Unmatched right values found");
            }
            if (lefts.size() > 1) {
                throw new IllegalArgumentException("Duplicate left values found");
            }

            if (rights.isEmpty()) {
                throw new IllegalArgumentException("Unmatched left values found");
            }
            if (rights.size() > 1) {
                throw new IllegalArgumentException("Duplicate right values found");
            }

            return Stream.of(T2.of(lefts.iterator().next(), rights.iterator().next()));
        };
    }

    public <R> Stream<T2<L, Optional<R>>> leftOuterJoin(Index<K, R> other) {
        return matchAndMerge(other, leftOuterJoinMerger());
    }

    private <R> BiFunction<Set<L>, Set<R>, Stream<T2<L, Optional<R>>>> leftOuterJoinMerger() {
        return (lefts, rights) -> {
            if (rights.isEmpty()) {
                return lefts.stream().map(l -> T2.of(l, Optional.<R>empty()));
            }

            return lefts.stream().flatMap(l -> rights.stream().map(r -> T2.of(l, Optional.<R>of(r))));
        };
    }

    public <R> Stream<T2<Optional<L>, R>> rightOuterJoin(Index<K, R> other) {
        return matchAndMerge(other, rightOuterJoinMerger());
    }

    private <R> BiFunction<Set<L>, Set<R>, Stream<T2<Optional<L>, R>>> rightOuterJoinMerger() {
        return (lefts, rights) -> {

            if (lefts.isEmpty()) {
                return rights.stream().map(r -> T2.of(Optional.<L>empty(), r));
            }

            return rights.stream().flatMap(r -> lefts.stream().map(l -> T2.of(Optional.of(l), r)));
        };
    }

    public <R> Stream<T2<L, R>> innerJoin(Index<K, R> other) {
        return matchAndMerge(other, innerJoinMerger());
    }

    private <R> BiFunction<Set<L>, Set<R>, Stream<T2<L, R>>> innerJoinMerger() {
        return (lefts, rights) -> lefts.stream().flatMap(l ->
            rights.stream().map(r -> T2.of(l, r)));
    }

    public <R> Stream<T2<Optional<L>, Optional<R>>> outerJoin(Index<K, R> other) {
        return matchAndMerge(other, outerJoinMerger());
    }

    private <R> BiFunction<Set<L>, Set<R>, Stream<T2<Optional<L>, Optional<R>>>> outerJoinMerger() {
        return (lefts, rights) -> {

            if (lefts.isEmpty()) {
                return rights.stream().map(r -> T2.of(Optional.<L>empty(), Optional.of(r)));
            }

            if (rights.isEmpty()) {
                return lefts.stream().map(l -> T2.of(Optional.of(l), Optional.<R>empty()));
            }

            return lefts.stream().flatMap(l ->
                            rights.stream().map(r ->
                                            T2.of(Optional.<L>of(l), Optional.<R>of(r))
                            )
            );
        };
    }

    private <R> Stream<T2<Set<L>, Set<R>>> matchedSublists(Index<K, R> other) {
        return StreamSupport.stream(KeyMatchingSpliterator.over(
                comparator,
                entrySpliterator(),
                other.entrySpliterator(),
                Collections.emptySet(),
                Collections.emptySet()), false);
    }

    @Override
    public String toString() {
        return indexed.toString();
    }

}