package com.fasterxml.sort;

import com.fasterxml.sort.util.SegmentedBuffer;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Comparator;
import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;

public class IteratingSorter<T> extends SorterBase<T>
{
    // Set iff sort spilled to disk
    private List<File> _mergerInputs;
    private DataReader<T> _merger;


    protected IteratingSorter(SortConfig config,
                              DataReaderFactory<T> readerFactory,
                              DataWriterFactory<T> writerFactory,
                              Comparator<T> comparator)
    {
        super(config, readerFactory, writerFactory, comparator);
    }

    protected IteratingSorter() {
        super();
    }

    protected IteratingSorter(SortConfig config) {
        super(config);
    }

    /**
     * Method that will perform full sort on input data read using given
     * {@link DataReader}.
     *
     * Conversions to and from intermediate sort files is done
     * using {@link DataReaderFactory} and {@link DataWriterFactory} configured
     * for this sorter.
     *
     * @return Iterator if sorting complete and output is ready to be written; null if it was cancelled
     */
    public Iterator<T> sort(DataReader<T> inputReader)
        throws IOException
    {
        // Clean up any previous sort
        close();

        // First, pre-sort:
        _phase = Phase.PRE_SORTING;
        boolean inputClosed = false;

        SegmentedBuffer buffer = new SegmentedBuffer();
        _presortFileCount = 0;
        _sortRoundCount = -1;
        _currentSortRound = -1;

        Iterator<T> iterator = null;
        try {
            Object[] items = _readMax(inputReader, buffer, _config.getMaxMemoryUsage(), null);
            if (_checkForCancel()) {
                close();
                return null;
            }
            Arrays.sort(items, _rawComparator());
            T next = inputReader.readNext();
            /* Minor optimization: in case all entries might fit in
             * in-memory sort buffer, avoid writing intermediate file
             * and just write results directly.
             */
            if (next == null) {
                inputClosed = true;
                inputReader.close();
                _phase = Phase.SORTING;
                iterator = new CastingIterator<T>(Arrays.asList(items).iterator());
            } else { // but if more data than memory-buffer-full, do it right:
                List<File> presorted = new ArrayList<File>();
                presorted.add(_writePresorted(items));
                items = null; // it's a big array, clear refs as early as possible
                _presort(inputReader, buffer, next, presorted);
                inputClosed = true;
                inputReader.close();
                _phase = Phase.SORTING;
                if (_checkForCancel(presorted)) {
                    close();
                    return null;
                }
                _mergerInputs = presorted;
                _merger = merge(presorted);
                iterator = new MergerIterator<T>(_merger);
            }
        } finally {
            if (!inputClosed) {
                try {
                    inputReader.close();
                } catch (IOException e) {
                    // Ignore
                }
            }
        }
        if (_checkForCancel()) {
            close();
            return null;
        }
        _phase = Phase.COMPLETE;
        return iterator;
    }

    public void close() {
        if (_merger != null) {
            try {
                _merger.close();
            }
            catch (IOException e) {
                // Ignore
            }
        }
        if (_mergerInputs != null) {
            for (File input : _mergerInputs) {
                input.delete();
            }
        }
        _mergerInputs = null;
        _merger = null;
    }


    /*
    /**********************************************************************
    /* Exception API
    /**********************************************************************
    */

    public static class IterableSorterException extends RuntimeException {
        public IterableSorterException(IOException cause) {
            super(cause);
        }
    }


    /*
    /**********************************************************************
    /* Iterator API
    /**********************************************************************
    */

    private static class CastingIterator<T> implements Iterator<T> {
        private final Iterator<Object> _it;

        public CastingIterator(Iterator<Object> it) {
            _it = it;
        }

        @Override
        public boolean hasNext() {
            return _it.hasNext();
        }

        @SuppressWarnings("unchecked")
        @Override
        public T next() {
            return (T)_it.next();
        }

        @Override
        public void remove() {
            throw new UnsupportedOperationException();
        }
    }

    private static class MergerIterator<T> implements Iterator<T> {
        private final DataReader<T> _merger;
        private T _next;

        private MergerIterator(DataReader<T> merger) {
            _merger = merger;
        }

        private void prepNext() {
            if (_next != null) {
                try {
                    _next = _merger.readNext();
                } catch (IOException e) {
                    throw new IterableSorterException(e);
                }
            }
        }

        @Override
        public boolean hasNext() {
            prepNext();
            return (_next != null);
        }

        @Override
        public T next() {
            prepNext();
            if (_next == null) {
                throw new NoSuchElementException();
            }
            T t = _next;
            _next = null;
            return t;
        }

        @Override
        public void remove() {
            throw new UnsupportedOperationException();
        }
    }
}

@@ -71,16 +71,42 @@
 uffer;%0A%0A
+import java.io.Closeable;%0A
 import j
@@ -373,16 +373,37 @@
 rBase%3CT%3E
+ implements Closeable
 %0A%7B%0A    /
@@ -3661,24 +3661,227 @@
 tor;%0A    %7D%0A%0A
+%0A    /*%0A    /**********************************************************************%0A    /* Closeable API%0A    /**********************************************************************%0A    */%0A%0A    @Override%0A
     public v

Implement Closeable