MergeIterables.java

public class MergeIterables {
  private MergeIterables() {}

  public static <T extends Comparable<T>> Iterable<T> mergeIterables(final List<? extends Iterable<? extends T>> iterables) {
    return mergeIterables(iterables, Ordering.<T>natural());
  }

  public static <T> Iterable<T> mergeIterables(final List<? extends Iterable<? extends T>> iterables, final Comparator<T> comparator) {
    if (iterables.isEmpty()) {
      return Collections.emptyList();
    }
    return new Iterable<T>(){
      @Override
      public Iterator<T> iterator() {
        return mergeIterator(iterables, comparator);
      }
    };
  }

  private static <T> Iterator<T> mergeIterator(final List<? extends Iterable<? extends T>> iterables, final Comparator<T> comparator) {
    final PriorityQueue<IndexedItem<T>> workspace = new PriorityQueue<IndexedItem<T>>(iterables.size());
    final List<Iterator<? extends T>> iterators = Lists.newArrayListWithCapacity(iterables.size());
    for (int i = 0; i < iterables.size(); i++) {
      final Iterator<? extends T> iterator = iterables.get(i).iterator();
      iterators.add(iterator);
      workspace.add(getNextItem(iterator, i, comparator));
    }

    return new AbstractIterator<T>() {
      @Override
      protected T computeNext() {
        IndexedItem<T> next = workspace.poll();
        if (next == null || next.isEndOfData()) {
          return endOfData();
        }
        final T resultItem = next.item;

        final Iterator<? extends T> iterator = iterators.get(next.index);
        if (iterator.hasNext()) {
          next.item = iterator.next();
        } else {
          next = IndexedItem.endOfData(next.index);
        }
        workspace.add(next);
        return resultItem;
      }
    };
  }

  private static <T> IndexedItem<T> getNextItem(Iterator<? extends T> iterator, int iteratorIndex, final Comparator<T> comparator) {
    if (iterator.hasNext()) {
      return new IndexedItem<T>(iteratorIndex, iterator.next(), comparator);
    } else {
      return IndexedItem.endOfData(iteratorIndex);
    }
  }


private static class IndexedItem<T> implements Comparable<IndexedItem<T>> {
    private final int index;
    private T item;
    private final Comparator<T> comparator;

    private IndexedItem(int index, T item, Comparator<T> comparator) {
      this.index = index;
      this.item = item;
      this.comparator = comparator;
    }

    @Override
    public int compareTo(IndexedItem<T> otherItem) {
      if (otherItem.isEndOfData()) {
        return -1;
      }
      return comparator.compare(item, otherItem.item);
    }

    public boolean isEndOfData() {
      return false;
    }

    public static <T> IndexedItem<T> endOfData(int index) {
      return new EndOfData<T>(index);
    }

    /**
      End of data is a sentinel that's always at the end of the index
    **/
    private static class EndOfData<J> extends IndexedItem<J> {
      private EndOfData(int index) {
        super(index, null, null);
      }

      @Override
      public int compareTo(IndexedItem<J> otherItem) {
        return 1;
      }

      @Override
      public boolean isEndOfData() {
        return true;
      }
    }
  }
}