diff --git a/src/main/java/io/usethesource/vallang/impl/primitive/StringValue.java b/src/main/java/io/usethesource/vallang/impl/primitive/StringValue.java index e52c85a8..4c064c9a 100644 --- a/src/main/java/io/usethesource/vallang/impl/primitive/StringValue.java +++ b/src/main/java/io/usethesource/vallang/impl/primitive/StringValue.java @@ -33,6 +33,7 @@ import java.util.NoSuchElementException; import java.util.PrimitiveIterator; import java.util.PrimitiveIterator.OfInt; +import java.util.function.Function; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import org.checkerframework.checker.nullness.qual.Nullable; import io.usethesource.vallang.IString; @@ -1256,26 +1257,16 @@ public AbstractString rotateLeftRight() { @Override public OfInt iterator() { return new OfInt() { - final Deque todo = new ArrayDeque<>(depth); - OfInt currentLeaf = leftmostLeaf(todo, LazyConcatString.this).iterator(); + IteratorOfIterators current = new IteratorOfIterators<>(leftMostIterator(), IStringTreeNode::iterator); @Override public boolean hasNext() { - return currentLeaf.hasNext(); /* || !todo.isEmpty() is unnecessary due to post-condition of nextInt() */ + return current.hasNext(); } @Override public int nextInt() { - int next = currentLeaf.nextInt(); - - if (!currentLeaf.hasNext() && !todo.isEmpty()) { - // now we back track to the previous node we went left from, - // take the right branch and continue with its first leaf: - currentLeaf = leftmostLeaf(todo, todo.pop()).iterator(); - } - - assert currentLeaf.hasNext() || todo.isEmpty(); - return next; + return current.currentIterator().nextInt(); } }; } @@ -1283,24 +1274,67 @@ public int nextInt() { @Override public Iterator iterateParts() { return new Iterator<> () { - final Deque todo = new ArrayDeque<>(depth); - Iterator currentLeaf = leftmostLeaf(todo, LazyConcatString.this).iterateParts(); + IteratorOfIterators> current = new IteratorOfIterators<>(leftMostIterator(), IStringTreeNode::iterateParts); + @Override public boolean hasNext() { - return currentLeaf.hasNext(); /* || !todo.isEmpty() is unnecessary due to post-condition of nextInt() */ + return current.hasNext(); } + @Override public CharBuffer next() { - var next = currentLeaf.next(); + return current.currentIterator().next(); + } + }; + } - if (!currentLeaf.hasNext() && !todo.isEmpty()) { - // now we back track to the previous node we went left from, - // take the right branch and continue with its first leaf: - currentLeaf = leftmostLeaf(todo, todo.pop()).iterateParts(); - } + private static class IteratorOfIterators> { + private final Iterator base; + private final Function nested; + private @Nullable T current = null; + + IteratorOfIterators(Iterator base, Function nested) { + this.base = base; + this.nested = nested; + } + + boolean hasNext() { + return base.hasNext() || (current != null && current.hasNext()); + } + + T currentIterator() { + if (current == null || !current.hasNext()) { + current = nested.apply(base.next()); + } + return current; + } + + } - assert currentLeaf.hasNext() || todo.isEmpty(); - return next; + private Iterator leftMostIterator() { + return new Iterator<>() { + final Deque todo = new ArrayDeque<>(depth); + IStringTreeNode nextNode = leftmostLeaf(todo, LazyConcatString.this); + + + @Override + public boolean hasNext() { + return nextNode != null; /* || !todo.isEmpty() is unnecessary due to post-condition of next() */ + } + + @Override + public IStringTreeNode next() { + var result = nextNode; + if (result == null) { + throw new NoSuchElementException(); + } + if (todo.isEmpty()) { + nextNode = null; + } + else { + nextNode = leftmostLeaf(todo, todo.pop()); + } + return result; } }; }