Implement getBlock for chunk batching extent

Also improve speed of comparators, by using ::comparingX and bitwise
ops.

worldedit-core/src/main/java/com/sk89q/worldedit/extent/reorder/ChunkBatchingExtent.java

@@ -19,22 +19,25 @@
 
 package com.sk89q.worldedit.extent.reorder;
 
+import com.google.common.collect.Table;
+import com.google.common.collect.TreeBasedTable;
 import com.sk89q.worldedit.WorldEditException;
 import com.sk89q.worldedit.extent.AbstractDelegateExtent;
 import com.sk89q.worldedit.extent.Extent;
 import com.sk89q.worldedit.function.operation.Operation;
 import com.sk89q.worldedit.function.operation.RunContext;
-import com.sk89q.worldedit.function.operation.SetLocatedBlocks;
 import com.sk89q.worldedit.math.BlockVector2;
 import com.sk89q.worldedit.math.BlockVector3;
-import com.sk89q.worldedit.util.collection.LocatedBlockList;
+import com.sk89q.worldedit.world.block.BaseBlock;
+import com.sk89q.worldedit.world.block.BlockState;
 import com.sk89q.worldedit.world.block.BlockStateHolder;
 
 import java.util.Comparator;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
-import java.util.SortedMap;
-import java.util.TreeMap;
+import java.util.Map;
+import java.util.Set;
 
 /**
  * A special extent that batches changes into Minecraft chunks. This helps
@@ -49,10 +52,12 @@ public class ChunkBatchingExtent extends AbstractDelegateExtent {
      * in. This allows for file caches to be used while loading the chunk.
      */
     private static final Comparator<BlockVector2> REGION_OPTIMIZED_SORT =
-            Comparator.comparing((BlockVector2 vec) -> vec.divide(32), BlockVector2.COMPARING_GRID_ARRANGEMENT)
+            Comparator.comparing((BlockVector2 vec) -> vec.shr(5), BlockVector2.COMPARING_GRID_ARRANGEMENT)
                     .thenComparing(BlockVector2.COMPARING_GRID_ARRANGEMENT);
 
-    private final SortedMap<BlockVector2, LocatedBlockList> batches = new TreeMap<>(REGION_OPTIMIZED_SORT);
+    private final Table<BlockVector2, BlockVector3, BaseBlock> batches =
+        TreeBasedTable.create(REGION_OPTIMIZED_SORT, BlockVector3.sortByCoordsYzx());
+    private final Set<BlockVector3> containedBlocks = new HashSet<>();
     private boolean enabled;
 
     public ChunkBatchingExtent(Extent extent) {
@@ -76,16 +81,51 @@ public class ChunkBatchingExtent extends AbstractDelegateExtent {
         return enabled;
     }
 
+    private BlockVector2 getChunkPos(BlockVector3 location) {
+        return location.shr(4).toBlockVector2();
+    }
+
+    private BlockVector3 getInChunkPos(BlockVector3 location) {
+        return BlockVector3.at(location.getX() & 15, location.getY(), location.getZ() & 15);
+    }
+
     @Override
     public <B extends BlockStateHolder<B>> boolean setBlock(BlockVector3 location, B block) throws WorldEditException {
         if (!enabled) {
             return getExtent().setBlock(location, block);
         }
-        BlockVector2 chunkPos = BlockVector2.at(location.getBlockX() >> 4, location.getBlockZ() >> 4);
-        batches.computeIfAbsent(chunkPos, k -> new LocatedBlockList()).add(location, block);
+        BlockVector2 chunkPos = getChunkPos(location);
+        BlockVector3 inChunkPos = getInChunkPos(location);
+        batches.put(chunkPos, inChunkPos, block.toBaseBlock());
+        containedBlocks.add(location);
         return true;
     }
 
+    @Override
+    public BlockState getBlock(BlockVector3 position) {
+        BaseBlock internal = getInternalBlock(position);
+        if (internal != null) {
+            return internal.toImmutableState();
+        }
+        return super.getBlock(position);
+    }
+
+    @Override
+    public BaseBlock getFullBlock(BlockVector3 position) {
+        BaseBlock internal = getInternalBlock(position);
+        if (internal != null) {
+            return internal;
+        }
+        return super.getFullBlock(position);
+    }
+
+    private BaseBlock getInternalBlock(BlockVector3 position) {
+        if (!containedBlocks.contains(position)) {
+            return null;
+        }
+        return batches.get(getChunkPos(position), getInChunkPos(position));
+    }
+
     @Override
     protected Operation commitBefore() {
         if (!commitRequired()) {
@@ -94,17 +134,21 @@ public class ChunkBatchingExtent extends AbstractDelegateExtent {
         return new Operation() {
 
             // we get modified between create/resume -- only create this on resume to prevent CME
-            private Iterator<LocatedBlockList> batchIterator;
+            private Iterator<Map<BlockVector3, BaseBlock>> batchIterator;
 
             @Override
             public Operation resume(RunContext run) throws WorldEditException {
                 if (batchIterator == null) {
-                    batchIterator = batches.values().iterator();
+                    batchIterator = batches.rowMap().values().iterator();
                 }
                 if (!batchIterator.hasNext()) {
                     return null;
                 }
-                new SetLocatedBlocks(getExtent(), batchIterator.next()).resume(run);
+                Map<BlockVector3, BaseBlock> next = batchIterator.next();
+                for (Map.Entry<BlockVector3, BaseBlock> block : next.entrySet()) {
+                    getExtent().setBlock(block.getKey(), block.getValue());
+                    containedBlocks.remove(block.getKey());
+                }
                 batchIterator.remove();
                 return this;
             }

worldedit-core/src/main/java/com/sk89q/worldedit/math/BlockVector2.java

@@ -19,7 +19,6 @@
 
 package com.sk89q.worldedit.math;
 
-import com.google.common.collect.ComparisonChain;
 import com.sk89q.worldedit.math.transform.AffineTransform;
 
 import java.util.Comparator;
@@ -48,12 +47,8 @@ public final class BlockVector2 {
      * cdef
      * </pre>
      */
-    public static final Comparator<BlockVector2> COMPARING_GRID_ARRANGEMENT = (a, b) -> {
-        return ComparisonChain.start()
-                .compare(a.getBlockZ(), b.getBlockZ())
-                .compare(a.getBlockX(), b.getBlockX())
-                .result();
-    };
+    public static final Comparator<BlockVector2> COMPARING_GRID_ARRANGEMENT =
+        Comparator.comparingInt(BlockVector2::getZ).thenComparingInt(BlockVector2::getX);
 
     public static BlockVector2 at(double x, double z) {
         return at((int) Math.floor(x), (int) Math.floor(z));
@@ -303,6 +298,27 @@ public final class BlockVector2 {
         return divide(n, n);
     }
 
+    /**
+     * Shift all components right.
+     *
+     * @param x the value to shift x by
+     * @param z the value to shift z by
+     * @return a new vector
+     */
+    public BlockVector2 shr(int x, int z) {
+        return at(this.x >> x, this.z >> z);
+    }
+
+    /**
+     * Shift all components right by {@code n}.
+     *
+     * @param n the value to shift by
+     * @return a new vector
+     */
+    public BlockVector2 shr(int n) {
+        return shr(n, n);
+    }
+
     /**
      * Get the length of the vector.
      *
@@ -532,5 +548,4 @@ public final class BlockVector2 {
     public String toString() {
         return "(" + x + ", " + z + ")";
     }
-
 }

worldedit-core/src/main/java/com/sk89q/worldedit/math/BlockVector3.java

@@ -19,13 +19,12 @@
 
 package com.sk89q.worldedit.math;
 
-import static com.google.common.base.Preconditions.checkArgument;
-
-import com.google.common.collect.ComparisonChain;
 import com.sk89q.worldedit.math.transform.AffineTransform;
 
 import java.util.Comparator;
 
+import static com.google.common.base.Preconditions.checkArgument;
+
 /**
  * An immutable 3-dimensional vector.
  */
@@ -64,13 +63,10 @@ public final class BlockVector3 {
 
     // thread-safe initialization idiom
     private static final class YzxOrderComparator {
-        private static final Comparator<BlockVector3> YZX_ORDER = (a, b) -> {
-            return ComparisonChain.start()
-                    .compare(a.y, b.y)
-                    .compare(a.z, b.z)
-                    .compare(a.x, b.x)
-                    .result();
-        };
+        private static final Comparator<BlockVector3> YZX_ORDER =
+            Comparator.comparingInt(BlockVector3::getY)
+                .thenComparingInt(BlockVector3::getZ)
+                .thenComparingInt(BlockVector3::getX);
     }
 
     /**
@@ -348,6 +344,28 @@ public final class BlockVector3 {
         return divide(n, n, n);
     }
 
+    /**
+     * Shift all components right.
+     *
+     * @param x the value to shift x by
+     * @param y the value to shift y by
+     * @param z the value to shift z by
+     * @return a new vector
+     */
+    public BlockVector3 shr(int x, int y, int z) {
+        return at(this.x >> x, this.y >> y, this.z >> z);
+    }
+
+    /**
+     * Shift all components right by {@code n}.
+     *
+     * @param n the value to shift by
+     * @return a new vector
+     */
+    public BlockVector3 shr(int n) {
+        return shr(n, n, n);
+    }
+
     /**
      * Get the length of the vector.
      *