Improve BotModule performance.

Several parts of bot module logic, often through the AIUtils helper class, will query or count over all actors in the world. This is not a fast operation and the AI tends to repeat it often. Introduce some ActorIndex classes that can maintain an index of actors in the world that match a query based on a mix of actor name, owner or trait. These indexes introduce some overhead to maintain, but allow the queries or counts that bot modules needs to perform to be greatly sped up, as the index means there is a much smaller starting set of actors to consider. This is beneficial to the bot logic as the TraitDictionary index maintained by the world works only in terms of traits and doesn't allow the bot logic to perform a sufficiently selective lookup. This is because the bot logic is usually defined in terms of actor names rather than traits.
2024-02-24 18:44:13 +00:00
parent d4457a4028
commit dc0f26a1cd
9 changed files with 273 additions and 70 deletions
--- a/OpenRA.Mods.Common/AIUtils.cs
+++ b/OpenRA.Mods.Common/AIUtils.cs
@@ -41,31 +41,14 @@ namespace OpenRA.Mods.Common
 				.Select(a => a.Trait);
 		}

-		public static IEnumerable<Actor> GetActorsWithTrait<T>(World world)
+		public static int CountActorsWithNameAndTrait<T>(string actorName, Player owner)
 		{
-			return world.ActorsHavingTrait<T>();
+			return owner.World.ActorsHavingTrait<T>().Count(a => a.Owner == owner && a.Info.Name == actorName);
 		}

-		public static int CountActorsWithTrait<T>(string actorName, Player owner)
+		public static int CountActorByCommonName<T>(ActorIndex.OwnerAndNamesAndTrait<T> actorIndex)
 		{
-			return GetActorsWithTrait<T>(owner.World).Count(a => a.Owner == owner && a.Info.Name == actorName);
-		}
-
-		public static int CountActorByCommonName(HashSet<string> commonNames, Player owner)
-		{
-			return owner.World.Actors.Count(a => !a.IsDead && a.Owner == owner &&
-				commonNames.Contains(a.Info.Name));
-		}
-
-		public static int CountBuildingByCommonName(HashSet<string> buildings, Player owner)
-		{
-			return GetActorsWithTrait<Building>(owner.World)
-				.Count(a => a.Owner == owner && buildings.Contains(a.Info.Name));
-		}
-
-		public static ActorInfo GetInfoByCommonName(HashSet<string> names, Player owner)
-		{
-			return owner.World.Map.Rules.Actors.Where(k => names.Contains(k.Key)).Random(owner.World.LocalRandom).Value;
+			return actorIndex.Actors.Count(a => !a.IsDead);
 		}

 		public static void BotDebug(string format, params object[] args)
--- a/OpenRA.Mods.Common/ActorIndex.cs
+++ b/OpenRA.Mods.Common/ActorIndex.cs
@@ -0,0 +1,152 @@
+#region Copyright & License Information
+/*
+ * Copyright (c) The OpenRA Developers and Contributors
+ * This file is part of OpenRA, which is free software. It is made
+ * available to you under the terms of the GNU General Public License
+ * as published by the Free Software Foundation, either version 3 of
+ * the License, or (at your option) any later version. For more
+ * information, see COPYING.
+ */
+#endregion
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+
+namespace OpenRA.Mods.Common
+{
+	/// <summary>
+	/// Maintains an index of actors in the world.
+	/// </summary>
+	public abstract class ActorIndex : IDisposable
+	{
+		readonly World world;
+		readonly HashSet<Actor> actors = new();
+
+		public IReadOnlyCollection<Actor> Actors => actors;
+
+		ActorIndex(World world, IEnumerable<Actor> initialActorsToIndex)
+		{
+			this.world = world;
+			world.ActorAdded += AddActor;
+			world.ActorRemoved += RemoveActor;
+
+			actors.UnionWith(initialActorsToIndex);
+		}
+
+		protected abstract bool ShouldIndexActor(Actor actor);
+
+		void AddActor(Actor actor)
+		{
+			if (ShouldIndexActor(actor))
+				actors.Add(actor);
+		}
+
+		void RemoveActor(Actor actor)
+		{
+			actors.Remove(actor);
+		}
+
+		public void Dispose()
+		{
+			Dispose(true);
+			GC.SuppressFinalize(this);
+		}
+
+		protected virtual void Dispose(bool disposing)
+		{
+			if (disposing)
+			{
+				world.ActorAdded -= AddActor;
+				world.ActorRemoved -= RemoveActor;
+			}
+		}
+
+		// No OwnerAndTrait class is provided. As the world can provide actors by trait anyway,
+		// an additional filter on owner isn't sufficiently selective to justify the overhead of manging the index.
+		// Whereas a filter with actor names is much more selective, so OwnerAndNames is worthwhile.
+
+		/// <summary>
+		/// Maintains an index of actors in the world that
+		/// are owned by a given <see cref="Player"/>
+		/// and have one of the given <see cref="ActorInfo.Name"/>.
+		/// </summary>
+		public sealed class OwnerAndNames : ActorIndex
+		{
+			readonly HashSet<string> names;
+			readonly Player owner;
+
+			public OwnerAndNames(World world, IReadOnlyCollection<string> names, Player owner)
+				: base(world, ActorsToIndex(world, names.ToHashSet(), owner))
+			{
+				this.names = names.ToHashSet();
+				this.owner = owner;
+			}
+
+			static IEnumerable<Actor> ActorsToIndex(World world, HashSet<string> names, Player owner)
+			{
+				return world.Actors.Where(a => a.Owner == owner && names.Contains(a.Info.Name));
+			}
+
+			protected override bool ShouldIndexActor(Actor actor)
+			{
+				return actor.Owner == owner && names.Contains(actor.Info.Name);
+			}
+		}
+
+		/// <summary>
+		/// Maintains an index of actors in the world that
+		/// have one of the given <see cref="ActorInfo.Name"/>
+		/// and have the trait of type <typeparamref name="T"/>.
+		/// </summary>
+		public sealed class NamesAndTrait<T> : ActorIndex
+		{
+			readonly HashSet<string> names;
+
+			public NamesAndTrait(World world, IReadOnlyCollection<string> names)
+				: base(world, ActorsToIndex(world, names.ToHashSet()))
+			{
+				this.names = names.ToHashSet();
+			}
+
+			static IEnumerable<Actor> ActorsToIndex(World world, HashSet<string> names)
+			{
+				return world.ActorsHavingTrait<T>().Where(a => names.Contains(a.Info.Name));
+			}
+
+			protected override bool ShouldIndexActor(Actor actor)
+			{
+				return names.Contains(actor.Info.Name) && actor.TraitOrDefault<T>() != null;
+			}
+		}
+
+		/// <summary>
+		/// Maintains an index of actors in the world that
+		/// are owned by a given <see cref="Player"/>,
+		/// have one of the given <see cref="ActorInfo.Name"/>
+		/// and have the trait of type <typeparamref name="T"/>.
+		/// </summary>
+		public sealed class OwnerAndNamesAndTrait<T> : ActorIndex
+		{
+			readonly HashSet<string> names;
+			readonly Player owner;
+
+			public OwnerAndNamesAndTrait(World world, IReadOnlyCollection<string> names, Player owner)
+				: base(world, ActorsToIndex(world, names.ToHashSet(), owner))
+			{
+				this.names = names.ToHashSet();
+				this.owner = owner;
+			}
+
+			static IEnumerable<Actor> ActorsToIndex(World world, HashSet<string> names, Player owner)
+			{
+				return world.ActorsHavingTrait<T>().Where(a => a.Owner == owner && names.Contains(a.Info.Name));
+			}
+
+			protected override bool ShouldIndexActor(Actor actor)
+			{
+				return actor.Owner == owner && names.Contains(actor.Info.Name) && actor.TraitOrDefault<T>() != null;
+			}
+		}
+	}
+}
--- a/OpenRA.Mods.Common/Traits/BotModules/BaseBuilderBotModule.cs
+++ b/OpenRA.Mods.Common/Traits/BotModules/BaseBuilderBotModule.cs
@@ -139,12 +139,11 @@ namespace OpenRA.Mods.Common.Traits
 	}

 	public class BaseBuilderBotModule : ConditionalTrait<BaseBuilderBotModuleInfo>, IGameSaveTraitData,
-		IBotTick, IBotPositionsUpdated, IBotRespondToAttack, IBotRequestPauseUnitProduction
+		IBotTick, IBotPositionsUpdated, IBotRespondToAttack, IBotRequestPauseUnitProduction, INotifyActorDisposing
 	{
 		public CPos GetRandomBaseCenter()
 		{
-			var randomConstructionYard = world.Actors.Where(a => a.Owner == player &&
-				Info.ConstructionYardTypes.Contains(a.Info.Name))
+			var randomConstructionYard = constructionYardBuildings.Actors
 				.RandomOrDefault(world.LocalRandom);

 			return randomConstructionYard?.Location ?? initialBaseCenter;
@@ -166,12 +165,21 @@ namespace OpenRA.Mods.Common.Traits
 		readonly BaseBuilderQueueManager[] builders;
 		int currentBuilderIndex = 0;

+		readonly ActorIndex.OwnerAndNamesAndTrait<Building> refineryBuildings;
+		readonly ActorIndex.OwnerAndNamesAndTrait<Building> powerBuildings;
+		readonly ActorIndex.OwnerAndNamesAndTrait<Building> constructionYardBuildings;
+		readonly ActorIndex.OwnerAndNamesAndTrait<Building> barracksBuildings;
+
 		public BaseBuilderBotModule(Actor self, BaseBuilderBotModuleInfo info)
 			: base(info)
 		{
 			world = self.World;
 			player = self.Owner;
 			builders = new BaseBuilderQueueManager[info.BuildingQueues.Count + info.DefenseQueues.Count];
+			refineryBuildings = new ActorIndex.OwnerAndNamesAndTrait<Building>(world, info.RefineryTypes, player);
+			powerBuildings = new ActorIndex.OwnerAndNamesAndTrait<Building>(world, info.PowerTypes, player);
+			constructionYardBuildings = new ActorIndex.OwnerAndNamesAndTrait<Building>(world, info.ConstructionYardTypes, player);
+			barracksBuildings = new ActorIndex.OwnerAndNamesAndTrait<Building>(world, info.BarracksTypes, player);
 		}

 		protected override void Created(Actor self)
@@ -200,7 +208,7 @@ namespace OpenRA.Mods.Common.Traits
 			DefenseCenter = newLocation;
 		}

-		bool IBotRequestPauseUnitProduction.PauseUnitProduction => !IsTraitDisabled && !HasAdequateRefineryCount;
+		bool IBotRequestPauseUnitProduction.PauseUnitProduction => !IsTraitDisabled && !HasAdequateRefineryCount();

 		void IBotTick.BotTick(IBot bot)
 		{
@@ -305,13 +313,16 @@ namespace OpenRA.Mods.Common.Traits
 		}

 		// Require at least one refinery, unless we can't build it.
-		public bool HasAdequateRefineryCount =>
+		public bool HasAdequateRefineryCount() =>
 			Info.RefineryTypes.Count == 0 ||
-			AIUtils.CountBuildingByCommonName(Info.RefineryTypes, player) >= MinimumRefineryCount ||
-			AIUtils.CountBuildingByCommonName(Info.PowerTypes, player) == 0 ||
-			AIUtils.CountBuildingByCommonName(Info.ConstructionYardTypes, player) == 0;
+			AIUtils.CountActorByCommonName(refineryBuildings) >= MinimumRefineryCount() ||
+			AIUtils.CountActorByCommonName(powerBuildings) == 0 ||
+			AIUtils.CountActorByCommonName(constructionYardBuildings) == 0;

-		int MinimumRefineryCount => AIUtils.CountBuildingByCommonName(Info.BarracksTypes, player) > 0 ? Info.InititalMinimumRefineryCount + Info.AdditionalMinimumRefineryCount : Info.InititalMinimumRefineryCount;
+		int MinimumRefineryCount() =>
+			AIUtils.CountActorByCommonName(barracksBuildings) > 0
+			? Info.InititalMinimumRefineryCount + Info.AdditionalMinimumRefineryCount
+			: Info.InititalMinimumRefineryCount;

 		List<MiniYamlNode> IGameSaveTraitData.IssueTraitData(Actor self)
 		{
@@ -338,5 +349,13 @@ namespace OpenRA.Mods.Common.Traits
 			if (defenseCenterNode != null)
 				DefenseCenter = FieldLoader.GetValue<CPos>("DefenseCenter", defenseCenterNode.Value.Value);
 		}
+
+		void INotifyActorDisposing.Disposing(Actor self)
+		{
+			refineryBuildings.Dispose();
+			powerBuildings.Dispose();
+			constructionYardBuildings.Dispose();
+			barracksBuildings.Dispose();
+		}
 	}
 }
--- a/OpenRA.Mods.Common/Traits/BotModules/BotModuleLogic/BaseBuilderQueueManager.cs
+++ b/OpenRA.Mods.Common/Traits/BotModules/BotModuleLogic/BaseBuilderQueueManager.cs
@@ -255,7 +255,7 @@ namespace OpenRA.Mods.Common.Traits
 			}

 			// Next is to build up a strong economy
-			if (!baseBuilder.HasAdequateRefineryCount)
+			if (!baseBuilder.HasAdequateRefineryCount())
 			{
 				var refinery = GetProducibleBuilding(baseBuilder.Info.RefineryTypes, buildableThings);
 				if (refinery != null && HasSufficientPowerForActor(refinery))
--- a/OpenRA.Mods.Common/Traits/BotModules/CaptureManagerBotModule.cs
+++ b/OpenRA.Mods.Common/Traits/BotModules/CaptureManagerBotModule.cs
@@ -44,7 +44,7 @@ namespace OpenRA.Mods.Common.Traits
 		public override object Create(ActorInitializer init) { return new CaptureManagerBotModule(init.Self, this); }
 	}

-	public class CaptureManagerBotModule : ConditionalTrait<CaptureManagerBotModuleInfo>, IBotTick
+	public class CaptureManagerBotModule : ConditionalTrait<CaptureManagerBotModuleInfo>, IBotTick, INotifyActorDisposing
 	{
 		readonly World world;
 		readonly Player player;
@@ -55,6 +55,8 @@ namespace OpenRA.Mods.Common.Traits
 		// Units that the bot already knows about and has given a capture order. Any unit not on this list needs to be given a new order.
 		readonly List<Actor> activeCapturers = new();

+		readonly ActorIndex.OwnerAndNamesAndTrait<Captures> capturingActors;
+
 		public CaptureManagerBotModule(Actor self, CaptureManagerBotModuleInfo info)
 			: base(info)
 		{
@@ -67,6 +69,8 @@ namespace OpenRA.Mods.Common.Traits
 			unitCannotBeOrderedOrIsIdle = a => a.Owner != player || a.IsDead || !a.IsInWorld || a.IsIdle;

 			maximumCaptureTargetOptions = Math.Max(1, Info.MaximumCaptureTargetOptions);
+
+			capturingActors = new ActorIndex.OwnerAndNamesAndTrait<Captures>(world, Info.CapturingActorTypes, player);
 		}

 		protected override void TraitEnabled(Actor self)
@@ -105,11 +109,8 @@ namespace OpenRA.Mods.Common.Traits

 			activeCapturers.RemoveAll(unitCannotBeOrderedOrIsIdle);

-			var newUnits = world.ActorsHavingTrait<IPositionable>()
-				.Where(a => a.Owner == player && !activeCapturers.Contains(a));
-
-			var capturers = newUnits
-				.Where(a => a.IsIdle && Info.CapturingActorTypes.Contains(a.Info.Name) && a.Info.HasTraitInfo<CapturesInfo>())
+			var capturers = capturingActors.Actors
+				.Where(a => a.IsIdle && a.Info.HasTraitInfo<IPositionableInfo>() && !activeCapturers.Contains(a))
 				.Select(a => new TraitPair<CaptureManager>(a, a.TraitOrDefault<CaptureManager>()))
 				.Where(tp => tp.Trait != null)
 				.ToArray();
@@ -154,5 +155,10 @@ namespace OpenRA.Mods.Common.Traits
 				activeCapturers.Add(capturer.Actor);
 			}
 		}
+
+		void INotifyActorDisposing.Disposing(Actor self)
+		{
+			capturingActors.Dispose();
+		}
 	}
 }
--- a/OpenRA.Mods.Common/Traits/BotModules/HarvesterBotModule.cs
+++ b/OpenRA.Mods.Common/Traits/BotModules/HarvesterBotModule.cs
@@ -37,7 +37,7 @@ namespace OpenRA.Mods.Common.Traits
 		public override object Create(ActorInitializer init) { return new HarvesterBotModule(init.Self, this); }
 	}

-	public class HarvesterBotModule : ConditionalTrait<HarvesterBotModuleInfo>, IBotTick
+	public class HarvesterBotModule : ConditionalTrait<HarvesterBotModuleInfo>, IBotTick, INotifyActorDisposing
 	{
 		sealed class HarvesterTraitWrapper
 		{
@@ -59,6 +59,8 @@ namespace OpenRA.Mods.Common.Traits
 		readonly Player player;
 		readonly Func<Actor, bool> unitCannotBeOrdered;
 		readonly Dictionary<Actor, HarvesterTraitWrapper> harvesters = new();
+		readonly ActorIndex.OwnerAndNamesAndTrait<Building> refineries;
+		readonly ActorIndex.OwnerAndNamesAndTrait<Harvester> harvestersIndex;

 		IResourceLayer resourceLayer;
 		ResourceClaimLayer claimLayer;
@@ -71,6 +73,8 @@ namespace OpenRA.Mods.Common.Traits
 			world = self.World;
 			player = self.Owner;
 			unitCannotBeOrdered = a => a.Owner != self.Owner || a.IsDead || !a.IsInWorld;
+			refineries = new ActorIndex.OwnerAndNamesAndTrait<Building>(world, info.RefineryTypes, player);
+			harvestersIndex = new ActorIndex.OwnerAndNamesAndTrait<Harvester>(world, info.HarvesterTypes, player);
 		}

 		protected override void Created(Actor self)
@@ -102,7 +106,6 @@ namespace OpenRA.Mods.Common.Traits
 			scanForIdleHarvestersTicks = Info.ScanForIdleHarvestersInterval;

 			// Find new harvesters
-			// TODO: Look for a more performance-friendly way to update this list
 			var newHarvesters = world.ActorsHavingTrait<Harvester>().Where(a => !unitCannotBeOrdered(a) && !harvesters.ContainsKey(a));
 			foreach (var a in newHarvesters)
 				harvesters[a] = new HarvesterTraitWrapper(a);
@@ -133,10 +136,15 @@ namespace OpenRA.Mods.Common.Traits
 			var unitBuilder = requestUnitProduction.FirstEnabledTraitOrDefault();
 			if (unitBuilder != null && Info.HarvesterTypes.Count > 0)
 			{
-				var harvInfo = AIUtils.GetInfoByCommonName(Info.HarvesterTypes, player);
-				var harvCountTooLow = AIUtils.CountActorByCommonName(Info.HarvesterTypes, player) < AIUtils.CountBuildingByCommonName(Info.RefineryTypes, player);
-				if (harvCountTooLow && unitBuilder.RequestedProductionCount(bot, harvInfo.Name) == 0)
-					unitBuilder.RequestUnitProduction(bot, harvInfo.Name);
+				var harvCountTooLow =
+					AIUtils.CountActorByCommonName(harvestersIndex) <
+					AIUtils.CountActorByCommonName(refineries);
+				if (harvCountTooLow)
+				{
+					var harvesterType = Info.HarvesterTypes.Random(world.LocalRandom);
+					if (unitBuilder.RequestedProductionCount(bot, harvesterType) == 0)
+						unitBuilder.RequestUnitProduction(bot, harvesterType);
+				}
 			}
 		}

@@ -157,5 +165,11 @@ namespace OpenRA.Mods.Common.Traits

 			return Target.FromCell(world, path[0]);
 		}
+
+		void INotifyActorDisposing.Disposing(Actor self)
+		{
+			refineries.Dispose();
+			harvestersIndex.Dispose();
+		}
 	}
 }
--- a/OpenRA.Mods.Common/Traits/BotModules/McvManagerBotModule.cs
+++ b/OpenRA.Mods.Common/Traits/BotModules/McvManagerBotModule.cs
@@ -47,12 +47,12 @@ namespace OpenRA.Mods.Common.Traits
 		public override object Create(ActorInitializer init) { return new McvManagerBotModule(init.Self, this); }
 	}

-	public class McvManagerBotModule : ConditionalTrait<McvManagerBotModuleInfo>, IBotTick, IBotPositionsUpdated, IGameSaveTraitData
+	public class McvManagerBotModule : ConditionalTrait<McvManagerBotModuleInfo>,
+		IBotTick, IBotPositionsUpdated, IGameSaveTraitData, INotifyActorDisposing
 	{
 		public CPos GetRandomBaseCenter()
 		{
-			var randomConstructionYard = world.Actors.Where(a => a.Owner == player &&
-				Info.ConstructionYardTypes.Contains(a.Info.Name))
+			var randomConstructionYard = constructionYards.Actors
 				.RandomOrDefault(world.LocalRandom);

 			return randomConstructionYard?.Location ?? initialBaseCenter;
@@ -60,6 +60,9 @@ namespace OpenRA.Mods.Common.Traits

 		readonly World world;
 		readonly Player player;
+		readonly ActorIndex.OwnerAndNamesAndTrait<Transforms> mcvs;
+		readonly ActorIndex.OwnerAndNamesAndTrait<Building> constructionYards;
+		readonly ActorIndex.OwnerAndNamesAndTrait<Building> mcvFactories;

 		IBotPositionsUpdated[] notifyPositionsUpdated;
 		IBotRequestUnitProduction[] requestUnitProduction;
@@ -73,6 +76,9 @@ namespace OpenRA.Mods.Common.Traits
 		{
 			world = self.World;
 			player = self.Owner;
+			mcvs = new ActorIndex.OwnerAndNamesAndTrait<Transforms>(world, info.McvTypes, player);
+			constructionYards = new ActorIndex.OwnerAndNamesAndTrait<Building>(world, info.ConstructionYardTypes, player);
+			mcvFactories = new ActorIndex.OwnerAndNamesAndTrait<Building>(world, info.McvFactoryTypes, player);
 		}

 		protected override void Created(Actor self)
@@ -108,15 +114,12 @@ namespace OpenRA.Mods.Common.Traits
 				DeployMcvs(bot, true);

 				// No construction yards - Build a new MCV
-				if (ShouldBuildMCV())
+				var unitBuilder = requestUnitProduction.FirstEnabledTraitOrDefault();
+				if (unitBuilder != null && Info.McvTypes.Count > 0 && ShouldBuildMCV())
 				{
-					var unitBuilder = requestUnitProduction.FirstEnabledTraitOrDefault();
-					if (unitBuilder != null)
-					{
-						var mcvInfo = AIUtils.GetInfoByCommonName(Info.McvTypes, player);
-						if (unitBuilder.RequestedProductionCount(bot, mcvInfo.Name) == 0)
-							unitBuilder.RequestUnitProduction(bot, mcvInfo.Name);
-					}
+					var mcvType = Info.McvTypes.Random(world.LocalRandom);
+					if (unitBuilder.RequestedProductionCount(bot, mcvType) == 0)
+						unitBuilder.RequestUnitProduction(bot, mcvType);
 				}
 			}
 		}
@@ -124,19 +127,19 @@ namespace OpenRA.Mods.Common.Traits
 		bool ShouldBuildMCV()
 		{
 			// Only build MCV if we don't already have one in the field.
-			var allowedToBuildMCV = AIUtils.CountActorByCommonName(Info.McvTypes, player) == 0;
+			var allowedToBuildMCV = AIUtils.CountActorByCommonName(mcvs) == 0;
 			if (!allowedToBuildMCV)
 				return false;

 			// Build MCV if we don't have the desired number of construction yards, unless we have no factory (can't build it).
-			return AIUtils.CountBuildingByCommonName(Info.ConstructionYardTypes, player) < Info.MinimumConstructionYardCount &&
-				AIUtils.CountBuildingByCommonName(Info.McvFactoryTypes, player) > 0;
+			return AIUtils.CountActorByCommonName(constructionYards) < Info.MinimumConstructionYardCount &&
+				AIUtils.CountActorByCommonName(mcvFactories) > 0;
 		}

 		void DeployMcvs(IBot bot, bool chooseLocation)
 		{
-			var newMCVs = world.ActorsHavingTrait<Transforms>()
-				.Where(a => a.Owner == player && a.IsIdle && Info.McvTypes.Contains(a.Info.Name));
+			var newMCVs = mcvs.Actors
+				.Where(a => a.IsIdle);

 			foreach (var mcv in newMCVs)
 				DeployMcv(bot, mcv, chooseLocation);
@@ -148,7 +151,9 @@ namespace OpenRA.Mods.Common.Traits
 			if (move)
 			{
 				// If we lack a base, we need to make sure we don't restrict deployment of the MCV to the base!
-				var restrictToBase = Info.RestrictMCVDeploymentFallbackToBase && AIUtils.CountBuildingByCommonName(Info.ConstructionYardTypes, player) > 0;
+				var restrictToBase =
+					Info.RestrictMCVDeploymentFallbackToBase &&
+					AIUtils.CountActorByCommonName(constructionYards) > 0;

 				var transformsInfo = mcv.Info.TraitInfo<TransformsInfo>();
 				var desiredLocation = ChooseMcvDeployLocation(transformsInfo.IntoActor, transformsInfo.Offset, restrictToBase);
@@ -221,5 +226,12 @@ namespace OpenRA.Mods.Common.Traits
 			if (initialBaseCenterNode != null)
 				initialBaseCenter = FieldLoader.GetValue<CPos>("InitialBaseCenter", initialBaseCenterNode.Value.Value);
 		}
+
+		void INotifyActorDisposing.Disposing(Actor self)
+		{
+			mcvs.Dispose();
+			constructionYards.Dispose();
+			mcvFactories.Dispose();
+		}
 	}
 }
--- a/OpenRA.Mods.Common/Traits/BotModules/SquadManagerBotModule.cs
+++ b/OpenRA.Mods.Common/Traits/BotModules/SquadManagerBotModule.cs
@@ -103,12 +103,13 @@ namespace OpenRA.Mods.Common.Traits
 		public override object Create(ActorInitializer init) { return new SquadManagerBotModule(init.Self, this); }
 	}

-	public class SquadManagerBotModule : ConditionalTrait<SquadManagerBotModuleInfo>, IBotEnabled, IBotTick, IBotRespondToAttack, IBotPositionsUpdated, IGameSaveTraitData
+	public class SquadManagerBotModule : ConditionalTrait<SquadManagerBotModuleInfo>,
+		IBotEnabled, IBotTick, IBotRespondToAttack, IBotPositionsUpdated, IGameSaveTraitData, INotifyActorDisposing
 	{
 		public CPos GetRandomBaseCenter()
 		{
-			var randomConstructionYard = World.Actors.Where(a => a.Owner == Player &&
-				Info.ConstructionYardTypes.Contains(a.Info.Name))
+			var randomConstructionYard = constructionYardBuildings.Actors
+				.Where(a => a.Owner == Player)
 				.RandomOrDefault(World.LocalRandom);

 			return randomConstructionYard?.Location ?? initialBaseCenter;
@@ -124,6 +125,7 @@ namespace OpenRA.Mods.Common.Traits
 		readonly HashSet<Actor> activeUnits = new();

 		public List<Squad> Squads = new();
+		readonly ActorIndex.NamesAndTrait<Building> constructionYardBuildings;

 		IBot bot;
 		IBotPositionsUpdated[] notifyPositionsUpdated;
@@ -143,6 +145,7 @@ namespace OpenRA.Mods.Common.Traits
 			Player = self.Owner;

 			unitCannotBeOrdered = a => a == null || a.Owner != Player || a.IsDead || !a.IsInWorld;
+			constructionYardBuildings = new ActorIndex.NamesAndTrait<Building>(World, info.ConstructionYardTypes);
 		}

 		// Use for proactive targeting.
@@ -427,7 +430,7 @@ namespace OpenRA.Mods.Common.Traits
 				return;

 			var allEnemyBaseBuilder = FindEnemies(
-				World.Actors.Where(a => Info.ConstructionYardTypes.Contains(a.Info.Name)),
+				constructionYardBuildings.Actors,
 				ownUnits[0])
 				.ToList();

@@ -568,5 +571,10 @@ namespace OpenRA.Mods.Common.Traits
 					Squads.Add(Squad.Deserialize(bot, this, n.Value));
 			}
 		}
+
+		void INotifyActorDisposing.Disposing(Actor self)
+		{
+			constructionYardBuildings.Dispose();
+		}
 	}
 }
--- a/OpenRA.Mods.Common/Traits/BotModules/UnitBuilderBotModule.cs
+++ b/OpenRA.Mods.Common/Traits/BotModules/UnitBuilderBotModule.cs
@@ -44,7 +44,8 @@ namespace OpenRA.Mods.Common.Traits
 		public override object Create(ActorInitializer init) { return new UnitBuilderBotModule(init.Self, this); }
 	}

-	public class UnitBuilderBotModule : ConditionalTrait<UnitBuilderBotModuleInfo>, IBotTick, IBotNotifyIdleBaseUnits, IBotRequestUnitProduction, IGameSaveTraitData
+	public class UnitBuilderBotModule : ConditionalTrait<UnitBuilderBotModuleInfo>,
+		IBotTick, IBotNotifyIdleBaseUnits, IBotRequestUnitProduction, IGameSaveTraitData, INotifyActorDisposing
 	{
 		public const int FeedbackTime = 30; // ticks; = a bit over 1s. must be >= netlag.

@@ -52,6 +53,7 @@ namespace OpenRA.Mods.Common.Traits
 		readonly Player player;

 		readonly List<string> queuedBuildRequests = new();
+		readonly ActorIndex.OwnerAndNames unitsToBuild;

 		IBotRequestPauseUnitProduction[] requestPause;
 		int idleUnitCount;
@@ -65,6 +67,7 @@ namespace OpenRA.Mods.Common.Traits
 		{
 			world = self.World;
 			player = self.Owner;
+			unitsToBuild = new ActorIndex.OwnerAndNames(world, info.UnitsToBuild.Keys, player);
 		}

 		protected override void Created(Actor self)
@@ -174,20 +177,21 @@ namespace OpenRA.Mods.Common.Traits
 			if (buildableThings.Length == 0)
 				return null;

-			var allUnits = world.Actors.Where(a => a.Owner == player && Info.UnitsToBuild.ContainsKey(a.Info.Name) && !a.IsDead).ToArray();
+			var allUnits = unitsToBuild.Actors.Where(a => !a.IsDead).ToArray();

 			ActorInfo desiredUnit = null;
 			var desiredError = int.MaxValue;
 			foreach (var unit in buildableThings)
 			{
-				if (!Info.UnitsToBuild.ContainsKey(unit.Name) || (Info.UnitDelays != null && Info.UnitDelays.TryGetValue(unit.Name, out var delay) && delay > world.WorldTick))
+				if (!Info.UnitsToBuild.TryGetValue(unit.Name, out var share) ||
+					(Info.UnitDelays != null && Info.UnitDelays.TryGetValue(unit.Name, out var delay) && delay > world.WorldTick))
 					continue;

 				var unitCount = allUnits.Count(a => a.Info.Name == unit.Name);
 				if (Info.UnitLimits != null && Info.UnitLimits.TryGetValue(unit.Name, out var count) && unitCount >= count)
 					continue;

-				var error = allUnits.Length > 0 ? unitCount * 100 / allUnits.Length - Info.UnitsToBuild[unit.Name] : -1;
+				var error = allUnits.Length > 0 ? unitCount * 100 / allUnits.Length - share : -1;
 				if (error < 0)
 					return HasAdequateAirUnitReloadBuildings(unit) ? unit : null;

@@ -213,8 +217,8 @@ namespace OpenRA.Mods.Common.Traits
 			if (rearmableInfo == null)
 				return true;

-			var countOwnAir = AIUtils.CountActorsWithTrait<IPositionable>(actorInfo.Name, player);
-			var countBuildings = rearmableInfo.RearmActors.Sum(b => AIUtils.CountActorsWithTrait<Building>(b, player));
+			var countOwnAir = AIUtils.CountActorsWithNameAndTrait<IPositionable>(actorInfo.Name, player);
+			var countBuildings = rearmableInfo.RearmActors.Sum(b => AIUtils.CountActorsWithNameAndTrait<Building>(b, player));
 			if (countOwnAir >= countBuildings)
 				return false;

@@ -249,5 +253,10 @@ namespace OpenRA.Mods.Common.Traits
 			if (idleUnitCountNode != null)
 				idleUnitCount = FieldLoader.GetValue<int>("IdleUnitCount", idleUnitCountNode.Value.Value);
 		}
+
+		void INotifyActorDisposing.Disposing(Actor self)
+		{
+			unitsToBuild.Dispose();
+		}
 	}
 }