Fix iterator implementation

CitizensDev · Nov 11, 2013 · d7e3d03 · d7e3d03
1 parent c4c1594
commit d7e3d03
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 27 deletions.
diff --git a/src/main/java/net/citizensnpcs/util/ByIdArray.java b/src/main/java/net/citizensnpcs/util/ByIdArray.java
@@ -78,7 +78,7 @@ public T get(int index) {
 
     @Override
     public Iterator<T> iterator() {
-        return new Itr();
+        return new Itr2();
     }
 
     public void put(int index, T t) {
@@ -141,11 +141,12 @@ public void trimToSize() {
         }
     }
 
-    private class Itr implements Iterator<T> {
-        private int expected = ByIdArray.this.modCount;
-        private int idx;
+    private class Itr2 implements Iterator<T> {
+        int cursor;
+        int expectedModCount = modCount;
+        int lastRet = -1;
 
-        private Itr() {
+        public Itr2() {
             if (size > 0) {
                 if (lowest > highest || highest == Integer.MIN_VALUE || highest >= elementData.length
                         || elementData[highest] == null) {
@@ -154,47 +155,54 @@ private Itr() {
                 if (lowest > highest || lowest >= elementData.length || elementData[lowest] == null) {
                     recalcLowest();
                 }
-                idx = lowest;
+                cursor = lowest;
             }
         }
 
         private void advance() {
             do {
-                idx++;
-            } while (idx != highest + 1 && elementData[idx] == null);
+                cursor++;
+            } while (cursor != highest + 1 && elementData[cursor] == null);
+        }
+
+        final void checkForComodification() {
+            if (modCount != expectedModCount)
+                throw new ConcurrentModificationException();
         }
 
         @Override
         public boolean hasNext() {
-            if (modCount != expected) {
-                throw new ConcurrentModificationException();
-            }
-            return size > 0 && highest >= idx;
+            return size > 0 && highest >= cursor;
         }
 
         @Override
         @SuppressWarnings("unchecked")
         public T next() {
-            if (modCount != expected)
-                throw new ConcurrentModificationException();
-            if (idx > highest)
-                throw new NoSuchElementException();
-            T next = (T) elementData[idx];
-            if (next == null)
+            checkForComodification();
+            int i = cursor;
+            if (cursor > highest)
                 throw new NoSuchElementException();
+            Object[] elementData = ByIdArray.this.elementData;
+            if (i >= elementData.length)
+                throw new ConcurrentModificationException();
             advance();
-            return next;
+            return (T) elementData[lastRet = i];
         }
 
         @Override
         public void remove() {
-            if (modCount != expected)
+            if (lastRet < 0)
+                throw new IllegalStateException();
+            checkForComodification();
+
+            try {
+                ByIdArray.this.fastRemove(lastRet);
+                cursor = lastRet;
+                lastRet = -1;
+                expectedModCount = modCount;
+            } catch (IndexOutOfBoundsException ex) {
                 throw new ConcurrentModificationException();
-            if (elementData[idx] == null)
-                throw new NoSuchElementException();
-            fastRemove(idx);
-            expected = modCount;
-            advance();
+            }
         }
     }
 

diff --git a/src/test/java/net/citizensnpcs/util/ByIdArrayTest.java b/src/test/java/net/citizensnpcs/util/ByIdArrayTest.java
@@ -58,9 +58,16 @@ public void testIteratorRemove() {
         itr.remove();
         itr.next();
         assertSize(array, 2);
-        assertThat(array.contains(20), is(false));
-        assertThat(array.get(10), is(1));
+        assertThat(array.contains(10), is(false));
+        assertThat(array.get(20), is(2));
         assertThat(array.get(30), is(3));
+
+        itr = array.iterator();
+        while (itr.hasNext()) {
+            itr.next();
+            itr.remove();
+        }
+        assertSize(array, 0);
     }
 
     @Test