Fix primary building flag not working properly in D2k

OpenRA.Mods.Common/Traits/Buildings/PrimaryBuilding.cs

@@ -35,6 +35,10 @@ namespace OpenRA.Mods.Common.Traits
 		[Desc("The speech notification to play when selecting a primary building.")]
 		public readonly string SelectionNotification = "PrimaryBuildingSelected";
 
+		[Desc("List of production queues for which the primary flag should be set.",
+			"If empty, the list given in the `Produces` property of the `Production` trait will be used.")]
+		public readonly string[] ProductionQueues = { };
+
 		public object Create(ActorInitializer init) { return new PrimaryBuilding(init.Self, this); }
 	}
 
@@ -86,7 +90,8 @@ namespace OpenRA.Mods.Common.Traits
 			{
 				// Cancel existing primaries
 				// TODO: THIS IS SHIT
-				foreach (var p in self.Info.TraitInfo<ProductionInfo>().Produces)
+				var queues = info.ProductionQueues.Length == 0 ? self.Info.TraitInfo<ProductionInfo>().Produces : info.ProductionQueues;
+				foreach (var q in queues)
 				{
 					foreach (var b in self.World
 							.ActorsWithTrait<PrimaryBuilding>()
@@ -94,7 +99,7 @@ namespace OpenRA.Mods.Common.Traits
 								a.Actor != self &&
 								a.Actor.Owner == self.Owner &&
 								a.Trait.IsPrimary &&
-								a.Actor.Info.TraitInfo<ProductionInfo>().Produces.Contains(p)))
+								a.Actor.Info.TraitInfo<ProductionInfo>().Produces.Contains(q)))
 						b.Trait.SetPrimaryProducer(b.Actor, false);
 				}
 

mods/d2k/rules/structures.yaml

@@ -102,6 +102,7 @@ construction_yard:
 		Palette: d2k
 	PrimaryBuilding:
 		PrimaryCondition: primary
+		ProductionQueues: Building
 	ProvidesPrerequisite@buildingname:
 	GrantConditionOnPrerequisite:
 		Prerequisites: upgrade.conyard
@@ -209,6 +210,7 @@ barracks:
 		Produces: Infantry, Upgrade
 	PrimaryBuilding:
 		PrimaryCondition: primary
+		ProductionQueues: Infantry
 	ProductionBar:
 	ProvidesPrerequisite@atreides:
 		Prerequisite: barracks.atreides
@@ -408,6 +410,7 @@ light_factory:
 		Produces: Vehicle, Upgrade
 	PrimaryBuilding:
 		PrimaryCondition: primary
+		ProductionQueues: Vehicle
 	ProductionBar:
 	ProvidesPrerequisite@atreides:
 		Prerequisite: light.atreides
@@ -492,6 +495,7 @@ heavy_factory:
 		Produces: Armor, Upgrade
 	PrimaryBuilding:
 		PrimaryCondition: primary
+		ProductionQueues: Armor
 	ProductionBar:
 	ProvidesPrerequisite@atreides:
 		Prerequisite: heavy.atreides
@@ -653,6 +657,7 @@ starport:
 	ProductionBar:
 	PrimaryBuilding:
 		PrimaryCondition: primary
+		ProductionQueues: Starport
 	ProvidesPrerequisite@atreides:
 		Prerequisite: starport.atreides
 		Factions: atreides
@@ -905,6 +910,7 @@ high_tech_factory:
 	ProductionBar:
 	PrimaryBuilding:
 		PrimaryCondition: primary
+		ProductionQueues: Aircraft
 	Exit:
 		SpawnOffset: 0,0,728
 		ExitCell: 0,0