From 7076f54b6ba28478e265db693a9d956a5dfd9e93 Mon Sep 17 00:00:00 2001
From: Sergei Tachenov <sergei.tachenov@jetbrains.com>
Date: Wed, 30 Aug 2023 14:29:25 +0300
Subject: [PATCH] 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
---
 .../ui/DependenciesPanel.java                 |  2 +-
 .../ui/TreeExpansionMonitor.java              | 72 +++++++++++++++++++
 2 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/platform/lang-impl/src/com/intellij/packageDependencies/ui/DependenciesPanel.java b/platform/lang-impl/src/com/intellij/packageDependencies/ui/DependenciesPanel.java
index 5754e04111c2..449c96eddf0d 100644
--- a/platform/lang-impl/src/com/intellij/packageDependencies/ui/DependenciesPanel.java
+++ b/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) {
diff --git a/platform/platform-impl/src/com/intellij/packageDependencies/ui/TreeExpansionMonitor.java b/platform/platform-impl/src/com/intellij/packageDependencies/ui/TreeExpansionMonitor.java
index 9cc090c01cad..55eaa6f76ff1 100644
--- a/platform/platform-impl/src/com/intellij/packageDependencies/ui/TreeExpansionMonitor.java
+++ b/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;
+        }
+      }
+    }
+  }
 }