Merge pull request #961 from SJuliez/baytreepersistent

Preserve expanded state of Large Aero Weapon Bays
MegaMek · Nov 16, 2021 · e4e28f8 · e4e28f8
2 parents 59087f0 + 418451e
commit e4e28f8
Showing 1 changed file with 41 additions and 10 deletions.
diff --git a/src/megameklab/com/ui/util/BayWeaponCriticalTree.java b/src/megameklab/com/ui/util/BayWeaponCriticalTree.java
@@ -22,15 +22,7 @@
 import java.awt.event.MouseAdapter;
 import java.awt.event.MouseEvent;
 import java.awt.event.MouseListener;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.Enumeration;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Optional;
-import java.util.Set;
-import java.util.StringJoiner;
-import java.util.Vector;
+import java.util.*;
 
 import javax.swing.JLabel;
 import javax.swing.JMenu;
@@ -90,6 +82,10 @@ public class BayWeaponCriticalTree extends JTree {
     private final EntitySource eSource;
     private final DefaultTreeModel model;
     private RefreshListener refresh;
+
+    /** Stores a unique transient ID for each Bay to allow restoring the expanded state when the loadout changes. */
+    private final Map<Mounted, Integer> bayIdMap = new HashMap<>();
+    private int bayIdCounter;
 
     public BayWeaponCriticalTree(int location, EntitySource eSource, RefreshListener refresh) {
         this(location, eSource, refresh, FORWARD);
@@ -129,9 +125,11 @@ public void updateRefresh(RefreshListener refresh) {
     }
 
     public void rebuild() {
+        List<Integer> expandedBays = getExpandedBayIds();
         TreeNode root = initRoot();
         model.setRoot(root);
         setRootVisible(root.getChildCount() == 0);
+        restoreExpandedBays(expandedBays);
     }
 
     private int slotCount(Mounted bay) {
@@ -608,6 +606,7 @@ private class BayNode extends EquipmentNode {
 
         BayNode(Mounted object) {
             super(object);
+            bayIdMap.computeIfAbsent(object, m -> bayIdCounter++);
         }
 
         public boolean isCapital() {
@@ -1498,4 +1497,36 @@ public void removeExported(String selection, int action) {
         refresh.refreshStatus();
         refresh.refreshSummary();
     }
-}
+
+    private void restoreExpandedBays(List<Integer> expandedBayIds) {
+        TreeNode root = (TreeNode) model.getRoot();
+        if ((root == null) || (expandedBayIds == null)) {
+            return;
+        }
+        for (int n = 0; n < model.getChildCount(root); n++) {
+            Object node = model.getChild(root, n);
+            if (node instanceof BayNode) {
+                if (expandedBayIds.contains(bayIdMap.get(((BayNode) node).getMounted()))) {
+                    Object[] pathObjs = new Object[2];
+                    pathObjs[0] = root;
+                    pathObjs[1] = node;
+                    expandPath(new TreePath(pathObjs));
+                }
+            }
+        }
+    }
+
+    private List<Integer> getExpandedBayIds() {
+        List<Integer> result = new ArrayList<>();
+        for (int i = 0; i < getRowCount(); i++) {
+            TreePath currPath = getPathForRow(i);
+            if (isExpanded(currPath)) {
+                Object entry = currPath.getLastPathComponent();
+                if (entry instanceof BayNode) {
+                    result.add(bayIdMap.get(((BayNode) entry).getMounted()));
+                }
+            }
+        }
+        return result;
+    }
+}