Improve block getter performance

src/main/java/net/minestom/server/instance/Chunk.java

@@ -20,7 +20,6 @@ import net.minestom.server.network.player.PlayerConnection;
 import net.minestom.server.tag.Tag;
 import net.minestom.server.tag.TagHandler;
 import net.minestom.server.utils.ArrayUtils;
-import net.minestom.server.utils.Position;
 import net.minestom.server.utils.chunk.ChunkSupplier;
 import net.minestom.server.world.biomes.Biome;
 import org.jetbrains.annotations.NotNull;
@@ -119,14 +118,6 @@ public abstract class Chunk implements BlockGetter, BlockSetter, Viewable, Ticka
     @Override
     public abstract void tick(long time);
 
-    /**
-     * Gets all the block entities in this chunk.
-     *
-     * @return the block entities in this chunk
-     */
-    @NotNull
-    public abstract Set<Integer> getBlockEntities();
-
     /**
      * Gets the last time that this chunk changed.
      * <p>

src/main/java/net/minestom/server/instance/DynamicChunk.java

@@ -1,6 +1,8 @@
 package net.minestom.server.instance;
 
 import com.extollit.gaming.ai.path.model.ColumnarOcclusionFieldList;
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
 import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap;
 import net.minestom.server.coordinate.Vec;
 import net.minestom.server.entity.pathfinding.PFBlock;
@@ -9,14 +11,15 @@ import net.minestom.server.instance.block.BlockHandler;
 import net.minestom.server.network.packet.server.play.ChunkDataPacket;
 import net.minestom.server.utils.chunk.ChunkUtils;
 import net.minestom.server.world.biomes.Biome;
+import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jglrxavpok.hephaistos.nbt.NBTCompound;
 
 import java.lang.ref.SoftReference;
 import java.util.Map;
-import java.util.Set;
 import java.util.TreeMap;
+import java.util.concurrent.TimeUnit;
 
 /**
  * Represents a {@link Chunk} which store each individual block in memory.
@@ -28,8 +31,7 @@ public class DynamicChunk extends Chunk {
     protected final TreeMap<Integer, Section> sectionMap = new TreeMap<>();
 
     // Key = ChunkUtils#getBlockIndex
-    protected final Int2ObjectOpenHashMap<BlockHandler> handlerMap = new Int2ObjectOpenHashMap<>();
+    protected final Int2ObjectOpenHashMap<BlockEntry> entries = new Int2ObjectOpenHashMap<>();
-    protected final Int2ObjectOpenHashMap<NBTCompound> nbtMap = new Int2ObjectOpenHashMap<>();
     protected final Int2ObjectOpenHashMap<BlockHandler> tickableMap = new Int2ObjectOpenHashMap<>();
 
     private long lastChangeTime;
@@ -56,17 +58,11 @@ public class DynamicChunk extends Chunk {
         final int index = ChunkUtils.getBlockIndex(x, y, z);
         // Handler
         final BlockHandler handler = block.handler();
-        if (handler != null) {
-            this.handlerMap.put(index, handler);
-        } else {
-            this.handlerMap.remove(index);
-        }
-        // Nbt
         final NBTCompound nbt = block.nbt();
-        if (nbt != null) {
+        if (handler != null || nbt != null) {
-            this.nbtMap.put(index, nbt);
+            this.entries.put(index, new BlockEntry(handler, nbt));
         } else {
-            this.nbtMap.remove(index);
+            this.entries.remove(index);
         }
         // Tickable
         if (handler != null && handler.isTickable()) {
@@ -115,23 +111,20 @@ public class DynamicChunk extends Chunk {
             return Block.AIR;
         }
         final int index = ChunkUtils.getBlockIndex(x, y, z);
-        final BlockHandler handler = handlerMap.get(index);
+        final var entry = entries.get(index);
-        final NBTCompound nbt = nbtMap.get(index);
+        if (entry != null) {
-        if (handler != null) {
+            final BlockHandler handler = entry.handler;
-            block = block.withHandler(handler);
+            final NBTCompound nbt = entry.nbtCompound;
-        }
+            if (handler != null) {
-        if (nbt != null) {
+                block = block.withHandler(handler);
-            block = block.withNbt(nbt);
+            }
+            if (nbt != null) {
+                block = block.withNbt(nbt);
+            }
         }
         return block;
     }
 
-    @NotNull
-    @Override
-    public Set<Integer> getBlockEntities() {
-        return nbtMap.keySet();
-    }
-
     @Override
     public long getLastChangeTime() {
         return lastChangeTime;
@@ -149,8 +142,7 @@ public class DynamicChunk extends Chunk {
         packet.chunkX = chunkX;
         packet.chunkZ = chunkZ;
         packet.sections = (Map<Integer, Section>) sectionMap.clone(); // TODO deep clone
-        packet.handlerMap = handlerMap.clone();
+        packet.entries = entries.clone();
-        packet.nbtMap = nbtMap.clone();
 
         this.cachedPacketTime = getLastChangeTime();
         this.cachedPacket = new SoftReference<>(packet);
@@ -164,21 +156,42 @@ public class DynamicChunk extends Chunk {
         for (var entry : sectionMap.entrySet()) {
             dynamicChunk.sectionMap.put(entry.getKey(), entry.getValue().clone());
         }
-        dynamicChunk.handlerMap.putAll(handlerMap);
+        dynamicChunk.entries.putAll(entries);
-        dynamicChunk.nbtMap.putAll(nbtMap);
-
         return dynamicChunk;
     }
 
     @Override
     public void reset() {
         this.sectionMap.values().forEach(Section::clear);
-        this.handlerMap.clear();
+        this.entries.clear();
-        this.nbtMap.clear();
     }
 
     private @NotNull Section retrieveSection(int y) {
         final int sectionIndex = ChunkUtils.getSectionAt(y);
         return getSection(sectionIndex);
     }
+
+    @ApiStatus.Internal
+    public static class BlockEntry {
+        private static final Cache<NBTCompound, NBTCompound> NBT_CACHE = Caffeine.newBuilder()
+                .expireAfterWrite(5, TimeUnit.MINUTES)
+                .weakValues()
+                .build();
+
+        private final BlockHandler handler;
+        private final NBTCompound nbtCompound;
+
+        public BlockEntry(BlockHandler handler, NBTCompound nbtCompound) {
+            this.handler = handler;
+            this.nbtCompound = nbtCompound != null ? NBT_CACHE.get(nbtCompound, compound -> nbtCompound) : null;
+        }
+
+        public BlockHandler handler() {
+            return handler;
+        }
+
+        public NBTCompound nbtCompound() {
+            return nbtCompound;
+        }
+    }
 }

src/main/java/net/minestom/server/instance/block/BlockImpl.java

@@ -1,7 +1,5 @@
 package net.minestom.server.instance.block;
 
-import com.github.benmanes.caffeine.cache.Cache;
-import com.github.benmanes.caffeine.cache.Caffeine;
 import net.minestom.server.registry.Registry;
 import net.minestom.server.tag.Tag;
 import org.jetbrains.annotations.NotNull;
@@ -12,15 +10,8 @@ import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
-import java.util.concurrent.TimeUnit;
 
 class BlockImpl implements Block {
-
-    private static final Cache<NBTCompound, NBTCompound> NBT_CACHE = Caffeine.newBuilder()
-            .expireAfterWrite(5, TimeUnit.MINUTES)
-            .weakValues()
-            .build();
-
     private final Registry.BlockEntry registry;
     private final Map<String, String> properties;
     private final NBTCompound nbt;
@@ -57,8 +48,7 @@ class BlockImpl implements Block {
     public @NotNull <T> Block withTag(@NotNull Tag<T> tag, @Nullable T value) {
         var compound = Objects.requireNonNullElseGet(nbt(), NBTCompound::new);
         tag.write(compound, value);
-        final var cachedNbt = NBT_CACHE.get(compound, c -> compound);
+        return new BlockImpl(registry, properties, compound, handler);
-        return new BlockImpl(registry, properties, cachedNbt, handler);
     }
 
     @Override

src/main/java/net/minestom/server/network/packet/server/play/ChunkDataPacket.java

@@ -5,6 +5,7 @@ import io.netty.buffer.Unpooled;
 import it.unimi.dsi.fastutil.ints.Int2LongRBTreeMap;
 import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap;
 import net.minestom.server.MinecraftServer;
+import net.minestom.server.instance.DynamicChunk;
 import net.minestom.server.instance.Section;
 import net.minestom.server.instance.block.BlockHandler;
 import net.minestom.server.network.packet.server.ServerPacket;
@@ -34,8 +35,7 @@ public class ChunkDataPacket implements ServerPacket, CacheablePacket {
     public int chunkX, chunkZ;
 
     public Map<Integer, Section> sections = new HashMap<>();
-    public Map<Integer, BlockHandler> handlerMap = new HashMap<>();
+    public Map<Integer, DynamicChunk.BlockEntry> entries = new HashMap<>();
-    public Map<Integer, NBTCompound> nbtMap = new HashMap<>();
 
     private static final byte CHUNK_SECTION_COUNT = 16;
     private static final int MAX_BITS_PER_ENTRY = 16;
@@ -124,17 +124,18 @@ public class ChunkDataPacket implements ServerPacket, CacheablePacket {
         blocks.release();
 
         // Block entities
-        if (handlerMap == null || handlerMap.isEmpty()) {
+        if (entries == null || entries.isEmpty()) {
             writer.writeVarInt(0);
         } else {
             List<NBTCompound> compounds = new ArrayList<>();
-            for (var entry : handlerMap.entrySet()) {
+            for (var entry : entries.entrySet()) {
                 final int index = entry.getKey();
-                final BlockHandler handler = entry.getValue();
+                final var blockEntry = entry.getValue();
+                final BlockHandler handler = blockEntry.handler();
                 final var blockEntityTags = handler.getBlockEntityTags();
                 if (blockEntityTags.isEmpty())
                     continue;
-                final var blockNbt = Objects.requireNonNullElseGet(nbtMap.get(index), NBTCompound::new);
+                final var blockNbt = Objects.requireNonNullElseGet(blockEntry.nbtCompound(), NBTCompound::new);
                 final var resultNbt = new NBTCompound();
 
                 for (Tag<?> tag : blockEntityTags) {
@@ -219,8 +220,8 @@ public class ChunkDataPacket implements ServerPacket, CacheablePacket {
 
             // Block entities
             final int blockEntityCount = reader.readVarInt();
-            handlerMap = new Int2ObjectOpenHashMap<>();
+
-            nbtMap = new Int2ObjectOpenHashMap<>();
+            entries = new Int2ObjectOpenHashMap<>();
             for (int i = 0; i < blockEntityCount; i++) {
                 NBTCompound tag = (NBTCompound) reader.readTag();
                 final String id = tag.getString("id");