From db4cb88dc6beb290a9e711f1d048fb62890c722e Mon Sep 17 00:00:00 2001
From: Aleksey Pivovarov <AMPivovarov@gmail.com>
Date: Mon, 13 Mar 2023 17:44:36 +0100
Subject: [PATCH] vcs: migrate GitStageTree to AsyncChangesTree

GitOrigin-RevId: 38cb861dff4d504ff88cb6289c97368f5ad944db
---
 .../src/git4idea/index/ui/GitStagePanel.kt    |  8 ++--
 .../src/git4idea/index/ui/GitStageTree.kt     | 42 ++++++++++---------
 2 files changed, 28 insertions(+), 22 deletions(-)

diff --git a/plugins/git4idea/src/git4idea/index/ui/GitStagePanel.kt b/plugins/git4idea/src/git4idea/index/ui/GitStagePanel.kt
index 6f30463131ed..6ebaf1ef0987 100644
--- a/plugins/git4idea/src/git4idea/index/ui/GitStagePanel.kt
+++ b/plugins/git4idea/src/git4idea/index/ui/GitStagePanel.kt
@@ -311,9 +311,11 @@ internal class GitStagePanel(private val tracker: GitStageTracker,
     fun editedCommitChanged() {
       rebuildTree()
 
-      commitPanel.editedCommit?.let {
-        val node = TreeUtil.findNodeWithObject(root, it)
-        node?.let { expandPath(TreeUtil.getPathFromRoot(node)) }
+      invokeAfterRefresh {
+        commitPanel.editedCommit?.let {
+          val node = TreeUtil.findNodeWithObject(root, it)
+          node?.let { expandPath(TreeUtil.getPathFromRoot(node)) }
+        }
       }
     }
 
diff --git a/plugins/git4idea/src/git4idea/index/ui/GitStageTree.kt b/plugins/git4idea/src/git4idea/index/ui/GitStageTree.kt
index 4270dd4da99d..c9d46e882240 100644
--- a/plugins/git4idea/src/git4idea/index/ui/GitStageTree.kt
+++ b/plugins/git4idea/src/git4idea/index/ui/GitStageTree.kt
@@ -45,7 +45,7 @@ import javax.swing.tree.DefaultTreeModel
 abstract class GitStageTree(project: Project,
                             private val settings: GitStageUiSettings,
                             parentDisposable: Disposable) :
-  ChangesTree(project, false, true) {
+  AsyncChangesTree(project, false, true) {
 
   protected abstract val state: GitStageTracker.State
   protected abstract val ignoredFilePaths: Map<VirtualFile, List<FilePath>>
@@ -88,26 +88,30 @@ abstract class GitStageTree(project: Project,
 
   protected abstract fun createHoverIcon(node: ChangesBrowserGitFileStatusNode): HoverIcon?
 
-  override fun rebuildTree() {
-    val builder = MyTreeModelBuilder(myProject, groupingSupport.grouping)
-
-    builder.createKindNode(NodeKind.STAGED)
-    builder.createKindNode(NodeKind.UNSTAGED)
-
-    state.forEachStatus(*NodeKind.values()) { root, status, kind ->
-      builder.insertStatus(root, status, kind)
-    }
-
-    if (settings.ignoredFilesShown()) {
-      builder.insertIgnoredPaths(ignoredFilePaths)
-    }
-
-    customizeTreeModel(builder)
-    updateTreeModel(builder.build())
-  }
-
   protected open fun customizeTreeModel(builder: TreeModelBuilder) = Unit
 
+  override val changesTreeModel: AsyncChangesTreeModel = GitStateTreeModel()
+
+  private inner class GitStateTreeModel : SimpleAsyncChangesTreeModel() {
+    override fun buildTreeModelSync(grouping: ChangesGroupingPolicyFactory): DefaultTreeModel {
+      val builder = MyTreeModelBuilder(myProject, grouping)
+
+      builder.createKindNode(NodeKind.STAGED)
+      builder.createKindNode(NodeKind.UNSTAGED)
+
+      state.forEachStatus(*NodeKind.values()) { root, status, kind ->
+        builder.insertStatus(root, status, kind)
+      }
+
+      if (settings.ignoredFilesShown()) {
+        builder.insertIgnoredPaths(ignoredFilePaths)
+      }
+
+      customizeTreeModel(builder)
+      return builder.build()
+    }
+  }
+
   override fun getData(dataId: String): Any? {
     return when {
       GitStageDataKeys.GIT_STAGE_TREE.`is`(dataId) -> this