IDEA-138983 - Container nodes can not preserve collapsed state

platform/remote-servers/impl/src/com/intellij/remoteServer/impl/runtime/ui/ServersToolWindowContent.java

@@ -67,6 +67,7 @@ public class ServersToolWindowContent extends JPanel implements Disposable {
   private final DefaultTreeModel myTreeModel;
   private TreeBuilderBase myBuilder;
   private AbstractTreeNode<?> myLastSelection;
+  private Set<Object> myCollapsedTreeNodeValues = new HashSet<Object>();
 
   private final Project myProject;
 
@@ -113,6 +114,36 @@ public class ServersToolWindowContent extends JPanel implements Disposable {
         return true;
       }
     }.installOn(myTree);
+    myTree.addTreeExpansionListener(new TreeExpansionListener() {
+
+      @Override
+      public void treeExpanded(TreeExpansionEvent event) {
+        Object value = getNodeValue(event);
+        if (value != null) {
+          myCollapsedTreeNodeValues.remove(value);
+        }
+      }
+
+      @Override
+      public void treeCollapsed(TreeExpansionEvent event) {
+        Object value = getNodeValue(event);
+        if (value != null) {
+          myCollapsedTreeNodeValues.add(value);
+        }
+      }
+
+      private Object getNodeValue(TreeExpansionEvent event) {
+        DefaultMutableTreeNode treeNode = ObjectUtils.tryCast(event.getPath().getLastPathComponent(), DefaultMutableTreeNode.class);
+        if (treeNode == null) {
+          return null;
+        }
+        AbstractTreeNode nodeDescriptor = ObjectUtils.tryCast(treeNode.getUserObject(), AbstractTreeNode.class);
+        if (nodeDescriptor == null) {
+          return null;
+        }
+        return nodeDescriptor.getValue();
+      }
+    });
 
     DefaultActionGroup popupActionGroup = new DefaultActionGroup();
     popupActionGroup.add(ActionManager.getInstance().getAction(SERVERS_TOOL_WINDOW_TOOLBAR));
@@ -183,9 +214,10 @@ public class ServersToolWindowContent extends JPanel implements Disposable {
     myBuilder = new TreeBuilderBase(myTree, structure, myTreeModel) {
       @Override
       protected boolean isAutoExpandNode(NodeDescriptor nodeDescriptor) {
-        return nodeDescriptor instanceof ServersTreeStructure.RemoteServerNode
-               || nodeDescriptor instanceof ServersTreeStructure.DeploymentNodeImpl
-               || nodeDescriptor instanceof ServersTreeStructure.GroupNode;
+        return (nodeDescriptor instanceof ServersTreeStructure.RemoteServerNode
+                || nodeDescriptor instanceof ServersTreeStructure.DeploymentNodeImpl
+                || nodeDescriptor instanceof ServersTreeStructure.GroupNode)
+               && (!myCollapsedTreeNodeValues.contains(((AbstractTreeNode)nodeDescriptor).getValue()));
       }
     };
     Disposer.register(this, myBuilder);

platform/remote-servers/impl/src/com/intellij/remoteServer/impl/runtime/ui/tree/ServersTreeStructure.java

@@ -43,8 +43,6 @@ import com.intellij.remoteServer.runtime.deployment.DeploymentTask;
 import com.intellij.ui.LayeredIcon;
 import com.intellij.ui.awt.RelativePoint;
 import com.intellij.util.containers.ContainerUtil;
-import com.intellij.util.containers.hash.*;
-import com.intellij.util.containers.hash.HashSet;
 import icons.RemoteServersIcons;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -52,7 +50,6 @@ import org.jetbrains.annotations.Nullable;
 import javax.swing.*;
 import java.awt.event.MouseEvent;
 import java.util.*;
-import java.util.HashMap;
 
 /**
  * @author michael.golubev
@@ -64,6 +61,9 @@ public class ServersTreeStructure extends AbstractTreeStructureBase {
   private final ServersTreeRootNode myRootElement;
   private final Project myProject;
 
+  private final Map<RemoteServer, Map<String, DeploymentGroup>> myServer2DeploymentGroups
+    = new HashMap<RemoteServer, Map<String, DeploymentGroup>>();
+
   public ServersTreeStructure(@NotNull Project project) {
     super(project);
     myProject = project;
@@ -154,15 +154,38 @@ public class ServersTreeStructure extends AbstractTreeStructureBase {
         return Collections.emptyList();
       }
 
+      Map<DeploymentGroup, GroupNode> group2node = new HashMap<DeploymentGroup, GroupNode>();
       final List<AbstractTreeNode> children = new ArrayList<AbstractTreeNode>();
-      Set<String> groups = new HashSet<String>();
       for (Deployment deployment : connection.getDeployments()) {
         final String groupName = deployment.getGroup();
         if (groupName == null) {
           children.add(new DeploymentNodeImpl(connection, this, deployment));
         }
-        else if (groups.add(groupName)) {
-          children.add(new GroupNode(connection, this, groupName));
+        else {
+          Map<String, DeploymentGroup> groups
+            = ContainerUtil.getOrCreate(myServer2DeploymentGroups, getServer(), new Factory<Map<String, DeploymentGroup>>() {
+            @Override
+            public Map<String, DeploymentGroup> create() {
+              return new HashMap<String, DeploymentGroup>();
+            }
+          });
+
+          final DeploymentGroup group
+            = ContainerUtil.getOrCreate(groups, groupName, new Factory<DeploymentGroup>() {
+            @Override
+            public DeploymentGroup create() {
+              return new DeploymentGroup(groupName);
+            }
+          });
+
+          ContainerUtil.getOrCreate(group2node, group, new Factory<GroupNode>() {
+            @Override
+            public GroupNode create() {
+              GroupNode result = new GroupNode(connection, RemoteServerNode.this, group);
+              children.add(result);
+              return result;
+            }
+          });
         }
       }
       return children;
@@ -447,19 +470,32 @@ public class ServersTreeStructure extends AbstractTreeStructureBase {
     }
   }
 
-  public class GroupNode extends AbstractTreeNode<String> implements ServersTreeNode {
+  private static class DeploymentGroup {
+
+    private final String myName;
+
+    private DeploymentGroup(String name) {
+      myName = name;
+    }
+
+    public String getName() {
+      return myName;
+    }
+  }
+
+  public class GroupNode extends AbstractTreeNode<DeploymentGroup> implements ServersTreeNode {
 
     @NotNull private final ServerConnection<?> myConnection;
     @NotNull private final RemoteServerNode myServerNode;
 
-    public GroupNode(@NotNull ServerConnection<?> connection, @NotNull RemoteServerNode serverNode, @NotNull String group) {
+    public GroupNode(@NotNull ServerConnection<?> connection, @NotNull RemoteServerNode serverNode, @NotNull DeploymentGroup group) {
       super(doGetProject(), group);
       myConnection = connection;
       myServerNode = serverNode;
     }
 
     @NotNull
-    public String getGroup() {
+    public DeploymentGroup getGroup() {
       return getValue();
     }
 
@@ -468,7 +504,7 @@ public class ServersTreeStructure extends AbstractTreeStructureBase {
     public Collection<? extends AbstractTreeNode> getChildren() {
       List<AbstractTreeNode> children = new ArrayList<AbstractTreeNode>();
       for (Deployment deployment : myConnection.getDeployments()) {
-        if (StringUtil.equals(getGroup(), deployment.getGroup())) {
+        if (StringUtil.equals(getGroup().getName(), deployment.getGroup())) {
           children.add(new DeploymentNodeImpl(myConnection, myServerNode, deployment));
         }
       }
@@ -478,7 +514,7 @@ public class ServersTreeStructure extends AbstractTreeStructureBase {
     @Override
     protected void update(PresentationData presentation) {
       presentation.setIcon(myServerNode.getServer().getType().getIcon());
-      presentation.setPresentableText(getGroup());
+      presentation.setPresentableText(getGroup().getName());
     }
   }
 }