Strip BA parts from non-BA units MegaMek#2396

sixlettervariables · Feb 6, 2021 · 0a95d4f · 0a95d4f
1 parent b1f5d76
commit 0a95d4f
Showing 1 changed file with 22 additions and 12 deletions.
diff --git a/MekHQ/src/mekhq/campaign/unit/Unit.java b/MekHQ/src/mekhq/campaign/unit/Unit.java
@@ -2285,13 +2285,15 @@ public void initializeParts(boolean addParts) {
                     partsToRemove.add(part);
                 }
             } else if (part instanceof BattleArmorSuit) {
-                if (((BattleArmorSuit) part).getTrooper() < locations.length) {
+                if ((entity instanceof BattleArmor)
+                        && ((BattleArmorSuit) part).getTrooper() < locations.length) {
                     locations[((BattleArmorSuit) part).getTrooper()] = part;
                 } else {
                     partsToRemove.add(part);
                 }
             } else if (part instanceof MissingBattleArmorSuit) {
-                if (((MissingBattleArmorSuit) part).getTrooper() < locations.length) {
+                if ((entity instanceof BattleArmor)
+                        && ((MissingBattleArmorSuit) part).getTrooper() < locations.length) {
                     locations[((MissingBattleArmorSuit) part).getTrooper()] = part;
                 } else {
                     partsToRemove.add(part);
@@ -2325,19 +2327,27 @@ public void initializeParts(boolean addParts) {
             } else if (part instanceof MissingJumpJet) {
                 jumpJets.put(((MissingJumpJet)part).getEquipmentNum(), part);
             }  else if (part instanceof BattleArmorEquipmentPart) {
-                Part[] parts = baEquipParts.get(((BattleArmorEquipmentPart)part).getEquipmentNum());
-                if (null == parts) {
-                    parts = new Part[((BattleArmor)entity).getSquadSize()];
+                if (!(entity instanceof BattleArmor)) {
+                    partsToRemove.add(part);
+                } else {
+                    Part[] parts = baEquipParts.get(((BattleArmorEquipmentPart)part).getEquipmentNum());
+                    if (null == parts) {
+                        parts = new Part[((BattleArmor)entity).getSquadSize()];
+                    }
+                    parts[((BattleArmorEquipmentPart)part).getTrooper()-BattleArmor.LOC_TROOPER_1] = part;
+                    baEquipParts.put(((BattleArmorEquipmentPart)part).getEquipmentNum(), parts);
                 }
-                parts[((BattleArmorEquipmentPart)part).getTrooper()-BattleArmor.LOC_TROOPER_1] = part;
-                baEquipParts.put(((BattleArmorEquipmentPart)part).getEquipmentNum(), parts);
             } else if (part instanceof MissingBattleArmorEquipmentPart) {
-                Part[] parts = baEquipParts.get(((MissingBattleArmorEquipmentPart)part).getEquipmentNum());
-                if (null == parts) {
-                    parts = new Part[((BattleArmor)entity).getSquadSize()];
+                if (!(entity instanceof BattleArmor)) {
+                    partsToRemove.add(part);
+                } else {
+                    Part[] parts = baEquipParts.get(((MissingBattleArmorEquipmentPart)part).getEquipmentNum());
+                    if (null == parts) {
+                        parts = new Part[((BattleArmor)entity).getSquadSize()];
+                    }
+                    parts[((MissingBattleArmorEquipmentPart)part).getTrooper()-BattleArmor.LOC_TROOPER_1] = part;
+                    baEquipParts.put(((MissingBattleArmorEquipmentPart)part).getEquipmentNum(), parts);
                 }
-                parts[((MissingBattleArmorEquipmentPart)part).getTrooper()-BattleArmor.LOC_TROOPER_1] = part;
-                baEquipParts.put(((MissingBattleArmorEquipmentPart)part).getEquipmentNum(), parts);
             } else if (part instanceof EquipmentPart) {
                 equipParts.put(((EquipmentPart)part).getEquipmentNum(), part);
             } else if (part instanceof MissingEquipmentPart) {