From feb4783988a84c63d6159eb4bc9d1a1bd9a61c6d Mon Sep 17 00:00:00 2001
From: themode <themode@outlook.fr>
Date: Sat, 26 Dec 2020 18:52:44 +0100
Subject: [PATCH] Small optimization for entity velocity

---
 .../server/collision/CollisionUtils.java       |  1 +
 .../net/minestom/server/entity/Entity.java     | 18 ++++++++++++------
 .../net/minestom/server/instance/Chunk.java    |  2 ++
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/src/main/java/net/minestom/server/collision/CollisionUtils.java b/src/main/java/net/minestom/server/collision/CollisionUtils.java
index e8428564a..4086de592 100644
--- a/src/main/java/net/minestom/server/collision/CollisionUtils.java
+++ b/src/main/java/net/minestom/server/collision/CollisionUtils.java
@@ -182,6 +182,7 @@ public class CollisionUtils {
      * @param newPosition     the future target position
      * @return the position with the world border collision applied (can be {@code newPosition} if not changed)
      */
+    @NotNull
     public static Position applyWorldBorder(@NotNull Instance instance,
                                             @NotNull Position currentPosition, @NotNull Position newPosition) {
         final WorldBorder worldBorder = instance.getWorldBorder();
diff --git a/src/main/java/net/minestom/server/entity/Entity.java b/src/main/java/net/minestom/server/entity/Entity.java
index 00dc15221..df4b85fc6 100644
--- a/src/main/java/net/minestom/server/entity/Entity.java
+++ b/src/main/java/net/minestom/server/entity/Entity.java
@@ -478,14 +478,17 @@ public abstract class Entity implements Viewable, EventHandler, DataContainer, P
                 // Stop here if the position is the same
                 final boolean updatePosition = !newPosition.isSimilar(position);
 
-                // Check chunk
-                if (!ChunkUtils.isLoaded(instance, newPosition.getX(), newPosition.getZ())) {
+                // World border collision
+                final Position finalVelocityPosition = CollisionUtils.applyWorldBorder(instance, position, newPosition);
+                final Chunk finalChunk = instance.getChunkAt(finalVelocityPosition);
+
+                // Entity shouldn't be updated when moving in an unloaded chunk
+                if (!ChunkUtils.isLoaded(finalChunk)) {
                     return;
                 }
 
-                // World border and apply the position
-                final Position finalVelocityPosition = CollisionUtils.applyWorldBorder(instance, position, newPosition);
-                if (finalVelocityPosition != null && updatePosition) {
+                // Apply the position
+                if (updatePosition) {
                     refreshPosition(finalVelocityPosition);
                 }
 
@@ -498,7 +501,10 @@ public abstract class Entity implements Viewable, EventHandler, DataContainer, P
                     float drag;
                     if (onGround) {
                         final BlockPosition blockPosition = position.toBlockPosition();
-                        final CustomBlock customBlock = instance.getCustomBlock(blockPosition);
+                        final CustomBlock customBlock = finalChunk.getCustomBlock(
+                                blockPosition.getX(),
+                                blockPosition.getY(),
+                                blockPosition.getZ());
                         if (customBlock != null) {
                             // Custom drag
                             drag = customBlock.getDrag(instance, blockPosition);
diff --git a/src/main/java/net/minestom/server/instance/Chunk.java b/src/main/java/net/minestom/server/instance/Chunk.java
index 39916fdbb..34d2fd817 100644
--- a/src/main/java/net/minestom/server/instance/Chunk.java
+++ b/src/main/java/net/minestom/server/instance/Chunk.java
@@ -266,6 +266,7 @@ public abstract class Chunk implements Viewable, DataContainer {
      * @param z the block Z
      * @return the {@link CustomBlock} at the position
      */
+    @Nullable
     public CustomBlock getCustomBlock(int x, int y, int z) {
         final short customBlockId = getCustomBlockId(x, y, z);
         return customBlockId != 0 ? BLOCK_MANAGER.getCustomBlock(customBlockId) : null;
@@ -277,6 +278,7 @@ public abstract class Chunk implements Viewable, DataContainer {
      * @param index the block index
      * @return the {@link CustomBlock} at the block index
      */
+    @Nullable
     protected CustomBlock getCustomBlock(int index) {
         final int x = ChunkUtils.blockIndexToChunkPositionX(index);
         final int y = ChunkUtils.blockIndexToChunkPositionY(index);