Implement restoring biomes, entities, and extended world heights (#1316)

worldedit-core/src/main/java/com/sk89q/worldedit/command/LegacySnapshotUtilCommands.java

@@ -47,9 +47,16 @@ class LegacySnapshotUtilCommands {
         this.we = we;
     }
 
+    //FAWE start - biome and entity restore
     void restore(
-            Actor actor, World world, LocalSession session, EditSession editSession,
+            Actor actor,
-            String snapshotName
+            World world,
+            LocalSession session,
+            EditSession editSession,
+            String snapshotName,
+            boolean restoreBiomes,
+            boolean restoreEntities
+            //FAWE end
     ) throws WorldEditException {
         LocalConfiguration config = we.getConfiguration();
 
@@ -108,8 +115,9 @@ class LegacySnapshotUtilCommands {
 
         try {
             // Restore snapshot
-            SnapshotRestore restore = new SnapshotRestore(chunkStore, editSession, region);
+            //FAWE start - biome and entity restore
-            //player.print(restore.getChunksAffected() + " chunk(s) will be loaded.");
+            SnapshotRestore restore = new SnapshotRestore(chunkStore, editSession, region, restoreBiomes, restoreEntities);
+            //FAWE end
 
             restore.restore();
 

worldedit-core/src/main/java/com/sk89q/worldedit/command/SnapshotUtilCommands.java

@@ -37,6 +37,7 @@ import com.sk89q.worldedit.world.snapshot.experimental.SnapshotRestore;
 import org.enginehub.piston.annotation.Command;
 import org.enginehub.piston.annotation.CommandContainer;
 import org.enginehub.piston.annotation.param.Arg;
+import org.enginehub.piston.annotation.param.Switch;
 
 import java.io.IOException;
 import java.net.URI;
@@ -68,13 +69,23 @@ public class SnapshotUtilCommands {
     public void restore(
             Actor actor, World world, LocalSession session, EditSession editSession,
             @Arg(name = "snapshot", desc = "The snapshot to restore", def = "")
-                    String snapshotName
+                    String snapshotName,
+            //FAWE start - biome and entity restore
+            @Switch(name = 'b', desc = "If biomes should be restored. If restoring from pre-1.15 to 1.15+, biomes may not be " +
+                    "exactly the same due to 3D biomes.")
+                    boolean restoreBiomes,
+            @Switch(name = 'e', desc = "If entities should be restored. Will cause issues with duplicate entities if all " +
+                    "original entities were not removed.")
+                    boolean restoreEntities
+            //FAWE end
     ) throws WorldEditException, IOException {
         LocalConfiguration config = we.getConfiguration();
         checkSnapshotsConfigured(config);
 
         if (config.snapshotRepo != null) {
-            legacy.restore(actor, world, session, editSession, snapshotName);
+            //FAWE start - biome and entity restore
+            legacy.restore(actor, world, session, editSession, snapshotName, restoreBiomes, restoreEntities);
+            //FAWE end
             return;
         }
 
@@ -116,8 +127,9 @@ public class SnapshotUtilCommands {
 
         try {
             // Restore snapshot
-            SnapshotRestore restore = new SnapshotRestore(snapshot, editSession, region);
+            //FAWE start - biome and entity restore
-            //player.print(restore.getChunksAffected() + " chunk(s) will be loaded.");
+            SnapshotRestore restore = new SnapshotRestore(snapshot, editSession, region, restoreBiomes, restoreEntities);
+            //FAWE end
 
             restore.restore();
 

worldedit-core/src/main/java/com/sk89q/worldedit/entity/BaseEntity.java

@@ -47,9 +47,7 @@ public class BaseEntity implements NbtValued {
 
     private final EntityType type;
     @Nullable
-    //FAWE start - use LZ<CBT> over CompoundTag
     private LazyReference<CompoundBinaryTag> nbtData;
-    //FAWE end
 
     /**
      * Create a new base entity.
@@ -95,6 +93,12 @@ public class BaseEntity implements NbtValued {
         setNbtReference(other.getNbtReference());
     }
 
+    @Nullable
+    @Override
+    public LazyReference<CompoundBinaryTag> getNbtReference() {
+        return nbtData;
+    }
+
     @Override
     public void setNbtReference(@Nullable LazyReference<CompoundBinaryTag> nbtData) {
         this.nbtData = nbtData;
@@ -113,12 +117,6 @@ public class BaseEntity implements NbtValued {
     public BaseEntity(CompoundTag tag) {
         this(EntityTypes.parse(tag.getString("Id")), tag);
     }
-
-    @Nullable
-    @Override
-    public LazyReference<CompoundBinaryTag> getNbtReference() {
-        return nbtData;
-    }
     //FAWE end
 
 }

worldedit-core/src/main/java/com/sk89q/worldedit/world/chunk/AnvilChunk.java

@@ -41,9 +41,7 @@ import java.util.Map;
 
 public class AnvilChunk implements Chunk {
 
-    //FAWE start - use CBT > CT
     private final CompoundBinaryTag rootTag;
-    //FAWE end
     private final byte[][] blocks;
     private final byte[][] blocksAdd;
     private final byte[][] data;
@@ -52,9 +50,6 @@ public class AnvilChunk implements Chunk {
 
     private Map<BlockVector3, CompoundBinaryTag> tileEntities;
 
-
-    //FAWE start
-
     /**
      * Construct the chunk with a compound tag.
      *
@@ -66,7 +61,6 @@ public class AnvilChunk implements Chunk {
     public AnvilChunk(CompoundTag tag) throws DataException {
         this(tag.asBinaryTag());
     }
-    //FAWE end
 
     /**
      * Construct the chunk with a compound tag.
@@ -84,7 +78,6 @@ public class AnvilChunk implements Chunk {
         blocksAdd = new byte[16][16 * 16 * 8];
         data = new byte[16][16 * 16 * 8];
 
-        //FAWE start - use *BinaryTag > *Tag
         ListBinaryTag sections = NbtUtils.getChildTag(rootTag, "Sections", BinaryTagTypes.LIST);
 
         for (BinaryTag rawSectionTag : sections) {
@@ -135,7 +128,6 @@ public class AnvilChunk implements Chunk {
             }
         }
     }
-    //FAWE end
 
     private int getBlockID(BlockVector3 position) throws DataException {
         int x = position.getX() - rootX * 16;
@@ -201,7 +193,6 @@ public class AnvilChunk implements Chunk {
      * Used to load the tile entities.
      */
     private void populateTileEntities() throws DataException {
-        //FAWE start - use *BinaryTag > *Tag
         ListBinaryTag tags = NbtUtils.getChildTag(rootTag, "TileEntities", BinaryTagTypes.LIST);
 
         tileEntities = new HashMap<>();
@@ -248,7 +239,6 @@ public class AnvilChunk implements Chunk {
             tileEntities.put(vec, values.build());
         }
     }
-    //FAWE end
 
     /**
      * Get the map of tags keyed to strings for a block's tile entity data. May
@@ -260,7 +250,6 @@ public class AnvilChunk implements Chunk {
      * @throws DataException thrown if there is a data error
      */
     @Nullable
-    //FAWE start - use *BinaryTag > * Tag
     private CompoundBinaryTag getBlockTileEntity(BlockVector3 position) throws DataException {
         if (tileEntities == null) {
             populateTileEntities();
@@ -289,6 +278,5 @@ public class AnvilChunk implements Chunk {
 
         return state.toBaseBlock();
     }
-    //FAWE end
 
 }

worldedit-core/src/main/java/com/sk89q/worldedit/world/chunk/AnvilChunk13.java

@@ -21,22 +21,29 @@ package com.sk89q.worldedit.world.chunk;
 
 import com.fastasyncworldedit.core.util.NbtUtils;
 import com.sk89q.jnbt.CompoundTag;
+import com.sk89q.worldedit.entity.BaseEntity;
 import com.sk89q.worldedit.math.BlockVector3;
 import com.sk89q.worldedit.registry.state.Property;
+import com.sk89q.worldedit.util.concurrency.LazyReference;
 import com.sk89q.worldedit.util.nbt.BinaryTag;
 import com.sk89q.worldedit.util.nbt.BinaryTagTypes;
 import com.sk89q.worldedit.util.nbt.CompoundBinaryTag;
 import com.sk89q.worldedit.util.nbt.IntBinaryTag;
 import com.sk89q.worldedit.util.nbt.ListBinaryTag;
 import com.sk89q.worldedit.world.DataException;
+import com.sk89q.worldedit.world.biome.BiomeType;
+import com.sk89q.worldedit.world.biome.BiomeTypes;
 import com.sk89q.worldedit.world.block.BaseBlock;
 import com.sk89q.worldedit.world.block.BlockState;
 import com.sk89q.worldedit.world.block.BlockType;
 import com.sk89q.worldedit.world.block.BlockTypes;
+import com.sk89q.worldedit.world.entity.EntityTypes;
 import com.sk89q.worldedit.world.storage.InvalidFormatException;
 
 import javax.annotation.Nullable;
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -44,17 +51,15 @@ import java.util.Map;
  */
 public class AnvilChunk13 implements Chunk {
 
-    //FAWE start - CBT > CT
+    protected final CompoundBinaryTag rootTag;
-    private final CompoundBinaryTag rootTag;
-    //FAWE end
     private final BlockState[][] blocks;
-    private final int rootX;
+    //FAWE start - biome and entity restore
-    private final int rootZ;
+    protected BiomeType[] biomes;
-
+    //FAWE end
     private Map<BlockVector3, CompoundBinaryTag> tileEntities;
-
+    //FAWE start - biome and entity restore
-
+    private List<BaseEntity> entities;
-    //FAWE start
+    //FAWE end
 
     /**
      * Construct the chunk with a compound tag.
@@ -67,7 +72,6 @@ public class AnvilChunk13 implements Chunk {
     public AnvilChunk13(CompoundTag tag) throws DataException {
         this(tag.asBinaryTag());
     }
-    //FAWE end
 
     /**
      * Construct the chunk with a compound tag.
@@ -78,12 +82,8 @@ public class AnvilChunk13 implements Chunk {
     public AnvilChunk13(CompoundBinaryTag tag) throws DataException {
         rootTag = tag;
 
-        rootX = NbtUtils.getChildTag(rootTag, "xPos", BinaryTagTypes.INT).value();
-        rootZ = NbtUtils.getChildTag(rootTag, "zPos", BinaryTagTypes.INT).value();
-
         blocks = new BlockState[16][];
 
-        //FAWE start - use *BinaryTag > *Tag
         ListBinaryTag sections = NbtUtils.getChildTag(rootTag, "Sections", BinaryTagTypes.LIST);
 
         for (BinaryTag rawSectionTag : sections) {
@@ -132,7 +132,6 @@ public class AnvilChunk13 implements Chunk {
                 }
                 palette[paletteEntryId] = blockState;
             }
-            //FAWE end
 
             // parse block states
             long[] blockStatesSerialized = NbtUtils.getChildTag(sectionTag, "BlockStates", BinaryTagTypes.LONG_ARRAY).value();
@@ -191,7 +190,6 @@ public class AnvilChunk13 implements Chunk {
         if (rootTag.get("TileEntities") == null) {
             return;
         }
-        //FAWE start - use *BinaryTag > *Tag
         ListBinaryTag tags = NbtUtils.getChildTag(rootTag, "TileEntities", BinaryTagTypes.LIST);
 
         for (BinaryTag tag : tags) {
@@ -208,7 +206,6 @@ public class AnvilChunk13 implements Chunk {
             BlockVector3 vec = BlockVector3.at(x, y, z);
             tileEntities.put(vec, t);
         }
-        //FAWE end
     }
 
     /**
@@ -221,7 +218,6 @@ public class AnvilChunk13 implements Chunk {
      * @throws DataException thrown if there is a data error
      */
     @Nullable
-    //FAWE start - use *BinaryTag > *Tag
     private CompoundBinaryTag getBlockTileEntity(BlockVector3 position) throws DataException {
         if (tileEntities == null) {
             populateTileEntities();
@@ -234,9 +230,11 @@ public class AnvilChunk13 implements Chunk {
 
     @Override
     public BaseBlock getBlock(BlockVector3 position) throws DataException {
-        int x = position.getX() - rootX * 16;
+        //FAWE start - simplified
+        int x = position.getX() & 15;
         int y = position.getY();
-        int z = position.getZ() - rootZ * 16;
+        int z = position.getZ() & 15;
+        //FAWE end
 
         int section = y >> 4;
         int yIndex = y & 0x0F;
@@ -256,6 +254,61 @@ public class AnvilChunk13 implements Chunk {
 
         return state.toBaseBlock();
     }
+    //FAWE start - biome and entity restore
+
+    @Override
+    public BiomeType getBiome(final BlockVector3 position) throws DataException {
+        if (biomes == null) {
+            populateBiomes();
+        }
+        int rx = position.getX() & 15;
+        int rz = position.getZ() & 15;
+        return biomes[rz << 4 | rx];
+    }
+
+    @Override
+    public List<BaseEntity> getEntities() throws DataException {
+        if (entities == null) {
+            populateEntities();
+        }
+        return entities;
+    }
+
+    /**
+     * Used to load the biomes.
+     */
+    private void populateEntities() throws DataException {
+        entities = new ArrayList<>();
+        if (rootTag.get("Entities") == null) {
+            return;
+        }
+        ListBinaryTag tags = NbtUtils.getChildTag(rootTag, "Entities", BinaryTagTypes.LIST);
+
+        for (BinaryTag tag : tags) {
+            if (!(tag instanceof CompoundBinaryTag)) {
+                throw new InvalidFormatException("CompoundTag expected in Entities");
+            }
+
+            CompoundBinaryTag t = (CompoundBinaryTag) tag;
+
+            entities.add(new BaseEntity(EntityTypes.get(t.getString("id")), LazyReference.computed(t)));
+        }
+
+    }
+
+    /**
+     * Used to load the biomes.
+     */
+    private void populateBiomes() throws DataException {
+        biomes = new BiomeType[256];
+        if (rootTag.get("Biomes") == null) {
+            return;
+        }
+        int[] stored = NbtUtils.getChildTag(rootTag, "Biomes", BinaryTagTypes.INT_ARRAY).value();
+        for (int i = 0; i < 256; i++) {
+            biomes[i] = BiomeTypes.getLegacy(stored[i]);
+        }
+    }
     //FAWE end
 
 }

worldedit-core/src/main/java/com/sk89q/worldedit/world/chunk/AnvilChunk15.java

@@ -0,0 +1,82 @@
+/*
+ * WorldEdit, a Minecraft world manipulation toolkit
+ * Copyright (C) sk89q <http://www.sk89q.com>
+ * Copyright (C) WorldEdit team and contributors
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.sk89q.worldedit.world.chunk;
+
+import com.fastasyncworldedit.core.util.NbtUtils;
+import com.sk89q.jnbt.CompoundTag;
+import com.sk89q.worldedit.math.BlockVector3;
+import com.sk89q.worldedit.util.nbt.BinaryTagTypes;
+import com.sk89q.worldedit.util.nbt.CompoundBinaryTag;
+import com.sk89q.worldedit.world.DataException;
+import com.sk89q.worldedit.world.biome.BiomeType;
+import com.sk89q.worldedit.world.biome.BiomeTypes;
+
+/**
+ * The chunk format for Minecraft 1.15 and newer
+ */
+//FAWE start - biome and entity restore
+public class AnvilChunk15 extends AnvilChunk13 {
+    //FAWE end
+
+    /**
+     * Construct the chunk with a compound tag.
+     *
+     * @param tag the tag to read
+     * @throws DataException on a data error
+     * @deprecated Use {@link #AnvilChunk15(CompoundBinaryTag)}
+     */
+    @Deprecated
+    public AnvilChunk15(CompoundTag tag) throws DataException {
+        super(tag);
+    }
+
+    /**
+     * Construct the chunk with a compound tag.
+     *
+     * @param tag the tag to read
+     * @throws DataException on a data error
+     */
+    public AnvilChunk15(CompoundBinaryTag tag) throws DataException {
+        super(tag);
+    }
+
+    @Override
+    public BiomeType getBiome(final BlockVector3 position) throws DataException {
+        if (biomes == null) {
+            populateBiomes();
+        }
+        int x = (position.getX() & 15) >> 2;
+        int y = position.getY() >> 2;
+        int z = (position.getZ() & 15) >> 2;
+        return biomes[y << 4 | z << 2 | x];
+    }
+
+    private void populateBiomes() throws DataException {
+        biomes = new BiomeType[1024];
+        if (rootTag.get("Biomes") == null) {
+            return;
+        }
+        int[] stored = NbtUtils.getChildTag(rootTag, "Biomes", BinaryTagTypes.INT_ARRAY).value();
+        for (int i = 0; i < 1024; i++) {
+            biomes[i] = BiomeTypes.getLegacy(stored[i]);
+        }
+    }
+
+}

worldedit-core/src/main/java/com/sk89q/worldedit/world/chunk/AnvilChunk16.java

@@ -28,9 +28,7 @@ import com.sk89q.worldedit.world.storage.InvalidFormatException;
 /**
  * The chunk format for Minecraft 1.16 and newer
  */
-public class AnvilChunk16 extends AnvilChunk13 {
+public class AnvilChunk16 extends AnvilChunk15 {
-
-    //FAWE start
 
     /**
      * Construct the chunk with a compound tag.
@@ -53,7 +51,6 @@ public class AnvilChunk16 extends AnvilChunk13 {
     public AnvilChunk16(CompoundBinaryTag tag) throws DataException {
         super(tag);
     }
-    //FAWE end
 
     @Override
     protected void readBlockStates(BlockState[] palette, long[] blockStatesSerialized, BlockState[] chunkSectionBlocks) throws

worldedit-core/src/main/java/com/sk89q/worldedit/world/chunk/AnvilChunk17.java

@@ -0,0 +1,314 @@
+/*
+ * WorldEdit, a Minecraft world manipulation toolkit
+ * Copyright (C) sk89q <http://www.sk89q.com>
+ * Copyright (C) WorldEdit team and contributors
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.sk89q.worldedit.world.chunk;
+
+import com.fastasyncworldedit.core.util.NbtUtils;
+import com.sk89q.jnbt.CompoundTag;
+import com.sk89q.worldedit.entity.BaseEntity;
+import com.sk89q.worldedit.math.BlockVector3;
+import com.sk89q.worldedit.registry.state.Property;
+import com.sk89q.worldedit.util.concurrency.LazyReference;
+import com.sk89q.worldedit.util.nbt.BinaryTag;
+import com.sk89q.worldedit.util.nbt.BinaryTagTypes;
+import com.sk89q.worldedit.util.nbt.CompoundBinaryTag;
+import com.sk89q.worldedit.util.nbt.IntBinaryTag;
+import com.sk89q.worldedit.util.nbt.ListBinaryTag;
+import com.sk89q.worldedit.world.DataException;
+import com.sk89q.worldedit.world.biome.BiomeType;
+import com.sk89q.worldedit.world.biome.BiomeTypes;
+import com.sk89q.worldedit.world.block.BaseBlock;
+import com.sk89q.worldedit.world.block.BlockState;
+import com.sk89q.worldedit.world.block.BlockType;
+import com.sk89q.worldedit.world.block.BlockTypes;
+import com.sk89q.worldedit.world.entity.EntityTypes;
+import com.sk89q.worldedit.world.storage.InvalidFormatException;
+
+import javax.annotation.Nullable;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Supplier;
+
+/**
+ * The chunk format for Minecraft 1.17
+ */
+public class AnvilChunk17 implements Chunk {
+
+    private final CompoundBinaryTag rootTag;
+    private final Supplier<CompoundBinaryTag> entityTagSupplier;
+    private BiomeType[] biomes;
+    private BlockState[][] blocks;
+    private Map<BlockVector3, CompoundBinaryTag> tileEntities;
+    private List<BaseEntity> entities;
+    // initialise with default values
+    private int minSectionPosition = 0;
+    private int maxSectionPosition = 15;
+    private int sectionCount = 16;
+
+    /**
+     * Construct the chunk with a compound tag.
+     *
+     * @param tag the tag to read
+     * @throws DataException on a data error
+     * @deprecated Use {@link #AnvilChunk17(CompoundBinaryTag, Supplier)}
+     */
+    @Deprecated
+    public AnvilChunk17(CompoundTag tag, Supplier<CompoundTag> entitiesTag) throws DataException {
+        this(tag.asBinaryTag(), () -> {
+            CompoundTag compoundTag = entitiesTag.get();
+            if (compoundTag == null) {
+                return null;
+            }
+            return compoundTag.asBinaryTag();
+        });
+    }
+
+    /**
+     * Construct the chunk with a compound tag.
+     *
+     * @param tag       the tag to read
+     * @param entityTag supplier for the entity compound tag found in the entities folder mca files. Not accessed unless
+     *                  {@link #getEntities()} is called
+     * @throws DataException on a data error
+     */
+    public AnvilChunk17(CompoundBinaryTag tag, Supplier<CompoundBinaryTag> entityTag) throws DataException {
+        rootTag = tag;
+        entityTagSupplier = entityTag;
+
+        blocks = new BlockState[16][]; // initialise with default length
+
+        ListBinaryTag sections = NbtUtils.getChildTag(rootTag, "Sections", BinaryTagTypes.LIST);
+
+        for (BinaryTag rawSectionTag : sections) {
+            if (!(rawSectionTag instanceof CompoundBinaryTag)) {
+                continue;
+            }
+
+            CompoundBinaryTag sectionTag = (CompoundBinaryTag) rawSectionTag;
+            if (sectionTag.get("Y") == null || sectionTag.get("BlockStates") == null) {
+                continue; // Empty section.
+            }
+
+            int y = NbtUtils.getChildTag(sectionTag, "Y", BinaryTagTypes.BYTE).value();
+            updateSectionIndexRange(y);
+
+            // parse palette
+            ListBinaryTag paletteEntries = sectionTag.getList("Palette", BinaryTagTypes.COMPOUND);
+            int paletteSize = paletteEntries.size();
+            if (paletteSize == 0) {
+                continue;
+            }
+            BlockState[] palette = new BlockState[paletteSize];
+            for (int paletteEntryId = 0; paletteEntryId < paletteSize; paletteEntryId++) {
+                CompoundBinaryTag paletteEntry = (CompoundBinaryTag) paletteEntries.get(paletteEntryId);
+                BlockType type = BlockTypes.get(paletteEntry.getString("Name"));
+                if (type == null) {
+                    throw new InvalidFormatException("Invalid block type: " + paletteEntry.getString("Name"));
+                }
+                BlockState blockState = type.getDefaultState();
+                if (paletteEntry.get("Properties") != null) {
+                    CompoundBinaryTag properties = NbtUtils.getChildTag(paletteEntry, "Properties", BinaryTagTypes.COMPOUND);
+                    for (Property<?> property : blockState.getStates().keySet()) {
+                        if (properties.get(property.getName()) != null) {
+                            String value = properties.getString(property.getName());
+                            try {
+                                blockState = getBlockStateWith(blockState, property, value);
+                            } catch (IllegalArgumentException e) {
+                                throw new InvalidFormatException("Invalid block state for " + blockState
+                                        .getBlockType()
+                                        .getId() + ", " + property.getName() + ": " + value);
+                            }
+                        }
+                    }
+                }
+                palette[paletteEntryId] = blockState;
+            }
+
+            // parse block states
+            long[] blockStatesSerialized = NbtUtils.getChildTag(sectionTag, "BlockStates", BinaryTagTypes.LONG_ARRAY).value();
+
+            BlockState[] chunkSectionBlocks = new BlockState[4096];
+            blocks[y - minSectionPosition] = chunkSectionBlocks;
+
+            readBlockStates(palette, blockStatesSerialized, chunkSectionBlocks);
+        }
+    }
+
+    private void updateSectionIndexRange(int layer) {
+        if (layer >= minSectionPosition && layer <= maxSectionPosition) {
+            return;
+        }
+        if (layer < minSectionPosition) {
+            int diff = minSectionPosition - layer;
+            sectionCount += diff;
+            BlockState[][] tmpBlocks = new BlockState[sectionCount][];
+            System.arraycopy(blocks, 0, tmpBlocks, diff, blocks.length);
+            blocks = tmpBlocks;
+            minSectionPosition = layer;
+        } else {
+            int diff = layer - maxSectionPosition;
+            sectionCount += diff;
+            BlockState[][] tmpBlocks = new BlockState[sectionCount][];
+            System.arraycopy(blocks, 0, tmpBlocks, 0, blocks.length);
+            blocks = tmpBlocks;
+            maxSectionPosition = layer;
+        }
+    }
+
+    protected void readBlockStates(BlockState[] palette, long[] blockStatesSerialized, BlockState[] chunkSectionBlocks) throws
+            InvalidFormatException {
+        PackedIntArrayReader reader = new PackedIntArrayReader(blockStatesSerialized);
+        for (int blockPos = 0; blockPos < chunkSectionBlocks.length; blockPos++) {
+            int index = reader.get(blockPos);
+            if (index >= palette.length) {
+                throw new InvalidFormatException("Invalid block state table entry: " + index);
+            }
+            chunkSectionBlocks[blockPos] = palette[index];
+        }
+    }
+
+    private <T> BlockState getBlockStateWith(BlockState source, Property<T> property, String value) {
+        return source.with(property, property.getValueFor(value));
+    }
+
+    /**
+     * Used to load the tile entities.
+     */
+    private void populateTileEntities() throws DataException {
+        tileEntities = new HashMap<>();
+        if (rootTag.get("TileEntities") == null) {
+            return;
+        }
+        ListBinaryTag tags = NbtUtils.getChildTag(rootTag, "TileEntities", BinaryTagTypes.LIST);
+
+        for (BinaryTag tag : tags) {
+            if (!(tag instanceof CompoundBinaryTag)) {
+                throw new InvalidFormatException("CompoundTag expected in TileEntities");
+            }
+
+            CompoundBinaryTag t = (CompoundBinaryTag) tag;
+
+            int x = ((IntBinaryTag) t.get("x")).value();
+            int y = ((IntBinaryTag) t.get("y")).value();
+            int z = ((IntBinaryTag) t.get("z")).value();
+
+            BlockVector3 vec = BlockVector3.at(x, y, z);
+            tileEntities.put(vec, t);
+        }
+    }
+
+    /**
+     * Get the map of tags keyed to strings for a block's tile entity data. May
+     * return null if there is no tile entity data. Not public yet because
+     * what this function returns isn't ideal for usage.
+     *
+     * @param position the position
+     * @return the compound tag for that position, which may be null
+     * @throws DataException thrown if there is a data error
+     */
+    @Nullable
+    private CompoundBinaryTag getBlockTileEntity(BlockVector3 position) throws DataException {
+        if (tileEntities == null) {
+            populateTileEntities();
+        }
+
+        return tileEntities.get(position);
+    }
+
+    @Override
+    public BaseBlock getBlock(BlockVector3 position) throws DataException {
+        int x = position.getX() & 15;
+        int y = position.getY();
+        int z = position.getZ() & 15;
+
+        int section = y >> 4;
+        int yIndex = y & 0x0F;
+
+        if (section < minSectionPosition || section > maxSectionPosition) {
+            throw new DataException("Chunk does not contain position " + position);
+        }
+
+        BlockState[] sectionBlocks = blocks[section - minSectionPosition];
+        BlockState state = sectionBlocks != null ? sectionBlocks[(yIndex << 8) | (z << 4) | x] : BlockTypes.AIR.getDefaultState();
+
+        CompoundBinaryTag tileEntity = getBlockTileEntity(position);
+
+        if (tileEntity != null) {
+            return state.toBaseBlock(tileEntity);
+        }
+
+        return state.toBaseBlock();
+    }
+
+    @Override
+    public BiomeType getBiome(final BlockVector3 position) throws DataException {
+        if (biomes == null) {
+            populateBiomes();
+        }
+        int x = (position.getX() & 15) >> 2;
+        int y = (position.getY() - (minSectionPosition << 4)) >> 2; // normalize
+        int z = (position.getZ() & 15) >> 2;
+        return biomes[y << 4 | z << 2 | x];
+    }
+
+    private void populateBiomes() throws DataException {
+        biomes = new BiomeType[64 * blocks.length];
+        if (rootTag.get("Biomes") == null) {
+            return;
+        }
+        int[] stored = NbtUtils.getChildTag(rootTag, "Biomes", BinaryTagTypes.INT_ARRAY).value();
+        for (int i = 0; i < 1024; i++) {
+            biomes[i] = BiomeTypes.getLegacy(stored[i]);
+        }
+    }
+
+    @Override
+    public List<BaseEntity> getEntities() throws DataException {
+        if (entities == null) {
+            populateEntities();
+        }
+        return entities;
+    }
+
+    /**
+     * Used to load the biomes.
+     */
+    private void populateEntities() throws DataException {
+        entities = new ArrayList<>();
+        CompoundBinaryTag entityTag;
+        if (entityTagSupplier == null || (entityTag = entityTagSupplier.get()) == null) {
+            return;
+        }
+        ListBinaryTag tags = NbtUtils.getChildTag(entityTag, "Entities", BinaryTagTypes.LIST);
+
+        for (BinaryTag tag : tags) {
+            if (!(tag instanceof CompoundBinaryTag)) {
+                throw new InvalidFormatException("CompoundTag expected in Entities");
+            }
+
+            CompoundBinaryTag t = (CompoundBinaryTag) tag;
+
+            entities.add(new BaseEntity(EntityTypes.get(t.getString("id")), LazyReference.computed(t)));
+        }
+
+    }
+
+}

worldedit-core/src/main/java/com/sk89q/worldedit/world/chunk/Chunk.java

@@ -19,10 +19,18 @@
 
 package com.sk89q.worldedit.world.chunk;
 
+import com.sk89q.worldedit.entity.BaseEntity;
+import com.sk89q.worldedit.entity.Entity;
 import com.sk89q.worldedit.math.BlockVector3;
+import com.sk89q.worldedit.util.Location;
 import com.sk89q.worldedit.world.DataException;
+import com.sk89q.worldedit.world.biome.BiomeType;
 import com.sk89q.worldedit.world.block.BaseBlock;
 
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
 /**
  * A 16 by 16 block chunk.
  */
@@ -37,4 +45,25 @@ public interface Chunk {
      */
     BaseBlock getBlock(BlockVector3 position) throws DataException;
 
+    //FAWE start - biome and entity restore
+    /**
+     * Get a biome.
+     *
+     * @param position the position of the block
+     * @return block the block
+     * @throws DataException thrown on data error
+     */
+    default BiomeType getBiome(BlockVector3 position) throws DataException {
+        return null;
+    }
+
+    /**
+     * Get the stored entities.
+     * @return list of stored entities
+     */
+    default List<BaseEntity> getEntities() throws DataException {
+        return Collections.emptyList();
+    }
+    //FAWE end
+
 }

worldedit-core/src/main/java/com/sk89q/worldedit/world/chunk/OldChunk.java

@@ -43,9 +43,7 @@ import java.util.Map;
  */
 public class OldChunk implements Chunk {
 
-    //FAWE start
     private final CompoundBinaryTag rootTag;
-    //FAWE end
     private final byte[] blocks;
     private final byte[] data;
     private final int rootX;
@@ -53,8 +51,6 @@ public class OldChunk implements Chunk {
 
     private Map<BlockVector3, CompoundBinaryTag> tileEntities;
 
-    //FAWE start
-
     /**
      * Construct the chunk with a compound tag.
      *
@@ -66,7 +62,6 @@ public class OldChunk implements Chunk {
     public OldChunk(CompoundTag tag) throws DataException {
         this(tag.asBinaryTag());
     }
-    //FAWE end
 
     /**
      * Construct the chunk with a compound tag.
@@ -74,7 +69,6 @@ public class OldChunk implements Chunk {
      * @param tag the tag
      * @throws DataException if there is an error getting the chunk data
      */
-    //FAWE start - use *BinaryTag > *Tag
     public OldChunk(CompoundBinaryTag tag) throws DataException {
         rootTag = tag;
 
@@ -211,6 +205,5 @@ public class OldChunk implements Chunk {
 
         return state.toBaseBlock();
     }
-    //FAWE end
 
 }

worldedit-core/src/main/java/com/sk89q/worldedit/world/chunk/package-info.java

@@ -0,0 +1,7 @@
+/**
+ * The following classes are FAWE additions:
+ *
+ * @see com.sk89q.worldedit.world.chunk.AnvilChunk15
+ * @see com.sk89q.worldedit.world.chunk.AnvilChunk17
+ */
+package com.sk89q.worldedit.world.chunk;

worldedit-core/src/main/java/com/sk89q/worldedit/world/snapshot/SnapshotRestore.java

@@ -22,10 +22,16 @@ package com.sk89q.worldedit.world.snapshot;
 import com.fastasyncworldedit.core.math.LocalBlockVectorSet;
 import com.sk89q.worldedit.EditSession;
 import com.sk89q.worldedit.MaxChangedBlocksException;
+import com.sk89q.worldedit.entity.BaseEntity;
 import com.sk89q.worldedit.math.BlockVector2;
 import com.sk89q.worldedit.math.BlockVector3;
 import com.sk89q.worldedit.regions.CuboidRegion;
 import com.sk89q.worldedit.regions.Region;
+import com.sk89q.worldedit.util.Location;
+import com.sk89q.worldedit.util.nbt.BinaryTag;
+import com.sk89q.worldedit.util.nbt.BinaryTagTypes;
+import com.sk89q.worldedit.util.nbt.CompoundBinaryTag;
+import com.sk89q.worldedit.util.nbt.ListBinaryTag;
 import com.sk89q.worldedit.world.DataException;
 import com.sk89q.worldedit.world.chunk.Chunk;
 import com.sk89q.worldedit.world.storage.ChunkStore;
@@ -48,6 +54,10 @@ public class SnapshotRestore {
     //FAWE end
     private final ChunkStore chunkStore;
     private final EditSession editSession;
+    //FAWE start - biome and entity restore
+    private final boolean restoreBiomes;
+    private final boolean restoreEntities;
+    //FAWE end
     private ArrayList<BlockVector2> missingChunks;
     private ArrayList<BlockVector2> errorChunks;
     private String lastErrorMessage;
@@ -60,8 +70,30 @@ public class SnapshotRestore {
      * @param region      The {@link Region} to restore to
      */
     public SnapshotRestore(ChunkStore chunkStore, EditSession editSession, Region region) {
+        //FAWE start - biome and entity restore
+        this(chunkStore, editSession, region, false, false);
+    }
+
+    /**
+     * Construct the snapshot restore operation.
+     *
+     * @param chunkStore      The {@link com.sk89q.worldedit.world.snapshot.experimental.Snapshot} to restore from
+     * @param editSession     The {@link EditSession} to restore to
+     * @param region          The {@link Region} to restore to
+     * @param restoreBiomes   If biomes should be restored
+     * @param restoreEntities If entities should be restored
+     */
+    public SnapshotRestore(
+            ChunkStore chunkStore,
+            EditSession editSession,
+            Region region,
+            boolean restoreBiomes,
+            boolean restoreEntities
+    ) {
         this.chunkStore = chunkStore;
         this.editSession = editSession;
+        this.restoreBiomes = restoreBiomes;
+        this.restoreEntities = restoreEntities;
 
         if (region instanceof CuboidRegion) {
             findNeededCuboidChunks(region);
@@ -69,6 +101,7 @@ public class SnapshotRestore {
             findNeededChunks(region);
         }
     }
+    //FAWE end
 
     /**
      * Find needed chunks in the axis-aligned bounding box of the region.
@@ -151,10 +184,35 @@ public class SnapshotRestore {
                 for (BlockVector3 pos : entry.getValue()) {
                     try {
                         editSession.setBlock(pos, chunk.getBlock(pos));
+                        //FAWE start - biome and entity restore
+                        if (restoreBiomes && (pos.getX() & 3) == 0 && (pos.getY() & 3) == 0 && (pos.getZ() & 3) == 0) {
+                            editSession.setBiome(pos, chunk.getBiome(pos));
+                        }
+                        //FAWE end
                     } catch (DataException e) {
                         // this is a workaround: just ignore for now
                     }
                 }
+                //FAWE start - biome and entity restore
+                if (restoreEntities) {
+                    try {
+                        for (BaseEntity entity : chunk.getEntities()) {
+                            CompoundBinaryTag tag = entity.getNbtReference().getValue();
+                            ListBinaryTag pos = tag.getList("Pos");
+                            ListBinaryTag rotation = tag.getList("Rotation");
+                            double x = pos.getDouble(0);
+                            double y = pos.getDouble(1);
+                            double z = pos.getDouble(2);
+                            float yRot = rotation.getFloat(0);
+                            float xRot = rotation.getFloat(1);
+                            Location location = new Location(editSession.getWorld(), x, y, z, yRot, xRot);
+                            editSession.createEntity(location, entity);
+                        }
+                    } catch (DataException e) {
+                        // this is a workaround: just ignore for now
+                    }
+                }
+                //FAWE end
             } catch (MissingChunkException me) {
                 missingChunks.add(chunkPos);
             } catch (IOException | DataException me) {

worldedit-core/src/main/java/com/sk89q/worldedit/world/snapshot/experimental/SnapshotRestore.java

@@ -21,10 +21,15 @@ package com.sk89q.worldedit.world.snapshot.experimental;
 
 import com.sk89q.worldedit.EditSession;
 import com.sk89q.worldedit.MaxChangedBlocksException;
+import com.sk89q.worldedit.entity.BaseEntity;
 import com.sk89q.worldedit.math.BlockVector2;
 import com.sk89q.worldedit.math.BlockVector3;
 import com.sk89q.worldedit.regions.CuboidRegion;
 import com.sk89q.worldedit.regions.Region;
+import com.sk89q.worldedit.util.Location;
+import com.sk89q.worldedit.util.nbt.BinaryTagTypes;
+import com.sk89q.worldedit.util.nbt.CompoundBinaryTag;
+import com.sk89q.worldedit.util.nbt.ListBinaryTag;
 import com.sk89q.worldedit.world.DataException;
 import com.sk89q.worldedit.world.chunk.Chunk;
 import com.sk89q.worldedit.world.storage.ChunkStore;
@@ -44,6 +49,10 @@ public class SnapshotRestore {
     private final Map<BlockVector2, ArrayList<BlockVector3>> neededChunks = new LinkedHashMap<>();
     private final Snapshot snapshot;
     private final EditSession editSession;
+    //FAWE start - biome and entity restore
+    private final boolean restoreBiomes;
+    private final boolean restoreEntities;
+    //FAWE end
     private ArrayList<BlockVector2> missingChunks;
     private ArrayList<BlockVector2> errorChunks;
     private String lastErrorMessage;
@@ -56,8 +65,30 @@ public class SnapshotRestore {
      * @param region      The {@link Region} to restore to
      */
     public SnapshotRestore(Snapshot snapshot, EditSession editSession, Region region) {
+        //FAWE start - biome and entity restore
+        this(snapshot, editSession, region, false, false);
+    }
+
+    /**
+     * Construct the snapshot restore operation.
+     *
+     * @param snapshot        The {@link Snapshot} to restore from
+     * @param editSession     The {@link EditSession} to restore to
+     * @param region          The {@link Region} to restore to
+     * @param restoreBiomes   If biomes should be restored
+     * @param restoreEntities If entities should be restored
+     */
+    public SnapshotRestore(
+            Snapshot snapshot,
+            EditSession editSession,
+            Region region,
+            boolean restoreBiomes,
+            boolean restoreEntities
+    ) {
         this.snapshot = snapshot;
         this.editSession = editSession;
+        this.restoreBiomes = restoreBiomes;
+        this.restoreEntities = restoreEntities;
 
         if (region instanceof CuboidRegion) {
             findNeededCuboidChunks(region);
@@ -65,6 +96,7 @@ public class SnapshotRestore {
             findNeededChunks(region);
         }
     }
+    //FAWE end
 
     /**
      * Find needed chunks in the axis-aligned bounding box of the region.
@@ -148,10 +180,35 @@ public class SnapshotRestore {
                 for (BlockVector3 pos : entry.getValue()) {
                     try {
                         editSession.setBlock(pos, chunk.getBlock(pos));
+                        //FAWE start - biome and entity restore
+                        if (restoreBiomes && (pos.getX() & 3) == 0 && (pos.getY() & 3) == 0 && (pos.getZ() & 3) == 0) {
+                            editSession.setBiome(pos, chunk.getBiome(pos));
+                        }
+                        //FAWE end
                     } catch (DataException e) {
                         // this is a workaround: just ignore for now
                     }
                 }
+                //FAWE start - biome and entity restore
+                if (restoreEntities) {
+                    try {
+                        for (BaseEntity entity : chunk.getEntities()) {
+                            CompoundBinaryTag tag = entity.getNbtReference().getValue();
+                            ListBinaryTag pos = tag.getList("Pos", BinaryTagTypes.LIST);
+                            ListBinaryTag rotation = tag.getList("Rotation", BinaryTagTypes.LIST);
+                            double x = pos.getDouble(0);
+                            double y = pos.getDouble(1);
+                            double z = pos.getDouble(2);
+                            float yRot = rotation.getFloat(0);
+                            float xRot = rotation.getFloat(1);
+                            Location location = new Location(editSession.getWorld(), x, y, z, yRot, xRot);
+                            editSession.createEntity(location, entity);
+                        }
+                    } catch (DataException e) {
+                        // this is a workaround: just ignore for now
+                    }
+                }
+                //FAWE end
             } catch (MissingChunkException me) {
                 missingChunks.add(chunkPos);
             } catch (IOException | DataException me) {

worldedit-core/src/main/java/com/sk89q/worldedit/world/storage/ChunkStore.java

@@ -20,12 +20,14 @@
 package com.sk89q.worldedit.world.storage;
 
 import com.sk89q.jnbt.CompoundTag;
+import com.sk89q.worldedit.internal.Constants;
 import com.sk89q.worldedit.math.BlockVector2;
 import com.sk89q.worldedit.math.BlockVector3;
 import com.sk89q.worldedit.world.DataException;
 import com.sk89q.worldedit.world.World;
 import com.sk89q.worldedit.world.chunk.Chunk;
 
+import javax.annotation.Nullable;
 import java.io.Closeable;
 import java.io.IOException;
 
@@ -84,6 +86,20 @@ public abstract class ChunkStore implements Closeable {
      */
     public abstract CompoundTag getChunkTag(BlockVector2 position, World world) throws DataException, IOException;
 
+    //FAWE start - biome and entity restore
+    /**
+     * Get the tag for the entities stored in a chunk from the entities folder. 1.17+ use only.
+     * If an error occurs, returns null.
+     *
+     * @param position the position of the chunk
+     * @return tag
+     */
+    @Nullable
+    public CompoundTag getEntitiesTag(BlockVector2 position, World world) {
+        return null;
+    }
+    //FAWE end
+
     /**
      * Get a chunk at a location.
      *
@@ -95,7 +111,17 @@ public abstract class ChunkStore implements Closeable {
      */
     public Chunk getChunk(BlockVector2 position, World world) throws DataException, IOException {
         CompoundTag rootTag = getChunkTag(position, world);
-        return ChunkStoreHelper.getChunk(rootTag);
+        //FAWE start - biome and entity restore
+        int dataVersion = rootTag.getInt("DataVersion");
+        if (dataVersion == 0) {
+            dataVersion = -1;
+        }
+        if (dataVersion >= Constants.DATA_VERSION_MC_1_17) {
+            return ChunkStoreHelper.getChunk(rootTag, () -> getEntitiesTag(position, world));
+        } else {
+            return ChunkStoreHelper.getChunk(rootTag);
+        }
+        //FAWE end
     }
 
     @Override

worldedit-core/src/main/java/com/sk89q/worldedit/world/storage/ChunkStoreHelper.java

@@ -31,13 +31,16 @@ import com.sk89q.worldedit.world.DataException;
 import com.sk89q.worldedit.world.DataFixer;
 import com.sk89q.worldedit.world.chunk.AnvilChunk;
 import com.sk89q.worldedit.world.chunk.AnvilChunk13;
+import com.sk89q.worldedit.world.chunk.AnvilChunk15;
 import com.sk89q.worldedit.world.chunk.AnvilChunk16;
+import com.sk89q.worldedit.world.chunk.AnvilChunk17;
 import com.sk89q.worldedit.world.chunk.Chunk;
 import com.sk89q.worldedit.world.chunk.OldChunk;
 
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.Map;
+import java.util.function.Supplier;
 
 public class ChunkStoreHelper {
 
@@ -69,6 +72,21 @@ public class ChunkStoreHelper {
      * @throws DataException if the rootTag is not valid chunk data
      */
     public static Chunk getChunk(CompoundTag rootTag) throws DataException {
+        //FAWE start - biome and entity restore
+        return getChunk(rootTag, () -> null);
+    }
+
+    /**
+     * Convert a chunk NBT tag into a {@link Chunk} implementation.
+     *
+     * @param rootTag     the root tag of the chunk
+     * @param entitiesTag supplier to provide entities tag. Only required for 1.17+ where entities are stored in a separate
+     *                    location
+     * @return a Chunk implementation
+     * @throws DataException if the rootTag is not valid chunk data
+     */
+    public static Chunk getChunk(CompoundTag rootTag, Supplier<CompoundTag> entitiesTag) throws DataException {
+        //FAWE end
         Map<String, Tag> children = rootTag.getValue();
         CompoundTag tag = null;
 
@@ -112,9 +130,19 @@ public class ChunkStoreHelper {
                 dataVersion = currentDataVersion;
             }
         }
+        //FAWE start - biome and entity restore
+        if (dataVersion >= Constants.DATA_VERSION_MC_1_17) {
+            return new AnvilChunk17(tag, entitiesTag);
+        }
+        //FAWE end
         if (dataVersion >= Constants.DATA_VERSION_MC_1_16) {
             return new AnvilChunk16(tag);
         }
+        //FAWE start - biome and entity restore
+        if (dataVersion >= Constants.DATA_VERSION_MC_1_15) {
+            return new AnvilChunk15(tag);
+        }
+        //FAWE end
         if (dataVersion >= Constants.DATA_VERSION_MC_1_13) {
             return new AnvilChunk13(tag);
         }

worldedit-core/src/main/java/com/sk89q/worldedit/world/storage/FileMcRegionChunkStore.java

@@ -21,6 +21,7 @@ package com.sk89q.worldedit.world.storage;
 
 import com.sk89q.worldedit.world.DataException;
 
+import javax.annotation.Nullable;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.FileNotFoundException;
@@ -42,11 +43,15 @@ public class FileMcRegionChunkStore extends McRegionChunkStore {
         this.path = path;
     }
 
+    //FAWE start - biome and entity restore
     @Override
-    protected InputStream getInputStream(String name, String world) throws IOException, DataException {
+    protected InputStream getInputStream(String name, String world, @Nullable String folderOverride) throws IOException,
+            DataException {
         Pattern ext = Pattern.compile(".*\\.mc[ra]$"); // allow either file extension, both work the same
         File file = null;
-        File[] files = new File(path, "region").listFiles();
+        String folder = folderOverride != null && !folderOverride.isEmpty() ? folderOverride : "region";
+        File[] files = new File(path, folder).listFiles();
+        //FAWE end
 
         if (files == null) {
             throw new FileNotFoundException();
@@ -56,7 +61,9 @@ public class FileMcRegionChunkStore extends McRegionChunkStore {
             String tempName = f.getName().replaceFirst("mcr$", "mca"); // matcher only does one at a time
             if (ext.matcher(f.getName()).matches() && name.equalsIgnoreCase(tempName)) {
                 // get full original path now
-                file = new File(path + File.separator + "region" + File.separator + f.getName());
+                //FAWE start - biome and entity restore
+                file = new File(path + File.separator + folder + File.separator + f.getName());
+                //FAWE end
                 break;
             }
         }

worldedit-core/src/main/java/com/sk89q/worldedit/world/storage/McRegionChunkStore.java

@@ -24,6 +24,7 @@ import com.sk89q.worldedit.math.BlockVector2;
 import com.sk89q.worldedit.world.DataException;
 import com.sk89q.worldedit.world.World;
 
+import javax.annotation.Nullable;
 import java.io.IOException;
 import java.io.InputStream;
 
@@ -45,7 +46,10 @@ public abstract class McRegionChunkStore extends ChunkStore {
         return "r." + (x >> 5) + "." + (z >> 5) + ".mca";
     }
 
-    protected McRegionReader getReader(BlockVector2 pos, String worldname) throws DataException, IOException {
+    //FAWE start - biome and entity restore
+    protected McRegionReader getReader(BlockVector2 pos, String worldname, @Nullable String folderOverride) throws DataException,
+            IOException {
+        //FAWE end
         String filename = getFilename(pos);
         if (curFilename != null) {
             if (curFilename.equals(filename)) {
@@ -57,7 +61,9 @@ public abstract class McRegionChunkStore extends ChunkStore {
                 }
             }
         }
-        InputStream stream = getInputStream(filename, worldname);
+        //FAWE start - biome and entity restore
+        InputStream stream = getInputStream(filename, worldname, folderOverride);
+        //FAWE end
         cachedReader = new McRegionReader(stream);
         //curFilename = filename;
         return cachedReader;
@@ -66,21 +72,40 @@ public abstract class McRegionChunkStore extends ChunkStore {
     @Override
     public CompoundTag getChunkTag(BlockVector2 position, World world) throws DataException, IOException {
         return ChunkStoreHelper.readCompoundTag(() -> {
-            McRegionReader reader = getReader(position, world.getName());
+            McRegionReader reader = getReader(position, world.getName(), null);
 
             return reader.getChunkInputStream(position);
         });
     }
 
+    //FAWE start - biome and entity restore
+    @Override
+    public CompoundTag getEntitiesTag(BlockVector2 position, World world) {
+        try {
+            return ChunkStoreHelper.readCompoundTag(() -> {
+                McRegionReader reader = getReader(position, world.getName(), "entities");
+
+                return reader.getChunkInputStream(position);
+            });
+        } catch (DataException | IOException e) {
+            return null;
+        }
+    }
+    //FAWE end
+
     /**
      * Get the input stream for a chunk file.
      *
-     * @param name      the name of the chunk file
+     * @param name           the name of the chunk file
-     * @param worldName the world name
+     * @param worldName      the world name
+     * @param folderOverride override folder to check. "entities" used for getting entities in 1.17+
      * @return an input stream
      * @throws IOException if there is an error getting the chunk data
      */
-    protected abstract InputStream getInputStream(String name, String worldName) throws IOException, DataException;
+    //FAWE start - biome and entity restore
+    protected abstract InputStream getInputStream(String name, String worldName, @Nullable String folderOverride) throws
+            IOException, DataException;
+    //FAWE end
 
     @Override
     public void close() throws IOException {

worldedit-core/src/main/java/com/sk89q/worldedit/world/storage/TrueZipMcRegionChunkStore.java

@@ -25,6 +25,7 @@ import com.sk89q.worldedit.world.DataException;
 import de.schlichtherle.util.zip.ZipEntry;
 import de.schlichtherle.util.zip.ZipFile;
 
+import javax.annotation.Nullable;
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
@@ -73,20 +74,18 @@ public class TrueZipMcRegionChunkStore extends McRegionChunkStore {
         zip = new ZipFile(zipFile);
     }
 
-    /**
-     * Get the input stream for a chunk file.
-     *
-     * @param name      the name
-     * @param worldName the world name
-     * @return an input stream
-     * @throws IOException   if there is an error getting the chunk data
-     * @throws DataException if there is an error getting the chunk data
-     */
     @Override
     @SuppressWarnings("unchecked")
-    protected InputStream getInputStream(String name, String worldName) throws IOException, DataException {
+    //FAWE start - biome and entity restore
+    protected InputStream getInputStream(String name, String worldName, @Nullable String folderOverride) throws IOException,
+            DataException {
         // Detect subfolder for the world's files
-        if (folder != null) {
+        if (folderOverride != null) {
+            if (!folderOverride.isEmpty()) {
+                name = folderOverride + "/" + name;
+            }
+        } else if (folder != null) {
+            //FAWE end
             if (!folder.isEmpty()) {
                 name = folder + "/" + name;
             }

worldedit-core/src/main/java/com/sk89q/worldedit/world/storage/ZippedMcRegionChunkStore.java

@@ -23,6 +23,7 @@ package com.sk89q.worldedit.world.storage;
 
 import com.sk89q.worldedit.world.DataException;
 
+import javax.annotation.Nullable;
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
@@ -73,10 +74,17 @@ public class ZippedMcRegionChunkStore extends McRegionChunkStore {
     }
 
     @Override
-    protected InputStream getInputStream(String name, String worldName) throws IOException, DataException {
+    //FAWE start - biome and entity restore
+    protected InputStream getInputStream(String name, String worldName, @Nullable String folderOverride) throws IOException,
+            DataException {
         // Detect subfolder for the world's files
-        if (folder != null) {
+        if (folderOverride != null) {
+            if (!folderOverride.isEmpty()) {
+                name = folderOverride + "/" + name;
+            }
+        } else if (folder != null) {
             if (!folder.isEmpty()) {
+                //FAWE end
                 name = folder + "/" + name;
             }
         } else {
@@ -93,7 +101,9 @@ public class ZippedMcRegionChunkStore extends McRegionChunkStore {
                             endIndex = entryName.lastIndexOf('\\');
                         }
                         folder = entryName.substring(0, endIndex);
-                        if (folder.endsWith("poi")) {
+                        //FAWE start - biome and entity restore
+                        if (folder.endsWith("poi") || folder.endsWith("entities")) {
+                            //FAWE end
                             continue;
                         }
                         name = folder + "/" + name;