CompositeIterator.java

package CompositeIteratorPattern;

import java.util.*;

public class CompositeIterator implements Iterator<MenuComponent> {
    Stack<Iterator<MenuComponent>> stack = new Stack<Iterator<MenuComponent>>();

    // Pass in iterator of the top-level menu
    public CompositeIterator(Iterator<MenuComponent> iterator) {
        stack.push(iterator);
    }

    public MenuComponent next() {
        if (hasNext()) {
            // Get the current iterator off the stack and get its next element
            // Stack peek returns the top element without removing it
            Iterator<MenuComponent> iterator = stack.peek();
            MenuComponent component = iterator.next();

            // Then add the iterator of that component to the stack
            stack.push(component.createIterator());

            return component;
        } else {
            return null;
        }
    }

    public boolean hasNext() {
        if (stack.empty()) {
            return false;
        } else {
            // Get iterator off the top of the stack and see if it has next
            Iterator<MenuComponent> iterator = stack.peek();
            if (!iterator.hasNext()) {
                // if it doesn't, pop it off the stack
                stack.pop();
                // Recurively call hasNext() to check the next iterator
                return hasNext();
            } else {
                // If it has a next element, return true
                return true;
            }
        }
    }
}