From 54da48fd3edba448eb27ba899c5df08f37a39b0f Mon Sep 17 00:00:00 2001
From: Konstantin Shandurenko <k.shandurenko@corp.vk.com>
Date: Mon, 22 Feb 2021 07:33:11 +0300
Subject: [PATCH] Caching entity target found by goal selectors

---
 .../server/entity/ai/GoalSelector.java         | 17 +++++++++++++++++
 .../server/entity/ai/goal/MeleeAttackGoal.java | 18 +++---------------
 2 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/src/main/java/net/minestom/server/entity/ai/GoalSelector.java b/src/main/java/net/minestom/server/entity/ai/GoalSelector.java
index cc47581cf..06bfd14fd 100644
--- a/src/main/java/net/minestom/server/entity/ai/GoalSelector.java
+++ b/src/main/java/net/minestom/server/entity/ai/GoalSelector.java
@@ -60,6 +60,23 @@ public abstract class GoalSelector {
         return null;
     }
 
+    /**
+     * Finds a target for an entity.
+     * If the current one is not present, falls back to {@link GoalSelector#findTarget()}
+     * and updates the target inside the entity.
+     *
+     * @return the target entity, null if not found
+     */
+    @Nullable
+    public Entity findAndUpdateTarget() {
+        Entity target = entityCreature.getTarget();
+        if (target == null) {
+            target = findTarget();
+            entityCreature.setTarget(target);
+        }
+        return target;
+    }
+
     /**
      * Gets the entity behind the goal selector.
      *
diff --git a/src/main/java/net/minestom/server/entity/ai/goal/MeleeAttackGoal.java b/src/main/java/net/minestom/server/entity/ai/goal/MeleeAttackGoal.java
index 0b0d3cd9a..02e414f7f 100644
--- a/src/main/java/net/minestom/server/entity/ai/goal/MeleeAttackGoal.java
+++ b/src/main/java/net/minestom/server/entity/ai/goal/MeleeAttackGoal.java
@@ -40,12 +40,12 @@ public class MeleeAttackGoal extends GoalSelector {
 
     @Override
     public boolean shouldStart() {
-        return getTarget() != null;
+        return findAndUpdateTarget() != null;
     }
 
     @Override
     public void start() {
-        final Entity target = getTarget();
+        final Entity target = findAndUpdateTarget();
         Check.notNull(target, "The target is not expected to be null!");
         final Position targetPosition = target.getPosition();
         entityCreature.getNavigator().setPathTo(targetPosition);
@@ -53,7 +53,7 @@ public class MeleeAttackGoal extends GoalSelector {
 
     @Override
     public void tick(long time) {
-        final Entity target = getTarget();
+        final Entity target = findAndUpdateTarget();
 
         this.stop = target == null;
 
@@ -88,16 +88,4 @@ public class MeleeAttackGoal extends GoalSelector {
         // Stop following the target
         entityCreature.getNavigator().setPathTo(null);
     }
-
-    /**
-     * Use {@link EntityCreature#getTarget()} or
-     * the entity target selectors to get the correct target
-     *
-     * @return the target of the entity
-     */
-    @Nullable
-    private Entity getTarget() {
-        final Entity target = entityCreature.getTarget();
-        return target == null ? findTarget() : target;
-    }
 }