From bc84e626886482a86d7e8cdd1a01d5522d9f2ea9 Mon Sep 17 00:00:00 2001
From: hondacrx <prelude_20032002@yahoo.com>
Date: Mon, 15 Nov 2021 23:22:08 -0500
Subject: [PATCH] Core/AI: Add some checks to catch if the AI changes
 mid-update for players.

Port From (https://github.com/TrinityCore/TrinityCore/commit/1015f307041b25b47a9bb4bc76ac4ed78c49d735)
---
 .../Game/Entities/Creature/Creature.Fields.cs   |  1 -
 Source/Game/Entities/Creature/Creature.cs       | 17 +----------------
 Source/Game/Entities/Unit/Unit.Fields.cs        |  1 +
 Source/Game/Entities/Unit/Unit.Pets.cs          |  1 -
 Source/Game/Entities/Unit/Unit.cs               | 17 +++++++++--------
 5 files changed, 11 insertions(+), 26 deletions(-)

diff --git a/Source/Game/Entities/Creature/Creature.Fields.cs b/Source/Game/Entities/Creature/Creature.Fields.cs
index d5c086135..90f29123a 100644
--- a/Source/Game/Entities/Creature/Creature.Fields.cs
+++ b/Source/Game/Entities/Creature/Creature.Fields.cs
@@ -55,7 +55,6 @@ namespace Game.Entities
         bool m_AlreadySearchedAssistance;
         bool m_cannotReachTarget;
         uint m_cannotReachTimer;
-        bool m_AI_locked;
 
         SpellSchoolMask m_meleeDamageSchoolMask;
         public uint m_originalEntry;
diff --git a/Source/Game/Entities/Creature/Creature.cs b/Source/Game/Entities/Creature/Creature.cs
index 6cbffaa56..9b1b92d5a 100644
--- a/Source/Game/Entities/Creature/Creature.cs
+++ b/Source/Game/Entities/Creature/Creature.cs
@@ -574,9 +574,7 @@ namespace Game.Entities
                     }
 
                     // do not allow the AI to be changed during update
-                    m_AI_locked = true;
-                    base.AIUpdateTick(diff);
-                    m_AI_locked = false;
+                    AIUpdateTick(diff);
 
                     if (!IsAlive())
                         break;
@@ -729,25 +727,12 @@ namespace Game.Entities
 
         bool DestoryAI()
         {
-            if (m_AI_locked)
-            {
-                Log.outDebug(LogFilter.Scripts, "DestroyAI: failed to destroy, locked.");
-                return false;
-            }
-
             SetAI(null);
             return true;
         }
 
         public bool InitializeAI(CreatureAI ai = null)
         {
-            // make sure nothing can change the AI during AI update
-            if (m_AI_locked)
-            {
-                Log.outDebug(LogFilter.Scripts, "InitializeAI: failed to init, locked.");
-                return false;
-            }
-
             InitializeMovementAI();
 
             SetAI(ai != null ? ai : AISelector.SelectAI(this));
diff --git a/Source/Game/Entities/Unit/Unit.Fields.cs b/Source/Game/Entities/Unit/Unit.Fields.cs
index caf43fa2c..8b99cf503 100644
--- a/Source/Game/Entities/Unit/Unit.Fields.cs
+++ b/Source/Game/Entities/Unit/Unit.Fields.cs
@@ -33,6 +33,7 @@ namespace Game.Entities
         //AI
         protected UnitAI i_AI;
         protected UnitAI i_disabledAI;
+        bool m_aiLocked;
 
         //Movement
         protected float[] m_speed_rate = new float[(int)UnitMoveType.Max];
diff --git a/Source/Game/Entities/Unit/Unit.Pets.cs b/Source/Game/Entities/Unit/Unit.Pets.cs
index 29a07ce06..b1cd34625 100644
--- a/Source/Game/Entities/Unit/Unit.Pets.cs
+++ b/Source/Game/Entities/Unit/Unit.Pets.cs
@@ -81,7 +81,6 @@ namespace Game.Entities
                 Cypher.Assert(newAI != null);
                 i_AI = newAI;
                 newAI.OnCharmed(true);
-                AIUpdateTick(0, true);
             }
             else
             {
diff --git a/Source/Game/Entities/Unit/Unit.cs b/Source/Game/Entities/Unit/Unit.cs
index 71cbc5e3e..0e9090ee0 100644
--- a/Source/Game/Entities/Unit/Unit.cs
+++ b/Source/Game/Entities/Unit/Unit.cs
@@ -1168,27 +1168,28 @@ namespace Game.Entities
 
         public virtual UnitAI GetAI() { return i_AI; }
 
-        public void AIUpdateTick(uint diff, bool force = false)
+        public void AIUpdateTick(uint diff)
         {
-            if (diff == 0) // some places call with diff = 0, which does nothing (for now), see PR #22296
-                return;
-
             UnitAI ai = GetAI();
             if (ai != null)
+            {
+                m_aiLocked = true;
                 ai.UpdateAI(diff);
+                m_aiLocked = false;
+            }
         }
 
         public void SetAI(UnitAI newAI)
         {
-            if (i_AI != null)
-                AIUpdateTick(0, true); // old AI gets a final tick if enabled
+            Cypher.Assert(!m_aiLocked, "Attempt to replace AI during AI update tick");
 
             i_AI = newAI;
-            AIUpdateTick(0, true); // new AI gets its initial tick
         }
 
         public void ScheduleAIChange()
         {
+            Cypher.Assert(!m_aiLocked, "Attempt to schedule AI change during AI update tick");
+
             bool charmed = IsCharmed();
             // if charm is applied, we can't have disabled AI already, and vice versa
             if (charmed)
@@ -1204,9 +1205,9 @@ namespace Game.Entities
 
         void RestoreDisabledAI()
         {
+            Cypher.Assert(!m_aiLocked, "Attempt to restore AI during UpdateAI tick");
             Cypher.Assert(IsPlayer() || i_disabledAI != null, "Attempt to restore disabled AI on creature without disabled AI");
             i_AI = i_disabledAI;
-            AIUpdateTick(0, true);
         }
 
         public bool IsPossessing()