net.morimekta.collect.collectors.InNumBatchesCollector Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of collect Show documentation

Unmodifiable collections for java. These collections have some specific behavior criteria that is currently not guaranteed by the native java unmodifiable collections, they are also set up with convenient builders.

There is a newer version: 5.0.0

Show newest version

package net.morimekta.collect.collectors;

import net.morimekta.collect.UnmodifiableList;

import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BiConsumer;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collector;

import static net.morimekta.collect.UnmodifiableList.toList;
import static net.morimekta.collect.UnmodifiableSet.setOf;

public class InNumBatchesCollector implements Collector>, List>> {
    private final AtomicInteger nextPos = new AtomicInteger();
    private final AtomicInteger numTotal = new AtomicInteger();
    private final int numBatches;

    public InNumBatchesCollector(int numBatches) {
        this.numBatches = numBatches;
    }

    @Override
    public Supplier>> supplier() {
        return () -> {
            ArrayList> batches = new ArrayList<>(numBatches);
            for (int i = 0; i < numBatches; ++i) {
                batches.add(new ArrayList<>());
            }
            return batches;
        };
    }

    @Override
    public BiConsumer>, T> accumulator() {
        return (batches, item) -> {
            int pos = nextPos.getAndUpdate(i -> (++i) % numBatches);
            batches.get(pos).add(item);
            numTotal.incrementAndGet();
        };
    }

    @Override
    public BinaryOperator>> combiner() {
        return (a, b) -> {
            // Merge the two lists so the batches matches the order
            // of the non-parallel inBatchesOf with (a1..an) + (b1..bn)
            // as the set of items. It's not extremely efficient, but
            // works fine as this is not optimized for parallel streams.
            for (int i = 0; i < numBatches; ++i) {
                List al = a.get(i);
                List bl = b.get(i);
                al.addAll(bl);
            }
            return a;
        };
    }

    @Override
    public Function>, List>> finisher() {
        return batches -> {
            if (numTotal.get() < numBatches) {
                batches = batches.subList(0, numTotal.get());
            }
            return batches.stream()
                          .map(UnmodifiableList::asList)
                          .collect(toList());
        };
    }

    @Override
    public Set characteristics() {
        return setOf(Characteristics.UNORDERED);
    }
}