From 6c2e0c0b335c8f741d502b1805a46cca5b415946 Mon Sep 17 00:00:00 2001
From: Anuken <arnukren@gmail.com>
Date: Mon, 14 Feb 2022 18:24:18 -0500
Subject: [PATCH] Saving of unit commands

---
 core/src/mindustry/ai/types/CommandAI.java          | 5 +++++
 core/src/mindustry/entities/comp/UnitComp.java      | 4 +++-
 core/src/mindustry/entities/units/AIController.java | 8 ++++++++
 3 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/core/src/mindustry/ai/types/CommandAI.java b/core/src/mindustry/ai/types/CommandAI.java
index ab7492ec59..0f8fd4b108 100644
--- a/core/src/mindustry/ai/types/CommandAI.java
+++ b/core/src/mindustry/ai/types/CommandAI.java
@@ -78,6 +78,11 @@ public class CommandAI extends AIController{
         }
     }
 
+    @Override
+    public boolean keepState(){
+        return true;
+    }
+
     @Override
     public Teamc findTarget(float x, float y, float range, boolean air, boolean ground){
         return attackTarget == null ? super.findTarget(x, y, range, air, ground) : attackTarget;
diff --git a/core/src/mindustry/entities/comp/UnitComp.java b/core/src/mindustry/entities/comp/UnitComp.java
index 5d4c0c9988..73dcdf5355 100644
--- a/core/src/mindustry/entities/comp/UnitComp.java
+++ b/core/src/mindustry/entities/comp/UnitComp.java
@@ -334,7 +334,9 @@ abstract class UnitComp implements Healthc, Physicsc, Hitboxc, Statusc, Teamc, I
     public void afterRead(){
         afterSync();
         //reset controller state
-        controller(type.createController(self()));
+        if(!(controller instanceof AIController ai && ai.keepState())){
+            controller(type.createController(self()));
+        }
     }
 
     @Override
diff --git a/core/src/mindustry/entities/units/AIController.java b/core/src/mindustry/entities/units/AIController.java
index 7c97aba9d9..37811aad21 100644
--- a/core/src/mindustry/entities/units/AIController.java
+++ b/core/src/mindustry/entities/units/AIController.java
@@ -45,6 +45,14 @@ public class AIController implements UnitController{
         updateMovement();
     }
 
+    /**
+     * @return whether controller state should not be reset after reading.
+     * Do not override unless you know exactly what you are doing.
+     * */
+    public boolean keepState(){
+        return false;
+    }
+
     public void stopShooting(){
         for(var mount : unit.mounts){
             //ignore mount controllable stats too, they should not shoot either