Flatten iterator for nested list

Question

Given a list which can can contain elements as well as lists, write an iterator to flatten a nested list. Please make this code better and suggest any improvements.

public class FlattenIterator implements Iterator {

    private final Stack<Iterator> iteratorStack;

    private Object next;

    public FlattenIterator(List list) {
        if (list == null) {
            throw new NullPointerException();
        }

        iteratorStack = new Stack<Iterator>();
        iteratorStack.push(list.iterator());
    }

    public void remove() {
        /* Not implemented */
    }


    private void moveToNext() {
        if ((next == null) && !iteratorStack.empty()) {
            if (iteratorStack.peek().hasNext()) {
                final Object next = iteratorStack.peek().next();
                if (next instanceof Iterator) {
                    iteratorStack.push((Iterator) next);
                    moveToNext();
                } else if (next instanceof Iterable) {
                    iteratorStack.push(((Iterable) next).iterator());
                    moveToNext();
                } else if (next instanceof Array) {
                    iteratorStack.push(Arrays.asList((Array) next).iterator());
                    moveToNext();
                } else {
                    this.next = next;
                }
            } else {
                iteratorStack.pop();
                moveToNext();
            }
        }
    }

    public boolean hasNext() {
        moveToNext();
        return next != null;
    }

    public Object next() {
        moveToNext();
        if (next == null) {
            throw new NoSuchElementException(); 
        } else {
            Object objectToReturn = next;
            next = null;
            return objectToReturn;
        }
    }
}

Answer 1

The complete lack of generics is a bit surprising, although I understand that this is difficult in this case. If you can change your architecture, you should be using the Composite Pattern instead of this absence of typing.

If the Composite Pattern is not an option, and if we want to use as much typing as possible, the class would start like

class FlatIterator<T> implements Iterator<T> {

  private final Stack<Iterator> iteratorStack = new Stack<Iterator>();
    // A stack of Iterators. Cannot express the exact type in Java generics,
    // as the types of the elements are:
    //    top: Iterator<T>
    //         ...
    // bottom: Iterator<Iterator<... Iterator<T> ...>>

  private T next;
  private boolean nextValid = false;
    // is true whenever the value in "next" is actually the next value.

  // Cannot accurately express the type for "collection".
  // However, there is no need to restrict us to "List".
  public FlatIterator(Iterable collection) {
    ...
    iteratorStack.push(collection.iterator());

    // fetch the first elem
    nextValid = moveToNext();
  }

  ...

An implementation that implicitly represents the stack by a linked List of FlatIterators can have more exact typing, if necessary.

The Iterator interface mandates that the remove method may be unimplemented, in which case it must throw an UnsupportedOperationException (see the docs). Of course, an implementation that simply delegates to the current iterator would be quite simple.

Your class does not consider that the leaf nodes in your input may be Iterable, but should not be iterated. My solution has an explicit type T, so that the recursive flattening can stop at values of these types. The moveToNext should probably look like:

// proceeds to assign the next item to "next"
// returns true if there was a next element
// returns false if no further items are available.
private boolean moveToNext() {
  while(!iteratorStack.empty()) {

    // remove depleted iterators
    if (!iteratorStack.peek().hasNext()) {
      iteratorStack.pop();
      continue;
    }

    // now an iterator sits on top        
    // consume next elem from topmost iterator
    final Object peek = iteratorStack.peek().next();

    // handle it according to type:
    if (peek instanceof T) {
      // yay, we found the next elem!
      next = (T) peek;
      return true;
    }
    else if (peek instanceof Iterable) {
      stack.push(peek.iterator());
      continue;
    }
    // array compatibility hacks would go here
    else {
      throw /* SomeAppropriateException */;
    }
  }
  // no further elems are available, all iterators are depleted
  return false;
}

– notice the absence of recursion, as that would improve neither readability, performance, or elegance when working with a stack.

And hasNext() is simply:

public boolean hasNext() {
  // if current element isn't valid, move up
  if (!nextValid) nextValid = moveToNext();
  return nextValid;
}

While next looks like:

public T next() {
  // if nothing is there, throw an error
  if (!hasNext()) throw NoSuchElementException();
  // return next elem and invalidate it
  nextValid = false;
  return next;
}

The advantage of not using null is that an Iterator may yield null as an allowed value. You cannot handle this in your implementation. The cost that I pay is an awkward extra variable, which has to be kept in sync.

Answer 2

You should probably review: com.google.common.collect.Iterators, especially com.google.common.collect.Iterators.concat() (com.google.guava:guava). That implementation uses a member variable to hold and track the current iterator, in addition to the container of ordered iterators to inspect.

If you initialize current iterator with an empty iterator value (again, see guava), you'll eliminate some problem edge cases.

moveToNext() should not use recursion. Instead, it ought to be a while loop that is determining the correct current iterator to use, and then returning the status of that iterator. After all, you've already got a Stack, why do you need to use the stack also?

Take a good look at that constructor again - you happen to be getting an iterator from List.iterator(), but there's no reason at all to care if it came from List, or some other collection.

private final Stack<Iterator> iteratorStack = new Stack<Iterator>();
public FlattenIterator(List list) { this(list.iterator);  }
public FlattenIterator(Iterator iterator) { iteratorStack.push(iterator); }

The Dependency Injection pattern would perhaps suggest that you should have a constructor that accepts a Stack<Iterator>, and factory methods that create the Stack from the initial source iterator. Not immediately clear - here, the stack is an implementation detail, so maybe it's not important. On the other hand, if you can imagine sometimes swapping in a different implementation, that will run extra logging/instrumentation/debug code, then maybe you'll want to have that constructor.

Why not FlattenIterator<T> extends Iterator<T>? Yes, the source is going to be mixed type, but you don't need to propagate that forward forever.

You might consider whether List is actually the right idea; here, you are using List<Object> as a poor man's Composite<T> (see Design Patterns, Gamma et al). What you are really doing, after all, is visiting all leaves of a tree, so maybe the tree should be explicit, rather than implicit in the list. If you have control over the caller, that's a design choice you might consider walking backwards.

It's not clear if remove() should be a no-op, as it is here, or throw an UnsupportedOperationException. At a minumum, there should be a comment explaining why it doesn't throw.

Answer 3

Your code can end up with StackOverflow exception when you add self reference to the list.

List list = new ArrayList();
list.add(list);
FlattenIterator flattenIterator = new FlattenIterator(list);
while (flattenIterator.hasNext()) {
  // do something
}