From 94b1dd3923d441f7cac789754e3afc1e08c50c72 Mon Sep 17 00:00:00 2001
From: baldurk <baldurk@baldurk.org>
Date: Wed, 19 May 2021 14:42:43 +0100
Subject: [PATCH] Add ability to update a working expansion struct instead of
 resetting

* This is mostly useful when using a filtering model so we want to save the
  expansion state of nodes that are temporarily filtered out
---
 qrenderdoc/Widgets/Extended/RDTreeView.cpp | 23 ++++++++++++-------
 qrenderdoc/Widgets/Extended/RDTreeView.h   | 26 ++++++++++++++++++----
 2 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/qrenderdoc/Widgets/Extended/RDTreeView.cpp b/qrenderdoc/Widgets/Extended/RDTreeView.cpp
index dc00957a5..6ebe66ca1 100644
--- a/qrenderdoc/Widgets/Extended/RDTreeView.cpp
+++ b/qrenderdoc/Widgets/Extended/RDTreeView.cpp
@@ -505,12 +505,10 @@ void RDTreeView::columnsAboutToBeMoved(const QModelIndex &sourceParent, int sour
   m_currentHoverIndex = QModelIndex();
 }
 
-void RDTreeView::saveExpansion(RDTreeViewExpansionState &state, const ExpansionKeyGen &keygen)
+void RDTreeView::updateExpansion(RDTreeViewExpansionState &state, const ExpansionKeyGen &keygen)
 {
-  state.clear();
-
   for(int i = 0; i < model()->rowCount(); i++)
-    saveExpansionFromRow(state, model()->index(i, 0), 0, keygen);
+    updateExpansionFromRow(state, model()->index(i, 0), 0, keygen);
 }
 
 void RDTreeView::applyExpansion(const RDTreeViewExpansionState &state, const ExpansionKeyGen &keygen)
@@ -587,12 +585,17 @@ void RDTreeView::copySelection()
   clipboard->setText(clipData.trimmed());
 }
 
-void RDTreeView::saveExpansionFromRow(RDTreeViewExpansionState &state, QModelIndex idx, uint seed,
-                                      const ExpansionKeyGen &keygen)
+void RDTreeView::updateExpansionFromRow(RDTreeViewExpansionState &state, QModelIndex idx, uint seed,
+                                        const ExpansionKeyGen &keygen)
 {
   if(!idx.isValid())
     return;
 
+  int rowcount = model()->rowCount(idx);
+
+  if(rowcount == 0)
+    return;
+
   uint key = keygen(idx, seed);
   if(isExpanded(idx))
   {
@@ -601,8 +604,12 @@ void RDTreeView::saveExpansionFromRow(RDTreeViewExpansionState &state, QModelInd
     // only recurse to children if this one is expanded - forget expansion state under collapsed
     // branches. Technically we're losing information here but it allows us to skip a full expensive
     // search
-    for(int i = 0; i < model()->rowCount(idx); i++)
-      saveExpansionFromRow(state, model()->index(i, 0, idx), seed, keygen);
+    for(int i = 0; i < rowcount; i++)
+      updateExpansionFromRow(state, model()->index(i, 0, idx), seed, keygen);
+  }
+  else
+  {
+    state.remove(key);
   }
 }
 
diff --git a/qrenderdoc/Widgets/Extended/RDTreeView.h b/qrenderdoc/Widgets/Extended/RDTreeView.h
index 1836ea953..71ee58d2a 100644
--- a/qrenderdoc/Widgets/Extended/RDTreeView.h
+++ b/qrenderdoc/Widgets/Extended/RDTreeView.h
@@ -123,13 +123,31 @@ public:
   }
   void setModel(QAbstractItemModel *model) override;
 
-  // state is the storage to save the expansion state into
+  // state is the storage to save the expansion state into. The state will be preserved but any rows
+  // which are not processed (either because they don't currently exist in the model or because they
+  // are below a currently collapsed node) will be preserved. Nodes which are known and collapsed
+  // will be removed from the set. This can be useful for preserving expansion across filtering that
+  // may temporarily remove nodes.
   // keygen is a function that will take the index of a row and a previous hash, and return the hash
   //   for that row.
-  void saveExpansion(RDTreeViewExpansionState &state, const ExpansionKeyGen &keygen);
+  void updateExpansion(RDTreeViewExpansionState &state, const ExpansionKeyGen &keygen);
+
+  // similar to updateExpansion but always starts from a clean slate so no previous data is
+  // preserved
+  void saveExpansion(RDTreeViewExpansionState &state, const ExpansionKeyGen &keygen)
+  {
+    state.clear();
+    updateExpansion(state, keygen);
+  }
   void applyExpansion(const RDTreeViewExpansionState &state, const ExpansionKeyGen &keygen);
 
   // convenience overloads for the simple case of using a single column's data as hash
+  void updateExpansion(RDTreeViewExpansionState &state, int keyColumn, int keyRole = Qt::DisplayRole)
+  {
+    updateExpansion(state, [keyColumn, keyRole](QModelIndex idx, uint seed) {
+      return qHash(idx.sibling(idx.row(), keyColumn).data(keyRole).toString(), seed);
+    });
+  }
   void saveExpansion(RDTreeViewExpansionState &state, int keyColumn, int keyRole = Qt::DisplayRole)
   {
     saveExpansion(state, [keyColumn, keyRole](QModelIndex idx, uint seed) {
@@ -198,8 +216,8 @@ private:
 
   QMap<uint, RDTreeViewExpansionState> m_Expansions;
 
-  void saveExpansionFromRow(RDTreeViewExpansionState &state, QModelIndex idx, uint seed,
-                            const ExpansionKeyGen &keygen);
+  void updateExpansionFromRow(RDTreeViewExpansionState &state, QModelIndex idx, uint seed,
+                              const ExpansionKeyGen &keygen);
   void applyExpansionToRow(const RDTreeViewExpansionState &state, QModelIndex idx, uint seed,
                            const ExpansionKeyGen &keygen);