Flattens elements in the Iterable of Iterable of ...

IterableFlattener.java

package test;

import static java.util.Arrays.asList;

import java.util.ArrayDeque;
import java.util.Collections;
import java.util.Deque;
import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.Objects;

public final class IterableFlattener
{
  public static void main(String[] args)
  {
    // For depth based
    testEqual(Collections.emptySet(), flatten(null, 0));
    testEqual(asList("Hello"), flatten(asList("Hello"), 0));
    testEqual(asList("Hello"), flatten(asList(asList("Hello")), 1));
    testEqual(asList("Hello", null, "Hello0"), flatten(asList(asList("Hello", null, "Hello0")), 1));
    testEqual(asList("Hello", "Hello0"), flatten(asList(asList(asList("Hello", "Hello0"))), 2));

    // For type based
    testEqual(Collections.emptySet(), flatten(null, String.class));
    testEqual(asList("Hello"), flatten(asList("Hello"), String.class));
    testEqual(asList("Hello"), flatten(asList(asList("Hello")), String.class));
    testEqual(asList("Hello", /* nulls skipped */"Hello0"), flatten(asList(asList("Hello", null, "Hello0")), String.class));

    System.out.printf("Completed");
  }

  private static void testEqual(Iterable<?> expected, Iterable<?> actual)
  {
    Iterator<?> it1 = expected.iterator(), it2 = actual.iterator();
    while (it1.hasNext() && it2.hasNext())
    {
      Object el1 = it1.next(), el2 = it2.next();
      if (!Objects.equals(el1, el2))
        throw new IllegalArgumentException("Element from expected: " + el1 + " is not equal to element from actual: " + el2);
    }
    if (it1.hasNext())
      throw new IllegalArgumentException("expected has more elements than actual");
    if (it2.hasNext())
      throw new IllegalArgumentException("actual has more elements than expected");
  }

  public static <T> Iterable<T> flatten(Iterable<?> iterable, int depth)
  {
    if (depth < 0)
      throw new IllegalArgumentException("depth[" + depth + "] should not be < 0");
    if (iterable == null)
      return Collections.emptySet();
    return flatten(iterable, 0, depth);
  }

  @SuppressWarnings("unchecked")
  private static <T> Iterable<T> flatten(Iterable<?> iterable, int currentDepth, int targetDepth)
  {
    if (currentDepth == targetDepth)
      return (Iterable<T>) iterable;

    return () -> new AbstractIterator<T>()
    {
      Iterator<?> parentIterator = iterable.iterator();
      Iterator<T> childIterator = Collections.emptyIterator();

      @Override
      protected T computeNext()
      {
        for (;;)
        {
          if (childIterator.hasNext())
            return childIterator.next();

          Object iterableFromParent = null;
          // Skip null Iterables
          while (parentIterator.hasNext() && (iterableFromParent = parentIterator.next()) == null);
          if (iterableFromParent == null)
            return endOfData(); // No more elements with Parent

          if (!(iterableFromParent instanceof Iterable))
            throw new IllegalArgumentException("Expected: " + Iterable.class.getTypeName()
                + ", Found: " + iterableFromParent.getClass().getTypeName()
                + ", Current depth: " + currentDepth + ", Target depth: " + targetDepth);

          Iterable<T> flattenedIterable = flatten((Iterable<?>) iterableFromParent, currentDepth + 1, targetDepth);
          childIterator = flattenedIterable.iterator();
        }
      }
    };
  }

  public static <T> Iterable<T> flatten(Iterable<?> iterable, Class<T> type)
  {
    if (type == null)
      throw new IllegalArgumentException("type can't be null");

    if (iterable == null)
      return Collections.emptySet();

    return () -> new AbstractIterator<T>()
    {
      Deque<Iterator<?>> itarators = new ArrayDeque<>(Collections.singleton(iterable.iterator()));
      Iterator<?> currentIterator = Collections.emptyIterator();

      @SuppressWarnings("unchecked")
      @Override
      protected T computeNext()
      {
        for (;;)
        {
          if (currentIterator.hasNext())
          {
            Object el = currentIterator.next();

            if (el == null)
              continue;

            if (type.isInstance(el))
              return (T) el;

            if (!(el instanceof Iterable))
              throw new IllegalArgumentException("Expected: " + Iterable.class.getTypeName() + " or " + type.getTypeName()
                  + ", Found: " + el.getClass().getTypeName());
            itarators.addLast(currentIterator);
            currentIterator = ((Iterable<?>) el).iterator();
            continue;
          }

          if (itarators.isEmpty())
            return endOfData();

          currentIterator = (Iterator<?>) itarators.removeLast();
        }
      }
    };
  }

  private static abstract class AbstractIterator<T> implements Iterator<T>
  {
    private boolean nextComputed, endOfData;
    private T data;

    @Override
    public final boolean hasNext()
    {
      if (!endOfData && !nextComputed)
      {
        data = computeNext();
        nextComputed = true;
      }
      return !endOfData;
    }

    @Override
    public final T next()
    {
      if (!hasNext())
        throw new NoSuchElementException();
      nextComputed = false;
      return data;
    }

    protected abstract T computeNext();

    protected final T endOfData()
    {
      endOfData = true;
      return null;
    }
  }
}