From 1746dd2e1ae7c201e0366d6b250276345bff3379 Mon Sep 17 00:00:00 2001
From: TomyLobo <tomylobo@nurfuerspam.de>
Date: Sun, 20 Oct 2013 19:06:10 +0200
Subject: [PATCH] Snapshot restore now obeys masks again.

---
 .../commands/SnapshotUtilCommands.java        |  4 +-
 .../worldedit/snapshots/SnapshotRestore.java  | 58 ++++++++++---------
 2 files changed, 33 insertions(+), 29 deletions(-)

diff --git a/src/main/java/com/sk89q/worldedit/commands/SnapshotUtilCommands.java b/src/main/java/com/sk89q/worldedit/commands/SnapshotUtilCommands.java
index cfeb5e430..7dde042b2 100644
--- a/src/main/java/com/sk89q/worldedit/commands/SnapshotUtilCommands.java
+++ b/src/main/java/com/sk89q/worldedit/commands/SnapshotUtilCommands.java
@@ -140,10 +140,10 @@ public class SnapshotUtilCommands {
 
         try {
             // Restore snapshot
-            SnapshotRestore restore = new SnapshotRestore(chunkStore, region);
+            SnapshotRestore restore = new SnapshotRestore(chunkStore, editSession, region);
             //player.print(restore.getChunksAffected() + " chunk(s) will be loaded.");
 
-            restore.restore(editSession);
+            restore.restore();
 
             if (restore.hadTotalFailure()) {
                 String error = restore.getLastErrorMessage();
diff --git a/src/main/java/com/sk89q/worldedit/snapshots/SnapshotRestore.java b/src/main/java/com/sk89q/worldedit/snapshots/SnapshotRestore.java
index 4b20aabe4..5f76d3c35 100644
--- a/src/main/java/com/sk89q/worldedit/snapshots/SnapshotRestore.java
+++ b/src/main/java/com/sk89q/worldedit/snapshots/SnapshotRestore.java
@@ -47,12 +47,16 @@ public class SnapshotRestore {
     /**
      * Store a list of chunks that are needed and the points in them.
      */
-    private Map<BlockVector2D, ArrayList<Vector>> neededChunks =
+    private final Map<BlockVector2D, ArrayList<Vector>> neededChunks =
             new LinkedHashMap<BlockVector2D, ArrayList<Vector>>();
     /**
      * Chunk store.
      */
-    private ChunkStore chunkStore;
+    private final ChunkStore chunkStore;
+    /**
+     * Edit session.
+     */
+    private final EditSession editSession;
     /**
      * Count of the number of missing chunks.
      */
@@ -69,11 +73,13 @@ public class SnapshotRestore {
     /**
      * Construct the snapshot restore operation.
      *
-     * @param chunkStore
-     * @param region
+     * @param chunkStore The {@link ChunkStore} to restore from
+     * @param editSession The {@link EditSession} to restore to
+     * @param region The {@link Region} to restore to
      */
-    public SnapshotRestore(ChunkStore chunkStore, Region region) {
+    public SnapshotRestore(ChunkStore chunkStore, EditSession editSession, Region region) {
         this.chunkStore = chunkStore;
+        this.editSession = editSession;
 
         if (region instanceof CuboidRegion) {
             findNeededCuboidChunks(region);
@@ -83,9 +89,9 @@ public class SnapshotRestore {
     }
 
     /**
-     * Find needed chunks in the cuboid of the region.
+     * Find needed chunks in the axis-aligned bounding box of the region.
      *
-     * @param region
+     * @param region The {@link Region} to iterate
      */
     private void findNeededCuboidChunks(Region region) {
         Vector min = region.getMinimumPoint();
@@ -97,14 +103,7 @@ public class SnapshotRestore {
             for (int y = min.getBlockY(); y <= max.getBlockY(); ++y) {
                 for (int z = min.getBlockZ(); z <= max.getBlockZ(); ++z) {
                     Vector pos = new Vector(x, y, z);
-                    BlockVector2D chunkPos = ChunkStore.toChunk(pos);
-
-                    // Unidentified chunk
-                    if (!neededChunks.containsKey(chunkPos)) {
-                        neededChunks.put(chunkPos, new ArrayList<Vector>());
-                    }
-
-                    neededChunks.get(chunkPos).add(pos);
+                    checkAndAddBlock(pos);
                 }
             }
         }
@@ -113,23 +112,30 @@ public class SnapshotRestore {
     /**
      * Find needed chunks in the region.
      *
-     * @param region
+     * @param region The {@link Region} to iterate
      */
     private void findNeededChunks(Region region) {
         // First, we need to group points by chunk so that we only need
         // to keep one chunk in memory at any given moment
         for (Vector pos : region) {
-            BlockVector2D chunkPos = ChunkStore.toChunk(pos);
-
-            // Unidentified chunk
-            if (!neededChunks.containsKey(chunkPos)) {
-                neededChunks.put(chunkPos, new ArrayList<Vector>());
-            }
-
-            neededChunks.get(chunkPos).add(pos);
+            checkAndAddBlock(pos);
         }
     }
 
+    private void checkAndAddBlock(Vector pos) {
+        if (editSession.getMask() != null && !editSession.getMask().matches(editSession, pos))
+            return;
+
+        BlockVector2D chunkPos = ChunkStore.toChunk(pos);
+
+        // Unidentified chunk
+        if (!neededChunks.containsKey(chunkPos)) {
+            neededChunks.put(chunkPos, new ArrayList<Vector>());
+        }
+
+        neededChunks.get(chunkPos).add(pos);
+    }
+
     /**
      * Get the number of chunks that are needed.
      *
@@ -142,11 +148,9 @@ public class SnapshotRestore {
     /**
      * Restores to world.
      *
-     * @param editSession
      * @throws MaxChangedBlocksException
      */
-    public void restore(EditSession editSession)
-            throws MaxChangedBlocksException {
+    public void restore() throws MaxChangedBlocksException {
 
         missingChunks = new ArrayList<Vector2D>();
         errorChunks = new ArrayList<Vector2D>();