001package com.identityworksllc.iiq.common.iterators;
002
003import sailpoint.tools.CloseableIterator;
004import sailpoint.tools.Util;
005
006import java.util.ArrayList;
007import java.util.Collections;
008import java.util.Iterator;
009import java.util.List;
010import java.util.NoSuchElementException;
011import java.util.Objects;
012
013/**
014 * An iterator that batches up the inputs into batches of size
015 * 'batchSize', then returns them in chunks. The last batch may,
016 * of course, be smaller than the batch size.
017 * 
018 * Similar to the Guava {@link com.google.common.collect.Iterators#partition(Iterator, int)}
019 *
020 * @param <ObjectType>
021 */
022public class BatchingIterator<ObjectType> implements AutoCloseable, CloseableIterator<List<ObjectType>>, Iterator<List<ObjectType>> {
023    /**
024     * The recommended batch size, which will be used for all but the final
025     * batch (which may be smaller, obviously).
026     */
027    private final int batchSize;
028
029    /**
030     * The cached list of objects, which is calculated on hasNext() and then
031     * returned on the next invocation of next().
032     */
033    private List<ObjectType> cachedList;
034
035    /**
036     * The list of inputs
037     */
038    private final Iterator<? extends ObjectType> input;
039
040    /**
041     * Initialized
042     */
043    private volatile boolean initialized;
044
045    /**
046     * Constructor
047     *
048     * @param input The iterator being wrapped by this transformer
049     * @param batchSize The batch size
050     */
051    public BatchingIterator(Iterator<? extends ObjectType> input, int batchSize) {
052        this.input = Objects.requireNonNull(input);
053        this.batchSize = batchSize;
054        this.initialized = false;
055    }
056
057    /**
058     * Closes the iterator by flushing the input iterator
059     */
060    @Override
061    public void close() {
062        Util.flushIterator(input);
063    }
064
065    /**
066     * Retrieves the next batch of items and returns true if the batch is not
067     * empty. Note that if this iterator wraps something slow (e.g., an iterator
068     * that is streaming from disk or something), hasNext() will take a while.
069     *
070     * You MUST call hasNext() before any call to next() will work properly.
071     *
072     * @return True if the batch is not empty and next() will return a value
073     */
074    @Override
075    public boolean hasNext() {
076        List<ObjectType> tempList = new ArrayList<>();
077        int amount = 0;
078        while (input.hasNext() && amount++ < batchSize) {
079            tempList.add(input.next());
080        }
081        cachedList = tempList;
082        this.initialized = true;
083        return cachedList.size() > 0;
084    }
085
086    /**
087     * Returns the next item, which is derived during hasNext(). The result
088     * is always an immutable copy of the internal state of this class.
089     *
090     * @return The next item
091     * @throws NoSuchElementException if hasNext() has not been invoked or if it returned false
092     */
093    @Override
094    public List<ObjectType> next() {
095        if (!this.initialized) {
096            throw new IllegalStateException("You must call hasNext() first");
097        }
098        if (cachedList == null) {
099            throw new NoSuchElementException("Iterator has been exhausted");
100        }
101        List<ObjectType> immutableCopy = Collections.unmodifiableList(new ArrayList<>(cachedList));
102        cachedList = null;
103        return immutableCopy;
104    }
105}