diff --git a/qrenderdoc/Widgets/Extended/RDTreeView.cpp b/qrenderdoc/Widgets/Extended/RDTreeView.cpp
index f165f7c5d..549b6a2aa 100644
--- a/qrenderdoc/Widgets/Extended/RDTreeView.cpp
+++ b/qrenderdoc/Widgets/Extended/RDTreeView.cpp
@@ -241,6 +241,58 @@ void RDTreeView::rowsAboutToBeRemoved(const QModelIndex &parent, int start, int
   m_currentHoverIndex = QModelIndex();
 }
 
+void RDTreeView::saveExpansion(RDTreeViewExpansionState &state, QString prefix, int keyColumn,
+                               int role)
+{
+  state.clear();
+
+  for(int i = 0; i < model()->rowCount(); i++)
+    saveExpansion(state, model()->index(i, keyColumn), qHash(prefix), keyColumn, role);
+}
+
+void RDTreeView::applySavedExpansion(const RDTreeViewExpansionState &state, QString prefix,
+                                     int keyColumn, int role)
+{
+  for(int i = 0; i < model()->rowCount(); i++)
+    applySavedExpansion(state, model()->index(i, keyColumn), qHash(prefix), keyColumn, role);
+}
+
+void RDTreeView::saveExpansion(RDTreeViewExpansionState &state, QModelIndex idx, uint seed,
+                               int keyColumn, int role)
+{
+  if(!idx.isValid())
+    return;
+
+  uint key = qHash(model()->data(idx, role).toString(), seed);
+  if(isExpanded(idx))
+  {
+    state.insert(key);
+
+    // 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++)
+      saveExpansion(state, model()->index(i, keyColumn, idx), key, keyColumn, role);
+  }
+}
+
+void RDTreeView::applySavedExpansion(const RDTreeViewExpansionState &state, QModelIndex idx,
+                                     uint seed, int keyColumn, int role)
+{
+  if(!idx.isValid())
+    return;
+
+  uint key = qHash(model()->data(idx, role).toString(), seed);
+  if(state.contains(key))
+  {
+    expand(idx);
+
+    // same as above - only recurse when we have a parent that's expanded.
+    for(int i = 0; i < model()->rowCount(idx); i++)
+      applySavedExpansion(state, model()->index(i, keyColumn, idx), key, keyColumn, role);
+  }
+}
+
 void RDTreeView::drawRow(QPainter *painter, const QStyleOptionViewItem &options,
                          const QModelIndex &index) const
 {
diff --git a/qrenderdoc/Widgets/Extended/RDTreeView.h b/qrenderdoc/Widgets/Extended/RDTreeView.h
index 178cd06dd..059c98659 100644
--- a/qrenderdoc/Widgets/Extended/RDTreeView.h
+++ b/qrenderdoc/Widgets/Extended/RDTreeView.h
@@ -31,6 +31,8 @@
 
 class RDTreeView;
 
+typedef QSet<uint> RDTreeViewExpansionState;
+
 class RDTreeViewDelegate : public ForwardingDelegate
 {
 private:
@@ -94,6 +96,11 @@ public:
   void setItemDelegate(QAbstractItemDelegate *delegate);
   QAbstractItemDelegate *itemDelegate() const;
 
+  void saveExpansion(RDTreeViewExpansionState &state, QString prefix, int keyColumn,
+                     int role = Qt::DisplayRole);
+  void applySavedExpansion(const RDTreeViewExpansionState &state, QString prefix, int keyColumn,
+                           int role = Qt::DisplayRole);
+
 signals:
   void leave(QEvent *e);
   void keyPress(QKeyEvent *e);
@@ -119,6 +126,11 @@ private:
   bool m_VisibleGridLines = true;
   bool m_TooltipElidedItems = true;
 
+  void saveExpansion(RDTreeViewExpansionState &state, QModelIndex idx, uint seed, int keyColumn,
+                     int role);
+  void applySavedExpansion(const RDTreeViewExpansionState &state, QModelIndex idx, uint seed,
+                           int keyColumn, int role);
+
   QAbstractItemDelegate *m_userDelegate = NULL;
   RDTreeViewDelegate *m_delegate;
 
diff --git a/qrenderdoc/Widgets/Extended/RDTreeWidget.cpp b/qrenderdoc/Widgets/Extended/RDTreeWidget.cpp
index 4648cc07f..240ed3a6f 100644
--- a/qrenderdoc/Widgets/Extended/RDTreeWidget.cpp
+++ b/qrenderdoc/Widgets/Extended/RDTreeWidget.cpp
@@ -855,28 +855,6 @@ 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 526e5f835..bbf828a36 100644
--- a/qrenderdoc/Widgets/Extended/RDTreeWidget.h
+++ b/qrenderdoc/Widgets/Extended/RDTreeWidget.h
@@ -29,8 +29,6 @@
 class RDTreeWidget;
 class RDTreeWidgetModel;
 
-typedef QSet<QString> RDTreeWidgetExpansionState;
-
 class RDTreeWidgetItem
 {
 public:
@@ -266,11 +264,6 @@ 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 17e5b5b9a..f70ab9422 100644
--- a/qrenderdoc/Windows/ShaderViewer.cpp
+++ b/qrenderdoc/Windows/ShaderViewer.cpp
@@ -1898,8 +1898,8 @@ void ShaderViewer::updateDebugging()
   {
     const QString stateprefix = lit("!!@");
 
-    RDTreeWidgetExpansionState expansion;
-    ui->locals->saveExpansion(expansion, stateprefix, ui->locals->invisibleRootItem(), 0);
+    RDTreeViewExpansionState expansion;
+    ui->locals->saveExpansion(expansion, stateprefix, 0);
 
     ui->locals->clear();
 
@@ -2040,7 +2040,7 @@ void ShaderViewer::updateDebugging()
     while(fakeroot.childCount() > 0)
       ui->locals->addTopLevelItem(fakeroot.takeChild(0));
 
-    ui->locals->applySavedExpansion(expansion, stateprefix, ui->locals->invisibleRootItem(), 0);
+    ui->locals->applySavedExpansion(expansion, stateprefix, 0);
   }
 
   if(ui->registers->topLevelItemCount() == 0)