Update SpongeSchematic format to version 2.

Allows saving and loading entities and biomes.

worldedit-bukkit/src/main/java/com/sk89q/worldedit/bukkit/BukkitServerInterface.java

@@ -65,6 +65,12 @@ public class BukkitServerInterface implements MultiUserPlatform {
         return BukkitRegistries.getInstance();
     }
 
+    @Override
+    public int getDataVersion() {
+        // TODO - add to adapter - CraftMagicNumbers#getDataVersion
+        return 1631;
+    }
+
     @Override
     public boolean isValidMobType(String type) {
         final EntityType entityType = EntityType.fromName(type);

worldedit-core/src/main/java/com/sk89q/jnbt/CompoundTagBuilder.java

@@ -181,6 +181,18 @@ public class CompoundTagBuilder {
         return put(key, new StringTag(value));
     }
 
+    /**
+     * Remove the given key from the compound tag. Does nothing if the key doesn't exist.
+     *
+     * @param key the key
+     * @return this object
+     */
+    public CompoundTagBuilder remove(String key) {
+        checkNotNull(key);
+        entries.remove(key);
+        return this;
+    }
+
     /**
      * Put all the entries from the given map into this map.
      *

worldedit-core/src/main/java/com/sk89q/worldedit/extension/platform/Platform.java

@@ -45,6 +45,13 @@ public interface Platform {
      */
     Registries getRegistries();
 
+    /**
+     * Gets the Minecraft data version being used by the platform.
+     *
+     * @return the data version
+     */
+    int getDataVersion();
+
     /**
      * Checks if a mob type is valid.
      *

worldedit-core/src/main/java/com/sk89q/worldedit/extent/clipboard/BlockArrayClipboard.java

@@ -171,7 +171,10 @@ public class BlockArrayClipboard implements Clipboard {
         if (biomes != null
                 && position.containedWithin(getMinimumPoint().toBlockVector2(), getMaximumPoint().toBlockVector2())) {
             BlockVector2 v = position.subtract(region.getMinimumPoint().toBlockVector2());
-            return biomes[v.getBlockX()][v.getBlockZ()];
+            BiomeType biomeType = biomes[v.getBlockX()][v.getBlockZ()];
+            if (biomeType != null) {
+                return biomeType;
+            }
         }
 
         return BiomeTypes.OCEAN;

worldedit-core/src/main/java/com/sk89q/worldedit/extent/clipboard/io/MCEditSchematicReader.java

@@ -53,6 +53,7 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -315,6 +316,7 @@ public class MCEditSchematicReader extends NBTSchematicReader {
             case "PigZombie": return "zombie_pigman";
             default: return id;
         }
+        return id;
     }
 
     private String convertBlockEntityId(String id) {

worldedit-core/src/main/java/com/sk89q/worldedit/extent/clipboard/io/SpongeSchematicReader.java

@@ -28,18 +28,27 @@ import com.sk89q.jnbt.ListTag;
 import com.sk89q.jnbt.NBTInputStream;
 import com.sk89q.jnbt.NamedTag;
 import com.sk89q.jnbt.ShortTag;
+import com.sk89q.jnbt.StringTag;
 import com.sk89q.jnbt.Tag;
 import com.sk89q.worldedit.WorldEdit;
 import com.sk89q.worldedit.WorldEditException;
+import com.sk89q.worldedit.entity.BaseEntity;
 import com.sk89q.worldedit.extension.input.InputParseException;
 import com.sk89q.worldedit.extension.input.ParserContext;
+import com.sk89q.worldedit.extension.platform.Capability;
 import com.sk89q.worldedit.extent.clipboard.BlockArrayClipboard;
 import com.sk89q.worldedit.extent.clipboard.Clipboard;
 import com.sk89q.worldedit.extent.clipboard.io.legacycompat.NBTCompatibilityHandler;
 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.world.biome.BiomeType;
+import com.sk89q.worldedit.world.biome.BiomeTypes;
 import com.sk89q.worldedit.world.block.BlockState;
+import com.sk89q.worldedit.world.entity.EntityType;
+import com.sk89q.worldedit.world.entity.EntityTypes;
+import com.sk89q.worldedit.world.storage.NBTConversions;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -89,6 +98,15 @@ public class SpongeSchematicReader extends NBTSchematicReader {
         int version = requireTag(schematic, "Version", IntTag.class).getValue();
         if (version == 1) {
             return readVersion1(schematicTag);
+        } else if (version == 2) {
+            int dataVersion = requireTag(schematic, "DataVersion", IntTag.class).getValue();
+            if (dataVersion > WorldEdit.getInstance().getPlatformManager()
+                    .queryCapability(Capability.WORLD_EDITING).getDataVersion()) {
+                // maybe should just warn? simple schematics may be compatible still.
+                throw new IOException("Schematic was made in a newer Minecraft version. Data may be incompatible.");
+            }
+            BlockArrayClipboard clip = readVersion1(schematicTag);
+            return readVersion2(clip, schematicTag);
         }
         throw new IOException("This schematic version is currently not supported");
     }
@@ -123,10 +141,11 @@ public class SpongeSchematicReader extends NBTSchematicReader {
             region = new CuboidRegion(origin, origin.add(width, height, length).subtract(BlockVector3.ONE));
         }
 
+        // Note: these aren't actually required by the spec, but we require them I guess?
         int paletteMax = requireTag(schematic, "PaletteMax", IntTag.class).getValue();
         Map<String, Tag> paletteObject = requireTag(schematic, "Palette", CompoundTag.class).getValue();
         if (paletteObject.size() != paletteMax) {
-            throw new IOException("Differing given palette size to actual size");
+            throw new IOException("Block palette size does not match expected size.");
         }
 
         Map<Integer, BlockState> palette = new HashMap<>();
@@ -168,8 +187,8 @@ public class SpongeSchematicReader extends NBTSchematicReader {
 
         int index = 0;
         int i = 0;
-        int value = 0;
+        int value;
-        int varintLength = 0;
+        int varintLength;
         while (i < blocks.length) {
             value = 0;
             varintLength = 0;
@@ -185,7 +204,7 @@ public class SpongeSchematicReader extends NBTSchematicReader {
                 }
                 i++;
             }
-            // index = (y * length + z) * width + x
+            // index = (y * length * width) + (z * width) + x
             int y = index / (width * length);
             int z = (index % (width * length)) / width;
             int x = (index % (width * length)) % width;
@@ -219,6 +238,99 @@ public class SpongeSchematicReader extends NBTSchematicReader {
         return clipboard;
     }
 
+    private Clipboard readVersion2(BlockArrayClipboard version1, CompoundTag schematicTag) throws IOException {
+        Map<String, Tag> schematic = schematicTag.getValue();
+        if (schematic.containsKey("BiomeData")) {
+            readBiomes(version1, schematic);
+        }
+        if (schematic.containsKey("Entities")) {
+            readEntities(version1, schematic);
+        }
+        return version1;
+    }
+
+    private void readBiomes(BlockArrayClipboard clipboard, Map<String, Tag> schematic) throws IOException {
+        ByteArrayTag dataTag = requireTag(schematic, "BiomeData", ByteArrayTag.class);
+        // TODO for now, we just assume if biomedata is present, palette will be as well.
+        // atm the spec doesn't actually require palettes
+        IntTag maxTag = requireTag(schematic, "BiomePaletteMax", IntTag.class);
+        CompoundTag paletteTag = requireTag(schematic, "BiomePalette", CompoundTag.class);
+
+        Map<Integer, BiomeType> palette = new HashMap<>();
+        if (maxTag.getValue() != paletteTag.getValue().size()) {
+            throw new IOException("Biome palette size does not match expected size.");
+        }
+        Map<String, Tag> paletteEntries = paletteTag.getValue();
+
+        for (Map.Entry<String, Tag> palettePart : paletteEntries.entrySet()) {
+            BiomeType biome = BiomeTypes.get(palettePart.getKey());
+            if (biome == null) {
+                log.warn("Unknown biome type '" + palettePart.getKey() + "' in palette. Are you missing a mod or using a schematic made in a newer version of Minecraft?");
+            }
+            Tag idTag = palettePart.getValue();
+            if (!(idTag instanceof IntTag)) {
+                throw new IOException("Biome mapped to non-Int tag.");
+            }
+            palette.put(((IntTag) idTag).getValue(), biome);
+        }
+
+        int width = clipboard.getDimensions().getX();
+
+        byte[] biomes = dataTag.getValue();
+        int biomeIndex = 0;
+        int biomeJ = 0;
+        int bVal;
+        int varIntLength;
+        while (biomeJ < biomes.length) {
+            bVal = 0;
+            varIntLength = 0;
+
+            while (true) {
+                bVal |= (biomes[biomeJ] & 127) << (varIntLength++ * 7);
+                if (varIntLength > 5) {
+                    throw new RuntimeException("VarInt too big (probably corrupted data)");
+                }
+                if (((biomes[biomeJ] & 128) != 128)) {
+                    biomeJ++;
+                    break;
+                }
+                biomeJ++;
+            }
+            int z = biomeIndex / width;
+            int x = biomeIndex % width;
+            BiomeType type = palette.get(bVal);
+            clipboard.setBiome(clipboard.getMinimumPoint().toBlockVector2().add(x, z), type);
+            biomeIndex++;
+        }
+    }
+
+    private void readEntities(BlockArrayClipboard clipboard, Map<String, Tag> schematic) throws IOException {
+        List<Tag> entList = requireTag(schematic, "Entities", ListTag.class).getValue();
+        if (entList.isEmpty()) {
+            return;
+        }
+        for (Tag et : entList) {
+            if (!(et instanceof CompoundTag)) {
+                continue;
+            }
+            CompoundTag entityTag = (CompoundTag) et;
+            Map<String, Tag> tags = entityTag.getValue();
+            String id = requireTag(tags, "Id", StringTag.class).getValue();
+
+            EntityType entityType = EntityTypes.get(id);
+            if (entityType != null) {
+                Location location = NBTConversions.toLocation(clipboard,
+                        requireTag(tags, "Pos", ListTag.class),
+                        requireTag(tags, "Rotation", ListTag.class));
+                BaseEntity state = new BaseEntity(entityType,
+                        entityTag.createBuilder().putString("id", id).remove("Id").build());
+                clipboard.createEntity(location, state);
+            } else {
+                log.warn("Unknown entity when pasting schematic: " + id);
+            }
+        }
+    }
+
     @Override
     public void close() throws IOException {
         inputStream.close();

worldedit-core/src/main/java/com/sk89q/worldedit/extent/clipboard/io/SpongeSchematicWriter.java

@@ -21,8 +21,12 @@ package com.sk89q.worldedit.extent.clipboard.io;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
+import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
 import com.sk89q.jnbt.ByteArrayTag;
 import com.sk89q.jnbt.CompoundTag;
+import com.sk89q.jnbt.DoubleTag;
+import com.sk89q.jnbt.FloatTag;
 import com.sk89q.jnbt.IntArrayTag;
 import com.sk89q.jnbt.IntTag;
 import com.sk89q.jnbt.ListTag;
@@ -30,9 +34,16 @@ import com.sk89q.jnbt.NBTOutputStream;
 import com.sk89q.jnbt.ShortTag;
 import com.sk89q.jnbt.StringTag;
 import com.sk89q.jnbt.Tag;
+import com.sk89q.worldedit.WorldEdit;
+import com.sk89q.worldedit.entity.BaseEntity;
+import com.sk89q.worldedit.extension.platform.Capability;
 import com.sk89q.worldedit.extent.clipboard.Clipboard;
+import com.sk89q.worldedit.math.BlockVector2;
 import com.sk89q.worldedit.math.BlockVector3;
+import com.sk89q.worldedit.math.Vector3;
 import com.sk89q.worldedit.regions.Region;
+import com.sk89q.worldedit.util.Location;
+import com.sk89q.worldedit.world.biome.BiomeType;
 import com.sk89q.worldedit.world.block.BaseBlock;
 
 import java.io.ByteArrayOutputStream;
@@ -41,12 +52,15 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 /**
  * Writes schematic files using the Sponge schematic format.
  */
 public class SpongeSchematicWriter implements ClipboardWriter {
 
+    private static final int CURRENT_VERSION = 2;
+
     private static final int MAX_SIZE = Short.MAX_VALUE - Short.MIN_VALUE;
     private final NBTOutputStream outputStream;
 
@@ -63,17 +77,17 @@ public class SpongeSchematicWriter implements ClipboardWriter {
     @Override
     public void write(Clipboard clipboard) throws IOException {
         // For now always write the latest version. Maybe provide support for earlier if more appear.
-        outputStream.writeNamedTag("Schematic", new CompoundTag(write1(clipboard)));
+        outputStream.writeNamedTag("Schematic", new CompoundTag(write2(clipboard)));
     }
 
     /**
-     * Writes a version 1 schematic file.
+     * Writes a version 2 schematic file.
      *
      * @param clipboard The clipboard
      * @return The schematic map
      * @throws IOException If an error occurs
      */
-    private Map<String, Tag> write1(Clipboard clipboard) throws IOException {
+    private Map<String, Tag> write2(Clipboard clipboard) throws IOException {
         Region region = clipboard.getRegion();
         BlockVector3 origin = clipboard.getOrigin();
         BlockVector3 min = region.getMinimumPoint();
@@ -93,7 +107,9 @@ public class SpongeSchematicWriter implements ClipboardWriter {
         }
 
         Map<String, Tag> schematic = new HashMap<>();
-        schematic.put("Version", new IntTag(1));
+        schematic.put("Version", new IntTag(CURRENT_VERSION));
+        schematic.put("DataVersion", new IntTag(
+                WorldEdit.getInstance().getPlatformManager().queryCapability(Capability.WORLD_EDITING).getDataVersion()));
 
         Map<String, Tag> metadata = new HashMap<>();
         metadata.put("WEOffsetX", new IntTag(offset.getBlockX()));
@@ -129,10 +145,7 @@ public class SpongeSchematicWriter implements ClipboardWriter {
                     BlockVector3 point = BlockVector3.at(x0, y0, z0);
                     BaseBlock block = clipboard.getFullBlock(point);
                     if (block.getNbtData() != null) {
-                        Map<String, Tag> values = new HashMap<>();
+                        Map<String, Tag> values = new HashMap<>(block.getNbtData().getValue());
-                        for (Map.Entry<String, Tag> entry : block.getNbtData().getValue().entrySet()) {
-                            values.put(entry.getKey(), entry.getValue());
-                        }
 
                         values.remove("id"); // Remove 'id' if it exists. We want 'Id'
 
@@ -179,9 +192,101 @@ public class SpongeSchematicWriter implements ClipboardWriter {
         schematic.put("BlockData", new ByteArrayTag(buffer.toByteArray()));
         schematic.put("TileEntities", new ListTag(CompoundTag.class, tileEntities));
 
+        // version 2 stuff
+        if (clipboard.hasBiomes()) {
+            writeBiomes(clipboard, schematic);
+        }
+
+        if (!clipboard.getEntities().isEmpty()) {
+            writeEntities(clipboard, schematic);
+        }
+
         return schematic;
     }
 
+    private void writeBiomes(Clipboard clipboard, Map<String, Tag> schematic) {
+        BlockVector3 min = clipboard.getMinimumPoint();
+        int width = clipboard.getRegion().getWidth();
+        int length = clipboard.getRegion().getLength();
+
+        ByteArrayOutputStream buffer = new ByteArrayOutputStream(width * length);
+
+        int paletteMax = 0;
+        Map<String, Integer> palette = new HashMap<>();
+
+        for (int z = 0; z < length; z++) {
+            int z0 = min.getBlockZ() + z;
+            for (int x = 0; x < width; x++) {
+                int x0 = min.getBlockX() + x;
+                BlockVector2 pt = BlockVector2.at(x0, z0);
+                BiomeType biome = clipboard.getBiome(pt);
+
+                String biomeKey = biome.getId();
+                int biomeId;
+                if (palette.containsKey(biomeKey)) {
+                    biomeId = palette.get(biomeKey);
+                } else {
+                    biomeId = paletteMax;
+                    palette.put(biomeKey, biomeId);
+                    paletteMax++;
+                }
+
+                while ((biomeId & -128) != 0) {
+                    buffer.write(biomeId & 127 | 128);
+                    biomeId >>>= 7;
+                }
+                buffer.write(biomeId);
+            }
+        }
+
+        schematic.put("BiomePaletteMax", new IntTag(paletteMax));
+
+        Map<String, Tag> paletteTag = new HashMap<>();
+        palette.forEach((key, value) -> paletteTag.put(key, new IntTag(value)));
+
+        schematic.put("BiomePalette", new CompoundTag(paletteTag));
+        schematic.put("BiomeData", new ByteArrayTag(buffer.toByteArray()));
+    }
+
+    private void writeEntities(Clipboard clipboard, Map<String, Tag> schematic) {
+        List<CompoundTag> entities = clipboard.getEntities().stream().map(e -> {
+            BaseEntity state = e.getState();
+            if (state == null) {
+                return null;
+            }
+            Map<String, Tag> values = Maps.newHashMap();
+            CompoundTag rawData = state.getNbtData();
+            if (rawData != null) {
+                values.putAll(rawData.getValue());
+            }
+            values.remove("id");
+            values.put("Id", new StringTag(state.getType().getId()));
+            values.put("Pos", writeVector(e.getLocation().toVector()));
+            values.put("Rotation", writeRotation(e.getLocation()));
+
+            return new CompoundTag(values);
+        }).filter(e -> e != null).collect(Collectors.toList());
+        if (entities.isEmpty()) {
+            return;
+        }
+        schematic.put("Entities", new ListTag(CompoundTag.class, entities));
+    }
+
+    private Tag writeVector(Vector3 vector) {
+        List<DoubleTag> list = new ArrayList<DoubleTag>();
+        list.add(new DoubleTag(vector.getX()));
+        list.add(new DoubleTag(vector.getY()));
+        list.add(new DoubleTag(vector.getZ()));
+        return new ListTag(DoubleTag.class, list);
+    }
+
+    private Tag writeRotation(Location location) {
+        List<FloatTag> list = new ArrayList<FloatTag>();
+        list.add(new FloatTag(location.getYaw()));
+        list.add(new FloatTag(location.getPitch()));
+        return new ListTag(FloatTag.class, list);
+    }
+
     @Override
     public void close() throws IOException {
         outputStream.close();

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

@@ -52,7 +52,7 @@ public final class NBTConversions {
         return new Location(
                 extent,
                 positionTag.asDouble(0), positionTag.asDouble(1), positionTag.asDouble(2),
-                (float) directionTag.asDouble(0), (float) directionTag.asDouble(1));
+                directionTag.getFloat(0), directionTag.getFloat(1));
     }
 
 }

worldedit-forge/src/main/java/com/sk89q/worldedit/forge/ForgePlatform.java

@@ -64,6 +64,12 @@ class ForgePlatform extends AbstractPlatform implements MultiUserPlatform {
         return ForgeRegistries.getInstance();
     }
 
+    @Override
+    public int getDataVersion() {
+        // TODO technically available as WorldInfo#field_209227_p but requires a world ref?
+        return 1631;
+    }
+
     @Override
     public boolean isValidMobType(String type) {
         return net.minecraftforge.registries.ForgeRegistries.ENTITIES.containsKey(new ResourceLocation(type));

worldedit-sponge/src/main/java/com/sk89q/worldedit/sponge/SpongePlatform.java

@@ -68,6 +68,12 @@ class SpongePlatform extends AbstractPlatform implements MultiUserPlatform {
         return SpongeRegistries.getInstance();
     }
 
+    @Override
+    public int getDataVersion() {
+        // TODO add to adapter - org.spongepowered.common.data.util.DataUtil#MINECRAFT_DATA_VERSION
+        return 1631;
+    }
+
     @Override
     public boolean isValidMobType(String type) {
         return Sponge.getRegistry().getType(EntityType.class, type).isPresent();