working weight balanced tree

assira.commons/src/main/java/net/ranides/assira/collection/sets/WBTreeSet.java

@@ -49,7 +49,8 @@ public class WBTreeSet<K> extends ASortedSet<K> {
     @Override
     public ListIterator<K> iterator(K begin) {
         WBTree.Node<K> after = tree.find(begin, WBTree.FindMode.GT);
-        return new TreeIterator(after);
+        int index = WBTree.indexOf(after);
+        return new TreeIterator(index < 0 ? tree.size() : index);
     }
 
     @Override
@@ -86,6 +87,17 @@ public class WBTreeSet<K> extends ASortedSet<K> {
             super(begin, bottom, end, top);
         }
 
+        @Override
+        public final int size() {
+            WBTree.Node<K> first = $first();
+            WBTree.Node<K> last = $last();
+            if(first == null || last == null) {
+                return 0;
+            } else {
+                return WBTree.indexOf(last) - WBTree.indexOf(first) + 1;
+            }
+        }
+
         @Override
         public ListIterator<K> iterator() {
             return new SIterator();
@@ -93,296 +105,228 @@ public class WBTreeSet<K> extends ASortedSet<K> {
 
         @Override
         public ListIterator<K> iterator(K from) {
+            WBTree.Node<K> fromNode;
             if(bottom) {
-                return new SIterator(tree.find(from, WBTree.FindMode.GT));
+                fromNode = tree.find(from, WBTree.FindMode.GT);
+            }
+            else if(compare(from, begin) > 0) {
+                fromNode = tree.find(from, WBTree.FindMode.GT);
+            } else {
+                fromNode = tree.find(begin, WBTree.FindMode.GE);
             }
-            if(compare(from, begin) > 0) {
-                return new SIterator(tree.find(from, WBTree.FindMode.GT));
+            if(fromNode == null) {
+                return new SIterator(size());
             } else {
-                return new SIterator(tree.find(begin, WBTree.FindMode.GE));
+                return new SIterator(WBTree.indexOf(fromNode) - WBTree.indexOf($first()));
+            }
+        }
+
+        private WBTree.Node<K> $first() {
+            WBTree.Node<K> out = bottom ? tree.first() : tree.find(begin, WBTree.FindMode.GE);
+            if(out == null || !inside(out.value)) {
+                return null;
+            }
+            return out;
+        }
+
+        public WBTree.Node<K> $last() {
+            WBTree.Node<K> out = top ? tree.last() : tree.find(end, WBTree.FindMode.LT);
+            if(out == null || !inside(out.value)) {
+                return null;
             }
+            return out;
         }
 
         @Override
         public K first() {
-            if(bottom) {
-                K out = WBTreeSet.this.first();
-                if(!inside(out)) {
-                    throw new NoSuchElementException();
-                }
-                return out;
-            } else {
-                return iterator().next();
+            WBTree.Node<K> out = $first();
+            if(out == null) {
+                throw new NoSuchElementException();
             }
+            return out.value;
         }
 
         @Override
         public K last() {
-            if(top) {
-                K out = WBTreeSet.this.last();
-                if(!inside(out)) {
-                    throw new NoSuchElementException();
-                }
-                return out;
-            } else {
-                WBTree.Node<K> node = tree.find(end, WBTree.FindMode.LT);
-                if(node == null || !inside(node.value)) {
-                    throw new NoSuchElementException();
-                }
-                return node.value;
+            WBTree.Node<K> out = $last();
+            if(out == null) {
+                throw new NoSuchElementException();
             }
+            return out.value;
         }
 
         private class SIterator implements ListIterator<K> {
 
-            private WBTree.Node<K> current;
-            private WBTree.Node<K> prev;
-            private WBTree.Node<K> next;
-            private int index = 0;
+            private int index;
+            private int last;
 
-            public SIterator() {
-                if(bottom) {
-                    this.current = tree.first();
-                } else {
-                    this.current = tree.find(begin, WBTree.FindMode.LE);
-                    if(this.current == null) {
-                        this.current = tree.first();
-                    }
-                }
+            private final int minIndex;
+            private int maxIndex;
 
-                this.next = current;
-            }
+            private final int shift;
 
-            public SIterator(WBTree.Node<K> after) {
-                this.current = after;
-                this.next = current;
+            public SIterator() {
+                this(0);
             }
 
-            @Override
-            public boolean hasNext() {
-                if(next != null) {
-                    return inside(next.value);
-                }
-                if(current == null) {
-                    return false;
-                }
-                next = WBTree.next(current);
+            public SIterator(int index) {
+                this.index = index;
+                this.last = -1;
+                this.shift = index;
 
-                return next != null && inside(next.value);
+                WBTree.Node<K> first = $first();
+                WBTree.Node<K> last = $last();
+                if(first == null || last == null) {
+                    this.minIndex = 0;
+                    this.maxIndex = 0;
+                } else {
+                    this.minIndex = WBTree.indexOf(first);
+                    this.maxIndex = WBTree.indexOf(last) + 1;
+                }
             }
 
             @Override
-            public K next() {
-                if(next == null) {
-                    if(current == null) {
-                        throw new NoSuchElementException();
-                    }
-                    next = WBTree.next(current);
-                }
-
-                prev = current;
-                current = next;
-                next = null;
-
-                if(current == null || !inside(current.value)) {
-                    throw new NoSuchElementException();
-                }
-
-                index++;
-
-                return current.value;
+            public final boolean hasNext() {
+                return minIndex+index < maxIndex;
             }
 
             @Override
-            public boolean hasPrevious() {
-                if(prev != null) {
-                    return true;
-                }
-                if(current == null) {
-                    return false;
-                }
-                prev = WBTree.prev(current);
-
-                return prev != null && inside(prev.value);
+            public final boolean hasPrevious() {
+                return minIndex+index > minIndex;
             }
 
             @Override
-            public K previous() {
-                if(prev == null) {
-                    if(current == null) {
-                        throw new NoSuchElementException();
-                    }
-                    prev = WBTree.prev(current);
+            public final K next() {
+                if (!hasNext()) {
+                    throw new NoSuchElementException();
                 }
+                return tree.at((last = index++) + minIndex).value;
+            }
 
-                next = current;
-                current = prev;
-                prev = null;
-
-                if(current == null || !inside(current.value)) {
+            @Override
+            public final K previous() {
+                if (!hasPrevious()) {
                     throw new NoSuchElementException();
                 }
-
-                index--;
-
-                return current.value;
+                return tree.at((last = --index) + minIndex).value;
             }
 
             @Override
-            public int nextIndex() {
-                return index;
+            public final int nextIndex() {
+                return index - minIndex - shift;
             }
 
             @Override
-            public int previousIndex() {
-                return index - 1;
+            public final int previousIndex() {
+                return index - minIndex - shift - 1;
             }
 
             @Override
-            public void remove() {
-                if(current == null) {
-                    throw new IllegalStateException();
-                }
-                prev = WBTree.prev(current);
-                next = WBTree.next(current);
-
-                tree.removeNode(current);
-                current = null;
-                index--;
+            public final void add(K value) {
+                throw new UnsupportedOperationException();
             }
 
             @Override
-            public void set(K k) {
+            public final void set(K value) {
                 throw new UnsupportedOperationException();
             }
 
             @Override
-            public void add(K k) {
-                throw new UnsupportedOperationException();
+            public final void remove() {
+                if (last == -1) {
+                    throw new IllegalStateException();
+                }
+                tree.removeNode(tree.at(last + minIndex));
+                // If the last operation was a next(),
+                // we are removing an element *before* us,
+                // and we must decrease pos correspondingly.
+                if (last < index) {
+                    index--;
+                }
+                last = -1;
+                maxIndex--;
             }
+
         }
     }
 
     private class TreeIterator implements ListIterator<K> {
 
-        private WBTree.Node<K> current;
-        private WBTree.Node<K> prev;
-        private WBTree.Node<K> next;
-        private int index = 0;
+        private int index;
+
+        private int last;
+
+        private final int offset;
 
         public TreeIterator() {
-            this.current = tree.first();
-            this.next = current;
+            this(0);
         }
 
-        public TreeIterator(WBTree.Node<K> after) {
-            this.current = after;
-            this.next = current;
+        public TreeIterator(int index) {
+            this.last = -1;
+            this.index = index;
+            this.offset = index;
         }
 
         @Override
-        public boolean hasNext() {
-            if(next != null) {
-                return true;
-            }
-            if(current == null) {
-                return false;
-            }
-            next = WBTree.next(current);
-
-            return next != null;
+        public final boolean hasNext() {
+            return index < size();
         }
 
         @Override
-        public K next() {
-            if(next == null) {
-                if(current == null) {
-                    throw new NoSuchElementException();
-                }
-                next = WBTree.next(current);
-            }
-
-            prev = current;
-            current = next;
-            next = null;
-
-            if(current == null) {
-                throw new NoSuchElementException();
-            }
-
-            index++;
-
-            return current.value;
+        public final boolean hasPrevious() {
+            return index > 0;
         }
 
         @Override
-        public boolean hasPrevious() {
-            if(prev != null) {
-                return true;
-            }
-            if(current == null) {
-                return false;
+        public final K next() {
+            if (!hasNext()) {
+                throw new NoSuchElementException();
             }
-            prev = WBTree.prev(current);
-
-            return prev != null;
+            return tree.at(last = index++).value;
         }
 
         @Override
-        public K previous() {
-            if(prev == null) {
-                if(current == null) {
-                    throw new NoSuchElementException();
-                }
-                prev = WBTree.prev(current);
-            }
-
-            next = current;
-            current = prev;
-            prev = null;
-
-            if(current == null) {
+        public final K previous() {
+            if (!hasPrevious()) {
                 throw new NoSuchElementException();
             }
-
-            index--;
-
-            return current.value;
+            return tree.at(last = --index).value;
         }
 
         @Override
-        public int nextIndex() {
-            return index;
+        public final int nextIndex() {
+            return index - offset;
         }
 
         @Override
-        public int previousIndex() {
-            return index - 1;
+        public final int previousIndex() {
+            return index - offset - 1;
         }
 
         @Override
-        public void remove() {
-            if(current == null) {
-                throw new IllegalStateException();
-            }
-
-            index--;
-
-            prev = WBTree.prev(current);
-            next = WBTree.next(current);
-
-            tree.removeNode(current);
-            current = null;
+        public final void add(K value) {
+            throw new UnsupportedOperationException();
         }
 
         @Override
-        public void set(K k) {
+        public final void set(K value) {
             throw new UnsupportedOperationException();
         }
 
         @Override
-        public void add(K k) {
-            throw new UnsupportedOperationException();
+        public final void remove() {
+            if (last == -1) {
+                throw new IllegalStateException();
+            }
+            tree.removeNode(tree.at(last));
+            // If the last operation was a next(),
+            // we are removing an element *before* us,
+            // and we must decrease pos correspondingly.
+            if (last < index) {
+                index--;
+            }
+            last = -1;
         }
     }
 

assira.commons/src/test/java/net/ranides/assira/collection/sets/WBTreeSetTest.java

@@ -10,9 +10,11 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.ListIterator;
-import java.util.NoSuchElementException;
 
-import static org.junit.Assert.*;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 
 public class WBTreeSetTest {
 
@@ -115,49 +117,4 @@ public class WBTreeSetTest {
         assertEquals(4, tree.tailSet(3).size());
     }
 
-    @Test
-    public void iterator() {
-        WBTreeSet<Integer> tree = new WBTreeSet<>();
-
-        tree.add(5);
-        tree.add(3);
-//        tree.add(1);
-        tree.add(6);
-        tree.add(2);
-        tree.add(4);
-
-        ListIterator<Integer> itr = tree.iterator(3);
-        System.out.println(": " + itr.nextIndex());
-        System.out.println(">>> v = " + itr.next());
-
-        System.out.println(": " + itr.nextIndex());
-        System.out.println(">>> v = " + itr.next());
-
-        System.out.println(": " + itr.nextIndex());
-        System.out.println(">>> v = " + itr.next());
-
-        System.out.println(": " + itr.nextIndex());
-        System.out.println(">>> v = " + itr.next());
-
-        System.out.println(": " + itr.nextIndex());
-        System.out.println(">>> v = " + itr.next());
-
-        System.out.println(": " + itr.nextIndex());
-
-        assertThrows(NoSuchElementException.class, () -> {
-            itr.next();
-        });
-
-
-        itr.previous();
-
-        System.out.println("# " + itr.nextIndex());
-//        System.out.println(">>> v = " + itr.next());
-
-//        System.out.println(itr.nextIndex());
-//        System.out.println(">>> v = " + itr.next());
-//
-//        System.out.println(itr.nextIndex());
-//        System.out.println(">>> v = " + itr.next());
-    }
 }

assira.core/src/main/java/net/ranides/assira/collection/sets/ASortedSet.java

@@ -212,7 +212,7 @@ public abstract class ASortedSet<K> extends ASet<K> implements Set<K>, SortedSet
         }
 
         @Override
-        public final int size() {
+        public int size() {
             return IteratorUtils.size(iterator());
         }