From 5b3a12cd0cb5a20bbaf3a77e39953550257ba4cb Mon Sep 17 00:00:00 2001
From: baldurk <baldurk@baldurk.org>
Date: Fri, 8 Jun 2018 18:42:25 +0100
Subject: [PATCH] Save and restore tree expansion state when repopulating
 locals widget

---
 qrenderdoc/Widgets/Extended/RDTreeWidget.cpp | 22 ++++++++++++++++++++
 qrenderdoc/Widgets/Extended/RDTreeWidget.h   |  7 +++++++
 qrenderdoc/Windows/ShaderViewer.cpp          |  7 +++++++
 3 files changed, 36 insertions(+)

diff --git a/qrenderdoc/Widgets/Extended/RDTreeWidget.cpp b/qrenderdoc/Widgets/Extended/RDTreeWidget.cpp
index b1b277790..829cbbd75 100644
--- a/qrenderdoc/Widgets/Extended/RDTreeWidget.cpp
+++ b/qrenderdoc/Widgets/Extended/RDTreeWidget.cpp
@@ -855,6 +855,28 @@ void RDTreeWidget::setCurrentItem(RDTreeWidgetItem *node)
   setCurrentIndex(m_model->indexForItem(node, 0));
 }
 
+void RDTreeWidget::saveExpansion(RDTreeWidgetExpansionState &state, QString prefix,
+                                 RDTreeWidgetItem *root, int keyColumn)
+{
+  QModelIndex idx = m_model->indexForItem(root, keyColumn);
+  if(isExpanded(idx))
+    state.insert(prefix);
+
+  for(int i = 0; i < root->childCount(); i++)
+    saveExpansion(state, prefix + root->child(i)->text(keyColumn), root->child(i), keyColumn);
+}
+
+void RDTreeWidget::applySavedExpansion(RDTreeWidgetExpansionState &state, QString prefix,
+                                       RDTreeWidgetItem *root, int keyColumn)
+{
+  QModelIndex idx = m_model->indexForItem(root, keyColumn);
+  if(state.contains(prefix))
+    expand(idx);
+
+  for(int i = 0; i < root->childCount(); i++)
+    applySavedExpansion(state, prefix + root->child(i)->text(keyColumn), root->child(i), keyColumn);
+}
+
 RDTreeWidgetItem *RDTreeWidget::itemAt(const QPoint &p) const
 {
   return m_model->itemForIndex(indexAt(p));
diff --git a/qrenderdoc/Widgets/Extended/RDTreeWidget.h b/qrenderdoc/Widgets/Extended/RDTreeWidget.h
index 0587aa838..edfad33de 100644
--- a/qrenderdoc/Widgets/Extended/RDTreeWidget.h
+++ b/qrenderdoc/Widgets/Extended/RDTreeWidget.h
@@ -29,6 +29,8 @@
 class RDTreeWidget;
 class RDTreeWidgetModel;
 
+typedef QSet<QString> RDTreeWidgetExpansionState;
+
 class RDTreeWidgetItem
 {
 public:
@@ -262,6 +264,11 @@ public:
   void setSelectedItem(RDTreeWidgetItem *node);
   void setCurrentItem(RDTreeWidgetItem *node);
 
+  void saveExpansion(RDTreeWidgetExpansionState &state, QString prefix, RDTreeWidgetItem *root,
+                     int keyColumn);
+  void applySavedExpansion(RDTreeWidgetExpansionState &state, QString prefix,
+                           RDTreeWidgetItem *root, int keyColumn);
+
   RDTreeWidgetItem *itemAt(const QPoint &p) const;
   RDTreeWidgetItem *itemAt(int x, int y) const { return itemAt(QPoint(x, y)); }
   void expandItem(RDTreeWidgetItem *item);
diff --git a/qrenderdoc/Windows/ShaderViewer.cpp b/qrenderdoc/Windows/ShaderViewer.cpp
index 46201aed6..683340327 100644
--- a/qrenderdoc/Windows/ShaderViewer.cpp
+++ b/qrenderdoc/Windows/ShaderViewer.cpp
@@ -1723,6 +1723,11 @@ void ShaderViewer::updateDebugging()
 
   if(m_Trace->hasLocals)
   {
+    const QString stateprefix = lit("!!@");
+
+    RDTreeWidgetExpansionState expansion;
+    ui->locals->saveExpansion(expansion, stateprefix, ui->locals->invisibleRootItem(), 0);
+
     ui->locals->clear();
 
     const QString xyzw = lit("xyzw");
@@ -1851,6 +1856,8 @@ void ShaderViewer::updateDebugging()
 
     while(fakeroot.childCount() > 0)
       ui->locals->addTopLevelItem(fakeroot.takeChild(0));
+
+    ui->locals->applySavedExpansion(expansion, stateprefix, ui->locals->invisibleRootItem(), 0);
   }
 
   if(ui->registers->topLevelItemCount() == 0)