Implement remove on set/iter views of BlockMap

Fixes #1354.

(cherry picked from commit 637da62e34e676015aa7fbad490cf6f033e55b89)

worldedit-core/src/main/java/com/sk89q/worldedit/util/collection/BlockMap.java

@@ -19,7 +19,6 @@
 
 package com.sk89q.worldedit.util.collection;
 
-import com.google.common.collect.AbstractIterator;
 import com.sk89q.worldedit.math.BlockVector3;
 import com.sk89q.worldedit.world.block.BaseBlock;
 import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
@@ -36,6 +35,7 @@ import java.util.AbstractSet;
 import java.util.Collection;
 import java.util.Iterator;
 import java.util.Map;
+import java.util.NoSuchElementException;
 import java.util.Objects;
 import java.util.Set;
 import java.util.function.BiConsumer;
@@ -43,6 +43,7 @@ import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
+import static com.google.common.base.Preconditions.checkState;
 import static com.sk89q.worldedit.math.BitMath.fixSign;
 import static com.sk89q.worldedit.math.BitMath.mask;
 
@@ -245,26 +246,65 @@ public class BlockMap<V> extends AbstractMap<BlockVector3, V> {
             entrySet = es = new AbstractSet<Entry<BlockVector3, V>>() {
                 @Override
                 public Iterator<Entry<BlockVector3, V>> iterator() {
-                    return new AbstractIterator<Entry<BlockVector3, V>>() {
+                    return new Iterator<Entry<BlockVector3,V>>() {
 
                         private final ObjectIterator<Long2ObjectMap.Entry<Int2ObjectMap<V>>> primaryIterator
                             = Long2ObjectMaps.fastIterator(maps);
-                        private long currentGroupKey;
+                        private Long2ObjectMap.Entry<Int2ObjectMap<V>> currentPrimaryEntry;
                         private ObjectIterator<Int2ObjectMap.Entry<V>> secondaryIterator;
+                        private boolean finished;
+                        private LazyEntry next;
 
                         @Override
-                        protected Entry<BlockVector3, V> computeNext() {
+                        public boolean hasNext() {
-                            if (secondaryIterator == null || !secondaryIterator.hasNext()) {
+                            if (finished) {
-                                if (!primaryIterator.hasNext()) {
+                                return false;
-                                    return endOfData();
+                            }
+                            if (next == null) {
+                                LazyEntry proposedNext = computeNext();
+                                if (proposedNext == null) {
+                                    finished = true;
+                                    return false;
+                                }
+                                next = proposedNext;
+                            }
+                            return true;
                         }
 
-                                Long2ObjectMap.Entry<Int2ObjectMap<V>> next = primaryIterator.next();
+                        private LazyEntry computeNext() {
-                                currentGroupKey = next.getLongKey();
+                            if (secondaryIterator == null || !secondaryIterator.hasNext()) {
-                                secondaryIterator = Int2ObjectMaps.fastIterator(next.getValue());
+                                if (!primaryIterator.hasNext()) {
+                                    return null;
+                                }
+
+                                currentPrimaryEntry = primaryIterator.next();
+                                secondaryIterator = Int2ObjectMaps.fastIterator(currentPrimaryEntry.getValue());
+                                // be paranoid
+                                checkState(secondaryIterator.hasNext(),
+                                    "Should not have an empty map entry, it should have been removed!");
                             }
                             Int2ObjectMap.Entry<V> next = secondaryIterator.next();
-                            return new LazyEntry(currentGroupKey, next.getIntKey(), next.getValue());
+                            return new LazyEntry(currentPrimaryEntry.getLongKey(), next.getIntKey(), next.getValue());
+                        }
+
+                        @Override
+                        public Entry<BlockVector3, V> next() {
+                            if (!hasNext()) {
+                                throw new NoSuchElementException();
+                            }
+                            LazyEntry tmp = next;
+                            next = null;
+                            return tmp;
+                        }
+
+                        @Override
+                        public void remove() {
+                            secondaryIterator.remove();
+                            // ensure invariants hold
+                            if (currentPrimaryEntry.getValue().isEmpty()) {
+                                // the remove call cleared this map. call remove on the primary iter
+                                primaryIterator.remove();
+                            }
                         }
                     };
                 }
@@ -364,13 +404,14 @@ public class BlockMap<V> extends AbstractMap<BlockVector3, V> {
     @Override
     public V remove(Object key) {
         BlockVector3 vec = (BlockVector3) key;
-        Int2ObjectMap<V> activeMap = maps.get(toGroupKey(vec));
+        long groupKey = toGroupKey(vec);
+        Int2ObjectMap<V> activeMap = maps.get(groupKey);
         if (activeMap == null) {
             return null;
         }
         V removed = activeMap.remove(toInnerKey(vec));
         if (activeMap.isEmpty()) {
-            maps.remove(toGroupKey(vec));
+            maps.remove(groupKey);
         }
         return removed;
     }
@@ -429,7 +470,9 @@ public class BlockMap<V> extends AbstractMap<BlockVector3, V> {
         }
         if (o instanceof BlockMap) {
             // optimize by skipping entry translations:
-            return maps.equals(((BlockMap) o).maps);
+            @SuppressWarnings("unchecked")
+            BlockMap<V> other = (BlockMap<V>) o;
+            return maps.equals(other.maps);
         }
         return super.equals(o);
     }

worldedit-core/src/main/java/com/sk89q/worldedit/util/collection/Int2BaseBlockMap.java

@@ -91,6 +91,7 @@ class Int2BaseBlockMap extends AbstractInt2ObjectMap<BaseBlock> {
                         = Int2IntMaps.fastIterator(commonMap);
                     private final ObjectIterator<Entry<BaseBlock>> uncommonIter
                         = Int2ObjectMaps.fastIterator(uncommonMap);
+                    private boolean lastNextFromCommon = false;
 
                     @Override
                     public boolean hasNext() {
@@ -101,15 +102,26 @@ class Int2BaseBlockMap extends AbstractInt2ObjectMap<BaseBlock> {
                     public Entry<BaseBlock> next() {
                         if (commonIter.hasNext()) {
                             Int2IntMap.Entry e = commonIter.next();
+                            lastNextFromCommon = true;
                             return new BasicEntry<>(
                                 e.getIntKey(), assumeAsBlock(e.getIntValue())
                             );
                         }
                         if (uncommonIter.hasNext()) {
+                            lastNextFromCommon = false;
                             return uncommonIter.next();
                         }
                         throw new NoSuchElementException();
                     }
+
+                    @Override
+                    public void remove() {
+                        if (lastNextFromCommon) {
+                            commonIter.remove();
+                        } else {
+                            uncommonIter.remove();
+                        }
+                    }
                 };
             }
 

worldedit-core/src/test/java/com/sk89q/worldedit/util/collection/BlockMapTest.java

@@ -45,11 +45,11 @@ import org.mockito.MockitoAnnotations;
 
 import java.lang.reflect.Field;
 import java.util.AbstractMap;
+import java.util.Iterator;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
-import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import static com.google.common.base.Preconditions.checkNotNull;
@@ -63,24 +63,23 @@ import static org.junit.jupiter.api.Assumptions.assumeFalse;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
-import static org.mockito.Mockito.verifyZeroInteractions;
 import static org.mockito.Mockito.when;
 
 @DisplayName("An ordered block map")
 class BlockMapTest {
 /*
-    private static Platform mockedPlatform = mock(Platform.class);
+    private static final Platform MOCKED_PLATFORM = mock(Platform.class);
 
     @BeforeAll
     static void setupFakePlatform() {
-        when(mockedPlatform.getRegistries()).thenReturn(new BundledRegistries() {
+        when(MOCKED_PLATFORM.getRegistries()).thenReturn(new BundledRegistries() {
         });
-        when(mockedPlatform.getCapabilities()).thenReturn(ImmutableMap.of(
+        when(MOCKED_PLATFORM.getCapabilities()).thenReturn(ImmutableMap.of(
             Capability.WORLD_EDITING, Preference.PREFERRED,
             Capability.GAME_HOOKS, Preference.PREFERRED
         ));
         PlatformManager platformManager = WorldEdit.getInstance().getPlatformManager();
-        platformManager.register(mockedPlatform);
+        platformManager.register(MOCKED_PLATFORM);
 
         registerBlock("minecraft:air");
         registerBlock("minecraft:oak_wood");
@@ -88,7 +87,7 @@ class BlockMapTest {
 
     @AfterAll
     static void tearDownFakePlatform() throws Exception {
-        WorldEdit.getInstance().getPlatformManager().unregister(mockedPlatform);
+        WorldEdit.getInstance().getPlatformManager().unregister(MOCKED_PLATFORM);
         Field map = Registry.class.getDeclaredField("map");
         map.setAccessible(true);
         ((Map<?, ?>) map.get(BlockType.REGISTRY)).clear();
@@ -98,8 +97,6 @@ class BlockMapTest {
         BlockType.REGISTRY.register(id, new BlockType(id));
     }
 
-    @Mock
-    private Function<? super BlockVector3, ? extends BaseBlock> function;
     @Mock
     private BiFunction<? super BlockVector3, ? super BaseBlock, ? extends BaseBlock> biFunction;
     @Mock
@@ -110,7 +107,7 @@ class BlockMapTest {
     private final BaseBlock air = checkNotNull(BlockTypes.AIR).getDefaultState().toBaseBlock();
     private final BaseBlock oakWood = checkNotNull(BlockTypes.OAK_WOOD).getDefaultState().toBaseBlock();
 
-    private BlockMap<BaseBlock> map = BlockMap.createForBaseBlock();
+    private final BlockMap<BaseBlock> map = BlockMap.createForBaseBlock();
 
     @BeforeEach
     void setUp() {
@@ -186,14 +183,14 @@ class BlockMapTest {
         @DisplayName("never calls the forEach action")
         void neverCallsForEachAction() {
             map.forEach(biConsumer);
-            verifyZeroInteractions(biConsumer);
+            verifyNoMoreInteractions(biConsumer);
         }
 
         @Test
         @DisplayName("never calls the replaceAll function")
         void neverCallsReplaceAllFunction() {
             map.replaceAll(biFunction);
-            verifyZeroInteractions(biFunction);
+            verifyNoMoreInteractions(biFunction);
         }
 
         @Test
@@ -254,7 +251,7 @@ class BlockMapTest {
             assertEquals(air, map.merge(BlockVector3.ZERO, air, mergeFunction));
             assertEquals(1, map.size());
             assertEquals(air, map.get(BlockVector3.ZERO));
-            verifyZeroInteractions(mergeFunction);
+            verifyNoMoreInteractions(mergeFunction);
         }
 
     }
@@ -431,6 +428,26 @@ class BlockMapTest {
             assertEquals(0, map.size());
         }
 
+        @VariedVectorsProvider.Test
+        @DisplayName("keySet().remove(key) removes the entry from the map")
+        void keySetRemovePassesThrough(BlockVector3 vec) {
+            map.put(vec, air);
+            assertTrue(map.keySet().remove(vec));
+            assertEquals(0, map.size());
+        }
+
+        @VariedVectorsProvider.Test
+        @DisplayName("entrySet().iterator().remove() removes the entry from the map")
+        void entrySetIteratorRemovePassesThrough(BlockVector3 vec) {
+            map.put(vec, air);
+            Iterator<Map.Entry<BlockVector3, BaseBlock>> iterator = map.entrySet().iterator();
+            assertTrue(iterator.hasNext());
+            Map.Entry<BlockVector3, BaseBlock> entry = iterator.next();
+            assertEquals(entry.getKey(), vec);
+            iterator.remove();
+            assertEquals(0, map.size());
+        }
+
         @VariedVectorsProvider.Test(provideNonMatching = true)
         @DisplayName("remove(nonMatch) returns null")
         void removeNonMatchingKeyReturnsNull(BlockVector3 vec, BlockVector3 nonMatch) {