MegaMek · psikomonkie · May 15, 2025 · May 15, 2025 · May 15, 2025
@@ -1654,6 +1654,20 @@ public void addCampaignTransport(CampaignTransportType campaignTransportType, Un
         }
     }
 
+    /**
+     * This will update the transport in the transports list with current capacities. When a unit is added or removed
+     * from a transport, that information needs updated in the campaign transport map. This method will update the
+     * map for every {@code CampaignTransportType} for the given transport.
+     *
+     * @see Campaign#updateTransportInTransports(CampaignTransportType, Unit)
+     * @param transport             Unit
+     */
+    public void updateTransportInTransports(Unit transport) {
+        for (CampaignTransportType campaignTransportType : CampaignTransportType.values()) {
+            updateTransportInTransports (campaignTransportType, transport);
+        }
+    }
+
     /**
      * This will update the transport in the transports list with current capacities. When a unit is added or removed
      * from a transport, that information needs updated in the campaign transport map. This method takes the

@@ -87,11 +87,18 @@ public void updateTransportInTransporterMap(Unit transport) {
             return;
         }
         AbstractTransportedUnitsSummary transportedUnitsSummary = transport.getTransportedUnitsSummary(campaignTransportType);
-        for (TransporterType transporterType : transportedUnitsSummary.getTransportCapabilities()) {
+
+        // Let's make a list of all the transportTypes in the map, and all the transportTypes the unit has
+        Set<TransporterType> transporterTypes = new HashSet<>();
+        transporterTypes.addAll(transportedUnitsSummary.getTransportCapabilities());
+        transporterTypes.addAll(transportersMap.keySet());
+
+        // Now let's update the current transporterTypes
+        for (TransporterType transporterType : transporterTypes) {
             if (transportersMap.containsKey(transporterType)) {
                 Set<Double> oldCapacities = transportersMap.get(transporterType).keySet();
                 Double newCapacity = transportedUnitsSummary.getCurrentTransportCapacity(transporterType);
-                //First, if this is a new capacity for the map, let's manually add it
+                // First, if this is a new capacity for the map, let's manually add it
                 if (!oldCapacities.contains(newCapacity)) {
                     addTransporterToCapacityMap(transport, transporterType);
                 }
@@ -105,9 +112,14 @@ public void updateTransportInTransporterMap(Unit transport) {
                         addTransporterToCapacityMap(transport, transporterType);
                     }
                 }
+
+                // Finally, let's remove this from the map & get out if the transport doesn't have this transporterType
+                if (!transportedUnitsSummary.getTransportCapabilities().contains(transporterType)) {
+                    removeTransportFromCapacityMap(transport, transporterType, 0);
+                }
             }
             else {
-                logger.error(String.format("Invalid transporter type %s", transporterType));
+                addTransporterToCapacityMap (transport, transporterType);
             }
         }
     }
@@ -183,14 +195,18 @@ public void removeTransport(Unit transport) {
 
         for (TransporterType transporterTypeToRemove : toRemoveMap.keySet()) {
             double capacity = toRemoveMap.get(transporterTypeToRemove);
+            removeTransportFromCapacityMap(transport, transporterTypeToRemove, capacity);
+        }
+    }
 
-            transportersMap.get(transporterTypeToRemove).get(capacity).remove(transport.getId());
-            if (transportersMap.get(transporterTypeToRemove).get(capacity).isEmpty()) {
-                transportersMap.get(transporterTypeToRemove).remove(capacity);
-            }
-            if (transportersMap.get(transporterTypeToRemove).isEmpty()) {
-                transportersMap.remove(transporterTypeToRemove);
-            }
+    private void removeTransportFromCapacityMap(Unit transport, TransporterType transporterTypeToRemove, double capacity) {
+
+        transportersMap.get(transporterTypeToRemove).get(capacity).remove(transport.getId());
+        if (transportersMap.get(transporterTypeToRemove).get(capacity).isEmpty()) {
+            transportersMap.get(transporterTypeToRemove).remove(capacity);
+        }
+        if (transportersMap.get(transporterTypeToRemove).isEmpty()) {
+            transportersMap.remove(transporterTypeToRemove);
         }
     }
 }
@@ -1554,8 +1554,6 @@ private void complete() {
         // don't forget to switch entities!
         // ----------------- from here on oldUnit refers to the new entity -------------------------
         oldUnit.setEntity(newEntity);
-        // Bay capacities might have changed - reset them
-        oldUnit.initializeAllTransportSpace();
 
         // set up new parts
         ArrayList<Part> newParts = new ArrayList<>();
@@ -1697,6 +1695,11 @@ private void complete() {
         getCampaign().reloadGameEntities();
         C3Util.copyC3Networks(oldEntity, oldUnit.getEntity());
 
+        // Bay capacities might have changed - reset them
+        oldUnit.clearAllTransportSpace();
+        oldUnit.initializeAllTransportSpace();
+        campaign.updateTransportInTransports(oldUnit);
+
         // reload any soldiers
         for (Person soldier : soldiers) {
             if (!oldUnit.canTakeMoreGunners()) {

@@ -161,7 +161,7 @@ public boolean hasTransportCapacity(TransporterType transporterType) {
      * @return The current capacity of the transporter, or 0
      */
     @Override
-    public double getCurrentTransportCapacity(TransporterType transporterType) {
+    public double getCurrentTransportCapacity(@Nullable TransporterType transporterType) {
         return transportCapacity.getOrDefault(transporterType, 0.0);
     }
 
@@ -228,6 +228,14 @@ public void clearTransportedUnits() {
         clearTransportedEntities();
     }
 
+    /**
+     * Completely clears the capacity map. Helpful if the transportCapacity has a TransporterType for a Transporter
+     * the unit no longer has - such as after a refit.
+     */
+    public void clearTransportCapacityMap() {
+        transportCapacity = new HashMap<>();
+    }
+
     protected Set<Entity> clearTransportedEntities() {
         Set<Entity> transportedEntities = new HashSet<>();
         if (transport.getEntity() != null) {

@@ -384,6 +384,19 @@ public void initializeAllTransportSpace() {
         }
     }
 
+    public void clearAllTransportSpace() {
+        for (CampaignTransportType campaignTransportType : CampaignTransportType.values()) {
+            clearTransportSpace(campaignTransportType);
+        }
+    }
+
+    private void clearTransportSpace(CampaignTransportType campaignTransportType) {
+        AbstractTransportedUnitsSummary summary = getTransportedUnitsSummary(campaignTransportType);
+        if (summary != null) {
+            summary.clearTransportCapacityMap();
+        }
+    }
+
     private ShipTransportedUnitsSummary getShipTransportedUnitsSummary() {
         return (ShipTransportedUnitsSummary) getTransportedUnitsSummary(SHIP_TRANSPORT);
     }
@@ -445,7 +458,7 @@ private boolean hasTransportedUnitsType(CampaignTransportType campaignTransportT
      *
      * @return transported units summary of that type, or null
      */
-    public AbstractTransportedUnitsSummary getTransportedUnitsSummary(CampaignTransportType campaignTransportType) {
+    public @Nullable AbstractTransportedUnitsSummary getTransportedUnitsSummary(CampaignTransportType campaignTransportType) {
         for (AbstractTransportedUnitsSummary transportedUnitSummary : transportedUnitsSummaries) {
             if (transportedUnitSummary.getClass() == campaignTransportType.getTransportedUnitsSummaryType()) {
                 return transportedUnitSummary;

@@ -121,6 +121,7 @@ private void restoreUnit(Campaign campaign, Unit unit, Entity newEntity) {
         unit.removeParts();
 
         unit.initializeAllTransportSpace();
+        campaign.updateTransportInTransports(unit);
 
         unit.initializeParts(true);
         unit.runDiagnostic(false);

@@ -31,10 +31,8 @@
 import mekhq.campaign.unit.enums.TransporterType;
 import mekhq.campaign.utilities.CampaignTransportUtilities;
 import mekhq.utilities.MHQInternationalization;
-import mekhq.MekHQ;
 import mekhq.campaign.Campaign;
 import mekhq.campaign.enums.CampaignTransportType;
-import mekhq.campaign.event.UnitChangedEvent;
 import mekhq.campaign.unit.Unit;
 
 import javax.swing.*;
@@ -106,15 +104,7 @@ protected void transportMenuAction(ActionEvent evt, TransporterType transporterT
 
         }
         Set<Unit> oldTransports = transport.loadShipTransport(transporterType, units);
-        if (!oldTransports.isEmpty()) {
-            oldTransports.forEach(oldTransport -> campaign.updateTransportInTransports(campaignTransportType, oldTransport));
-            oldTransports.forEach(oldTransport -> MekHQ.triggerEvent(new UnitChangedEvent(transport)));
-        }
-        for (Unit unit : units) {
-            MekHQ.triggerEvent(new UnitChangedEvent(unit));
-        }
-        campaign.updateTransportInTransports(campaignTransportType, transport);
-        MekHQ.triggerEvent(new UnitChangedEvent(transport));
+        updateTransportsForTransportMenuAction(SHIP_TRANSPORT, transport, units, oldTransports);
     }
 
 }
@@ -31,10 +31,8 @@
 import mekhq.campaign.unit.enums.TransporterType;
 import mekhq.campaign.utilities.CampaignTransportUtilities;
 import mekhq.utilities.MHQInternationalization;
-import mekhq.MekHQ;
 import mekhq.campaign.Campaign;
 import mekhq.campaign.enums.CampaignTransportType;
-import mekhq.campaign.event.UnitChangedEvent;
 import mekhq.campaign.unit.Unit;
 
 import javax.swing.*;
@@ -106,14 +104,6 @@ protected void transportMenuAction(ActionEvent evt, TransporterType transporterT
 
         }
         Set<Unit> oldTransports = transport.loadTacticalTransport(transporterType, units);
-        if (!oldTransports.isEmpty()) {
-            oldTransports.forEach(oldTransport -> campaign.updateTransportInTransports(TACTICAL_TRANSPORT, oldTransport));
-            oldTransports.forEach(oldTransport -> MekHQ.triggerEvent(new UnitChangedEvent(transport)));
-        }
-        for (Unit unit : units) {
-            MekHQ.triggerEvent(new UnitChangedEvent(unit));
-        }
-        campaign.updateTransportInTransports(TACTICAL_TRANSPORT, transport);
-        MekHQ.triggerEvent(new UnitChangedEvent(transport));
+        updateTransportsForTransportMenuAction(TACTICAL_TRANSPORT, transport, units, oldTransports);
     }
 }
@@ -28,8 +28,10 @@
 
 package mekhq.gui.menus;
 
+import mekhq.MekHQ;
 import mekhq.campaign.Campaign;
 import mekhq.campaign.enums.CampaignTransportType;
+import mekhq.campaign.event.UnitChangedEvent;
 import mekhq.campaign.force.Force;
 import mekhq.campaign.unit.Unit;
 import mekhq.campaign.unit.enums.TransporterType;
@@ -171,4 +173,25 @@ private Set<JMenuItem> createTransportMenus(TransporterType transporterType, Set
      */
     protected abstract void transportMenuAction(ActionEvent evt, TransporterType transporterType, Unit transport, Set<Unit> units);
 
+    /**
+     * Shared updates used by {@link AssignForceToShipTransportMenu} and {@link AssignForceToTacticalTransportMenu}
+     * @param transport transport (Unit) that has loaded these units
+     * @param units units being assigned to the transport
+     * @param oldTransports transports (Unit) that had previously transported the units
+     */
+    protected void updateTransportsForTransportMenuAction(CampaignTransportType campaignTransportType, Unit transport,
+          Set<Unit> units, Set<Unit> oldTransports) {
+        if (!oldTransports.isEmpty()) {
+            oldTransports.forEach(oldTransport -> {
+                oldTransport.initializeAllTransportSpace();
+                campaign.updateTransportInTransports(campaignTransportType, oldTransport);
+            });
+            oldTransports.forEach(oldTransport -> MekHQ.triggerEvent(new UnitChangedEvent(transport)));
+        }
+        for (Unit unit : units) {
+            MekHQ.triggerEvent(new UnitChangedEvent(unit));
+        }
+        campaign.updateTransportInTransports(campaignTransportType, transport);
+        MekHQ.triggerEvent(new UnitChangedEvent(transport));
+    }
 }