correctly restore expanded and selected state of nodes in UsageView

UsageView/src/com/intellij/usages/impl/UsageViewImpl.java

@@ -29,6 +29,7 @@ import com.intellij.util.ui.tree.TreeUtil;
 import javax.swing.*;
 import javax.swing.tree.DefaultMutableTreeNode;
 import javax.swing.tree.TreePath;
+import javax.swing.tree.TreeNode;
 import java.awt.*;
 import java.awt.event.ActionEvent;
 import java.awt.event.ActionListener;
@@ -239,6 +240,8 @@ public class UsageViewImpl implements UsageView {
   }
 
   public void rulesChanged() {
+    final ArrayList<UsageState> states = new ArrayList<UsageState>();
+    captureUsagesExpandState(new TreePath(myTree.getModel().getRoot()), states);
     Collection<Usage> allUsages = myUsageNodes.keySet();
     reset();
     myBuilder.setGroupingRules(getGroupingRuleProvider().getActiveRules(myProject));
@@ -250,6 +253,35 @@ public class UsageViewImpl implements UsageView {
       }
       appendUsage(usage);
     }
+    restoreUsageExpandState(states);
+  }
+
+  private void captureUsagesExpandState(TreePath pathFrom, final Collection<UsageState> states) {
+    if (!myTree.isExpanded(pathFrom)) {
+      return;
+    }
+    final DefaultMutableTreeNode node = (DefaultMutableTreeNode)pathFrom.getLastPathComponent();
+    final int childCount = node.getChildCount();
+    for (int idx = 0; idx < childCount; idx++) {
+      final TreeNode child = node.getChildAt(idx);
+      if (child instanceof UsageNode) {
+        final Usage usage = ((UsageNode)child).getUsage();
+        if (usage != null) {
+          states.add(new UsageState(usage, myTree.getSelectionModel().isPathSelected(pathFrom.pathByAddingChild(child))));
+        }
+      }
+      else {
+        captureUsagesExpandState(pathFrom.pathByAddingChild(child), states);
+      }
+    }
+  }
+
+  private void restoreUsageExpandState(final Collection<UsageState> states) {
+    myTree.getSelectionModel().clearSelection();
+    for (Iterator<UsageState> it = states.iterator(); it.hasNext();) {
+      final UsageState usageState = it.next();
+      usageState.restore();
+    }
   }
 
   private void expandAll() {
@@ -376,13 +408,6 @@ public class UsageViewImpl implements UsageView {
 
   public void appendUsage(Usage usage) {
     final UsageNode node = myBuilder.appendUsage(usage);
-    if (node != null) {
-      if (myUsageNodes.size() == 0) { //first usage
-        TreePath usagePath = new TreePath(node.getPath());
-        myTree.expandPath(usagePath.getParentPath());
-        myTree.setSelectionPath(usagePath);
-      }
-    }
     myUsageNodes.put(usage, node);
   }
 
@@ -457,8 +482,20 @@ public class UsageViewImpl implements UsageView {
     SwingUtilities.invokeLater(new Runnable() {
       public void run() {
         flush();
+        showFirstUsage();
       }
     });
+
+  }
+
+  private void showFirstUsage() {
+    UsageViewTreeModelBuilder model = (UsageViewTreeModelBuilder)myTree.getModel();
+    final UsageNode firstUsageNode = model.getFirstUsageNode();
+    if (firstUsageNode != null) { //first usage;
+      TreePath usagePath = new TreePath(firstUsageNode.getPath());
+      myTree.expandPath(usagePath.getParentPath());
+      myTree.setSelectionPath(usagePath);
+    }
   }
 
   public void addButtonToLowerPane(final Runnable runnable, String text, char mnemonic) {
@@ -795,4 +832,33 @@ public class UsageViewImpl implements UsageView {
       }
     }
   }
+
+  private class UsageState {
+    private final Usage myUsage;
+    private final boolean mySelected;
+
+    public UsageState(final Usage usage) {
+      this(usage, false);
+    }
+
+    public UsageState(final Usage usage, boolean isSelected) {
+      myUsage = usage;
+      mySelected = isSelected;
+    }
+
+    public void restore() {
+      final UsageNode node = myUsageNodes.get(myUsage);
+      if (node == null) {
+        return;
+      }
+      final DefaultMutableTreeNode parentGroupingNode = (DefaultMutableTreeNode)node.getParent();
+      if (parentGroupingNode != null) {
+        final TreePath treePath = new TreePath(parentGroupingNode.getPath());
+        myTree.expandPath(treePath);
+        if (mySelected) {
+          myTree.addSelectionPath(treePath.pathByAddingChild(node));
+        }
+      }
+    }
+  }
 }

UsageView/src/com/intellij/usages/impl/UsageViewTreeCellRenderer.java

@@ -93,7 +93,7 @@ class UsageViewTreeCellRenderer extends ColoredTreeCellRenderer {
         append((String)userObject, SimpleTextAttributes.REGULAR_BOLD_ATTRIBUTES);
       }
       else {
-        append(userObject.toString(), SimpleTextAttributes.REGULAR_ATTRIBUTES);
+        append((userObject != null) ? userObject.toString() : "", SimpleTextAttributes.REGULAR_ATTRIBUTES);
       }
     }
     else {

UsageView/src/com/intellij/usages/impl/UsageViewTreeModelBuilder.java

@@ -5,6 +5,7 @@ import com.intellij.usages.UsageViewPresentation;
 
 import javax.swing.tree.DefaultMutableTreeNode;
 import javax.swing.tree.DefaultTreeModel;
+import javax.swing.tree.TreeNode;
 
 public class UsageViewTreeModelBuilder extends DefaultTreeModel {
   private final RootGroupNode myRootNode;
@@ -36,6 +37,25 @@ public class UsageViewTreeModelBuilder extends DefaultTreeModel {
     myRootNode.addNode(myTargetsNode);
   }
 
+  public UsageNode getFirstUsageNode() {
+    return (UsageNode)getFirstChildOfType(myRootNode, UsageNode.class);
+  }
+
+  private static TreeNode getFirstChildOfType(TreeNode parent, final Class type) {
+    final int childCount = parent.getChildCount();
+    for (int idx = 0; idx < childCount; idx++) {
+      final TreeNode child = parent.getChildAt(idx);
+      if (type.isAssignableFrom(child.getClass())) {
+        return child;
+      }
+      final TreeNode firstChildOfType = getFirstChildOfType(child, type);
+      if (firstChildOfType != null) {
+        return firstChildOfType;
+      }
+    }
+    return null;
+  }
+
   public boolean areTargetsValid() {
     if (myTargets.length == 0) return true;
     for (int i = 0; i < myTargetNodes.length; i++) {
@@ -57,7 +77,7 @@ public class UsageViewTreeModelBuilder extends DefaultTreeModel {
     }
 
     public void addNode(DefaultMutableTreeNode node) {
-      UsageViewTreeModelBuilder.this.insertNodeInto(node, this, getNodeInsertionIndex(node));
+      myTreeModel.insertNodeInto(node, this, getNodeInsertionIndex(node));
     }
   }
 }