FunctionalSetIterator.java

import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.*;

public class FunctionalSetIterator<E> implements Iterator<E> {
	private FunctionalSet<E> set;

	private SimpleFunctionalSet<E> current;

	private FunctionalSet<Object> used;

	private E prev;

	private boolean canRemove = false;

	public FunctionalSetIterator(FunctionalSet<E> set, SimpleFunctionalSet<E> head) {
		this.current = head;
		this.set = set;
		used = new FunctionalSet<Object>();
	}

	@Override
	public boolean hasNext() {
		return !(current instanceof EmptySet);
	}

	@Override
	public E next() {
		if (hasNext()) {
			E elem = ((AddSet<E>) current).getElement();
			used.add(elem);
			prev = elem;
			gotoNext();
			canRemove = true;
			return elem;
		} else {
			throw new NoSuchElementException();
		}
	}

	private void gotoNext() {
		boolean loop = true;
		while (loop) {
			loop = false;
			current = current.getRemainingSet();
			while (current instanceof RemoveSet) {
				used.add(((RemoveSet<E>) current).getObject());
				current = current.getRemainingSet();
			}
			if (current instanceof AddSet && used.contains(((AddSet<E>) current).getElement())) {
				loop = true;
			}
		}
	}

	@Override
	public void remove() {
		if (canRemove) {
			set.remove(prev);
		} else
			throw new IllegalStateException();
	}
}