From 80b92fb66778d6eb7d484343cbfecdf67eb1c625 Mon Sep 17 00:00:00 2001
From: Gustas <37534529+PunkPun@users.noreply.github.com>
Date: Sat, 7 Jan 2023 17:53:28 +0200
Subject: [PATCH] Fixed UnhardcodeBaseBuilderBotModule update rule

Update rules should should not read `modData.DefaultRules`
---
 .../UnhardcodeBaseBuilderBotModule.cs         | 81 ++++++++++++++++---
 1 file changed, 71 insertions(+), 10 deletions(-)

diff --git a/OpenRA.Mods.Common/UpdateRules/Rules/20210321/UnhardcodeBaseBuilderBotModule.cs b/OpenRA.Mods.Common/UpdateRules/Rules/20210321/UnhardcodeBaseBuilderBotModule.cs
index 90decf27e7..bee920a217 100644
--- a/OpenRA.Mods.Common/UpdateRules/Rules/20210321/UnhardcodeBaseBuilderBotModule.cs
+++ b/OpenRA.Mods.Common/UpdateRules/Rules/20210321/UnhardcodeBaseBuilderBotModule.cs
@@ -11,27 +11,88 @@
 
 using System.Collections.Generic;
 using System.Linq;
-using OpenRA.Mods.Common.Traits;
 
 namespace OpenRA.Mods.Common.UpdateRules.Rules
 {
 	public class UnhardcodeBaseBuilderBotModule : UpdateRule
 	{
+		MiniYamlNode defences;
+
+		// Excludes AttackBomber and AttackTDGunboatTurreted as actors with these AttackBase traits aren't supposed to be controlled.
+		readonly string[] attackBase = { "AttackLeap", "AttackPopupTurreted", "AttackAircraft", "AttackTesla", "AttackCharges", "AttackFollow", "AttackTurreted", "AttackFrontal", "AttackGarrisoned", "AttackOmni", "AttackSwallow" };
+		readonly string[] buildings = { "Building", "EnergyWall", "D2kBuilding" };
+
+		bool anyAdded;
+
 		public override string Name => "BaseBuilderBotModule got new fields to configure buildings that are defenses.";
 
 		public override string Description => "DefenseTypes were added.";
 
+		public override IEnumerable<string> BeforeUpdateActors(ModData modData, List<MiniYamlNode> resolvedActors)
+		{
+			var defences = new List<string>();
+
+			foreach (var actor in resolvedActors)
+			{
+				if (actor.Key.StartsWith('^'))
+					continue;
+
+				var isBuildable = false;
+				var isBuilding = false;
+				var canAttack = false;
+
+				foreach (var trait in actor.Value.Nodes)
+				{
+					if (trait.IsRemoval())
+						continue;
+
+					if (trait.KeyMatches("Buildable", includeRemovals: false))
+					{
+						isBuildable = true;
+						continue;
+					}
+
+					if (buildings.Any(v => trait.KeyMatches(v, includeRemovals: false)))
+					{
+						isBuilding = true;
+						continue;
+					}
+
+					if (attackBase.Any(ab => trait.KeyMatches(ab, includeRemovals: false)))
+						canAttack = true;
+				}
+
+				if (isBuildable && isBuilding && canAttack)
+				{
+					var name = actor.Key.ToLower();
+					if (!defences.Contains(name))
+						defences.Add(name);
+				}
+			}
+
+			this.defences = new MiniYamlNode("DefenseTypes", FieldSaver.FormatValue(defences));
+
+			yield break;
+		}
+
+		public override IEnumerable<string> AfterUpdate(ModData modData)
+		{
+			if (anyAdded)
+				yield return "`BaseBuilderBotModule` was unhardcoded and a new field added: `DefenseTypes`. Please verify the automated changes.";
+
+			anyAdded = false;
+		}
+
 		public override IEnumerable<string> UpdateActorNode(ModData modData, MiniYamlNode actorNode)
 		{
-			var addNodes = new List<MiniYamlNode>();
-
-			var defense = modData.DefaultRules.Actors.Values.Where(a => a.HasTraitInfo<BuildingInfo>() && a.HasTraitInfo<AttackBaseInfo>()).Select(a => a.Name);
-			var defensetypes = new MiniYamlNode("DefenseTypes", FieldSaver.FormatValue(defense.ToList()));
-			addNodes.Add(defensetypes);
-
-			foreach (var baseBuilderManager in actorNode.ChildrenMatching("BaseBuilderBotModule"))
-				foreach (var addNode in addNodes)
-					baseBuilderManager.AddNode(addNode);
+			foreach (var squadManager in actorNode.ChildrenMatching("BaseBuilderBotModule", includeRemovals: false))
+			{
+				if (!squadManager.ChildrenMatching(defences.Key, includeRemovals: false).Any())
+				{
+					squadManager.AddNode(defences);
+					anyAdded = true;
+				}
+			}
 
 			yield break;
 		}