IDEA-319883 TreeExpansionMonitor.restoreAsync

Use TreeUtil async API to restore the tree because with AsyncTreeModel
it's impossible to do otherwise: the model is simply not initialized straight away.

Since we only do it on a full model rebuild, cancelling the previous restore
is not necessary, as it's cancelled when the previous model is disposed.
However, we still check if the model is the same, as otherwise we don't
need to do extra work and also don't need to unfreeze at the end because
there's another restoreAsync running.

GitOrigin-RevId: bad47ce7e4a9a20e0734d9f4d6a7ea39592920a5

platform/lang-impl/src/com/intellij/packageDependencies/ui/DependenciesPanel.java

@@ -374,7 +374,7 @@ public final class DependenciesPanel extends JPanel implements Disposable, DataP
       Disposer.dispose(disposable);
     }
     tree.setModel(model);
-    monitor.restore();
+    monitor.restoreAsync();
   }
 
   private ActionGroup createTreePopupActions(boolean isRightTree) {

platform/platform-impl/src/com/intellij/packageDependencies/ui/TreeExpansionMonitor.java

@@ -2,7 +2,9 @@
 package com.intellij.packageDependencies.ui;
 
 import com.intellij.openapi.util.Comparing;
+import com.intellij.ui.tree.TreeVisitor;
 import com.intellij.util.ui.tree.TreeUtil;
+import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
 
 import javax.swing.*;
@@ -11,10 +13,12 @@ import javax.swing.event.TreeExpansionListener;
 import javax.swing.event.TreeSelectionEvent;
 import javax.swing.event.TreeSelectionListener;
 import javax.swing.tree.DefaultMutableTreeNode;
+import javax.swing.tree.TreeModel;
 import javax.swing.tree.TreeNode;
 import javax.swing.tree.TreePath;
 import java.util.*;
 import java.util.function.BiPredicate;
+import java.util.stream.Collectors;
 
 import static com.intellij.ui.tree.TreePathUtil.toTreePathArray;
 
@@ -111,6 +115,11 @@ public abstract class TreeExpansionMonitor<T> {
     myFrozen = false;
   }
 
+  @ApiStatus.Internal
+  public void restoreAsync() {
+    new AsyncRestorer().restore();
+  }
+
   protected abstract TreePath findPathByNode(final T node);
 
   public boolean isFreeze() {
@@ -121,4 +130,67 @@ public abstract class TreeExpansionMonitor<T> {
   private T getLastNode(@NotNull TreePath path) {
     return (T)path.getLastPathComponent();
   }
+
+  private class AsyncRestorer {
+
+    private final TreeModel initialModel = myTree.getModel();
+
+    private boolean isValidModel() {
+      return myTree.getModel() == initialModel;
+    }
+
+    void restore() {
+      freeze();
+      var nodesToExpand = myExpandedPaths.stream().map(TreeExpansionMonitor.this::getLastNode).collect(Collectors.toUnmodifiableSet());
+      var nodesToSelect = mySelectionNodes.stream().collect(Collectors.toUnmodifiableSet());
+      var pathsToExpand = new HashSet<TreePath>();
+      var pathsToSelect = new HashSet<TreePath>();
+      TreeUtil.promiseVisit(myTree, path -> {
+        if (!isValidModel()) {
+          return TreeVisitor.Action.INTERRUPT;
+        }
+        var node = getLastNode(path);
+        if (nodesToExpand.contains(node)) {
+          pathsToExpand.add(path);
+        }
+        if (nodesToSelect.contains(node)) {
+          pathsToSelect.add(path);
+        }
+        return TreeVisitor.Action.CONTINUE;
+      }).onSuccess(ignored -> {
+        TreeUtil.promiseExpand(myTree, pathsToExpand.stream().map(PathVisitor::new));
+      }).onSuccess(ignored -> {
+        TreeUtil.promiseSelect(myTree, pathsToSelect.stream().map(PathVisitor::new));
+      }).onProcessed(ignored -> {
+        if (isValidModel()) { // otherwise, there's another restoreAsync() running for the new model, it's responsible for unfreezing
+          unfreeze();
+        }
+      });
+    }
+
+    private class PathVisitor implements TreeVisitor {
+
+      private final @NotNull TreePath pathToActUpon;
+
+      PathVisitor(@NotNull TreePath pathToActUpon) {
+        this.pathToActUpon = pathToActUpon;
+      }
+
+      @Override
+      public @NotNull Action visit(@NotNull TreePath path) {
+        if (!isValidModel()) {
+          return Action.SKIP_SIBLINGS;
+        }
+        if (path.equals(pathToActUpon)) {
+          return Action.INTERRUPT;
+        }
+        else if (path.isDescendant(pathToActUpon)) {
+          return Action.CONTINUE;
+        }
+        else {
+          return Action.SKIP_CHILDREN;
+        }
+      }
+    }
+  }
 }