Keep the tree expansion state when refreshing the JSON tree

src/NppJsonViewer/JsonViewDlg.cpp

@@ -1,4 +1,5 @@
 #include <format>
+#include <functional>
 #include <regex>
 
 #include "JsonViewDlg.h"
@@ -333,14 +334,21 @@ void JsonViewDlg::ValidateJson()
     DrawJsonTree();
 }
 
-void JsonViewDlg::DrawJsonTree()
+void JsonViewDlg::DrawJsonTree(bool bPreserveExpansion)
 {
     UpdateTitle();
 
     // Disable all buttons and treeView
     std::vector<DWORD> ctrls = {IDC_BTN_REFRESH, IDC_BTN_VALIDATE, IDC_BTN_FORMAT, IDC_BTN_SEARCH, IDC_EDT_SEARCH};
     EnableControls(ctrls, false);
 
+    // Capture the expansion/selection state before the tree is thrown away, so
+    // that it can be re-applied onto the freshly built one.
+    TreeExpansionState expState;
+    const bool         bHasCurrentTree = m_pTreeView->GetRoot() && m_pTreeView->GetNodeCount() > 1;
+    if (bPreserveExpansion && bHasCurrentTree)
+        expState = CaptureExpansionState();
+
     HTREEITEM rootNode = nullptr;
     rootNode           = m_pTreeView->InitTree();
 
@@ -380,6 +388,11 @@ void JsonViewDlg::DrawJsonTree()
 
     m_pTreeView->Expand(rootNode);
 
+    // Re-apply the state of the previous tree. Paths that no longer exist
+    // (the document changed in the meantime) are silently dropped.
+    if (bPreserveExpansion && bHasCurrentTree && m_pTreeView->GetNodeCount() > 1)
+        ApplyExpansionState(expState);
+
     // Enable all buttons and treeView
     EnableControls(ctrls, true);
 }
@@ -565,6 +578,163 @@ void JsonViewDlg::SearchInTree()
     }
 }
 
+TreeExpansionState JsonViewDlg::CaptureExpansionState() const
+{
+    TreeExpansionState expState;
+
+    auto hRoot = m_pTreeView->GetRoot();
+    if (!hRoot)
+        return expState;
+
+    // Collect the key path of every node (root excluded) with its expanded flag
+    std::function<void(HTREEITEM, const std::vector<std::wstring>&)> walk;
+    walk = [&](HTREEITEM hParent, const std::vector<std::wstring>& parentKeys) {
+        for (HTREEITEM hChild = m_pTreeView->GetChildItem(hParent); hChild;
+             hChild           = m_pTreeView->GetNextSibling(hChild))
+        {
+            auto keys    = parentKeys;
+            auto nodeKey = GetPathKey(hChild);
+            keys.push_back(nodeKey);
+
+            expState.expandedPaths[TreeExpansionHelper::JoinPath(parentKeys, nodeKey)] = m_pTreeView->IsExpanded(hChild);
+
+            walk(hChild, keys);
+        }
+    };
+
+    walk(hRoot, {});
+
+    expState.selectedPath = GetCurrentSelectedPath();
+
+    return expState;
+}
+
+void JsonViewDlg::ApplyExpansionState(const TreeExpansionState& state)
+{
+    auto paths                                      = CollectExpandedPaths();
+    auto [pathsToExpand, pathToSelect]              = TreeExpansionHelper::MatchExpansion(state, paths);
+
+    for (const auto& path : pathsToExpand)
+    {
+        auto keys = TreeExpansionHelper::SplitPath(path);
+        if (!keys.empty())
+            ExpandByPath(keys);
+    }
+
+    if (!pathToSelect.empty())
+        SelectByPath(pathToSelect);
+}
+
+std::vector<std::wstring> JsonViewDlg::CollectExpandedPaths() const
+{
+    std::vector<std::wstring> paths;
+
+    auto hRoot = m_pTreeView->GetRoot();
+    if (!hRoot)
+        return paths;
+
+    std::function<void(HTREEITEM, const std::vector<std::wstring>&)> walk;
+    walk = [&](HTREEITEM hParent, const std::vector<std::wstring>& parentKeys) {
+        for (HTREEITEM hChild = m_pTreeView->GetChildItem(hParent); hChild;
+             hChild           = m_pTreeView->GetNextSibling(hChild))
+        {
+            auto nodeKey = GetPathKey(hChild);
+
+            auto keys = parentKeys;
+            keys.push_back(nodeKey);
+
+            paths.push_back(TreeExpansionHelper::JoinPath(parentKeys, nodeKey));
+
+            walk(hChild, keys);
+        }
+    };
+
+    walk(hRoot, {});
+
+    return paths;
+}
+
+std::vector<std::wstring> JsonViewDlg::GetCurrentSelectedPath() const
+{
+    std::vector<std::wstring> path;
+
+    auto hRoot     = m_pTreeView->GetRoot();
+    auto hSelected = m_pTreeView->GetSelection();
+    if (!hRoot || !hSelected || hSelected == hRoot)
+        return path;
+
+    // Walk up to the root and reverse the collected keys on the way back
+    std::vector<std::wstring> reversedKeys;
+    for (HTREEITEM h = hSelected; h && h != hRoot; h = m_pTreeView->GetParentItem(h))
+    {
+        reversedKeys.push_back(GetPathKey(h));
+    }
+
+    // Guard against a selection that does not belong to this tree anymore
+    if (m_pTreeView->GetParentItem(hSelected) == nullptr)
+        return {};
+
+    path.assign(reversedKeys.rbegin(), reversedKeys.rend());
+    return path;
+}
+
+auto JsonViewDlg::GetPathKey(HTREEITEM hti) const -> std::wstring
+{
+    auto key = m_pTreeView->GetNodeKey(hti);
+
+    // Remove the surrounding quotes of object keys: "name" -> name.
+    // Array indices ([0]) and unquoted keys are returned untouched.
+    if (key.size() >= 2 && key.front() == L'"' && key.back() == L'"')
+        key = key.substr(1, key.size() - 2);
+
+    return key;
+}
+
+auto JsonViewDlg::FindNodeByPath(const std::vector<std::wstring>& path) const -> HTREEITEM
+{
+    if (path.empty())
+        return nullptr;
+
+    auto hRoot = m_pTreeView->GetRoot();
+    if (!hRoot)
+        return nullptr;
+
+    HTREEITEM hCurrent = hRoot;
+    for (const auto& key : path)
+    {
+        HTREEITEM hNext = m_pTreeView->GetChildItem(hCurrent);
+        while (hNext && GetPathKey(hNext) != key)
+        {
+            hNext = m_pTreeView->GetNextSibling(hNext);
+        }
+
+        if (!hNext)
+            return nullptr;
+
+        hCurrent = hNext;
+    }
+
+    return hCurrent == hRoot ? nullptr : hCurrent;
+}
+
+void JsonViewDlg::ExpandByPath(const std::vector<std::wstring>& path)
+{
+    auto hNode = FindNodeByPath(path);
+    if (hNode)
+        m_pTreeView->Expand(hNode);
+}
+
+void JsonViewDlg::SelectByPath(const std::vector<std::wstring>& path)
+{
+    auto hNode = FindNodeByPath(path);
+    if (hNode)
+    {
+        // TreeView_SelectItem expands collapsed ancestors on its own, so the
+        // expansion state restored just before is left untouched.
+        m_pTreeView->SetSelection(hNode);
+    }
+}
+
 void JsonViewDlg::UpdateTitle()
 {
     const auto titleFileName = GetTitleFileName();
@@ -1115,7 +1285,7 @@ INT_PTR JsonViewDlg::run_dlgProc(UINT message, WPARAM wParam, LPARAM lParam)
         {
             // Handle Button events
         case IDC_BTN_REFRESH:
-            DrawJsonTree();
+            DrawJsonTree(true);
             break;
 
         case IDC_BTN_FORMAT:

src/NppJsonViewer/JsonViewDlg.h

@@ -14,6 +14,7 @@
 #include "JsonHandler.h"
 #include "JsonNode.h"
 #include "TreeHandler.h"
+#include "TreeExpansion.h"
 
 
 class JsonViewDlg
@@ -53,7 +54,7 @@ class JsonViewDlg
     void      AppendNodeCount(HTREEITEM node, unsigned elementCount, bool bArray) override;
 
 private:
-    void DrawJsonTree();
+    void DrawJsonTree(bool bPreserveExpansion = false);
     void ReDrawJsonTree(bool bForce = false);
     void HighlightAsJson(bool bForcefully = false) const;
     auto PopulateTreeUsingSax(HTREEITEM tree_root, const std::string& jsonText) -> std::optional<std::wstring>;
@@ -66,6 +67,24 @@ class JsonViewDlg
 
     void SearchInTree();
 
+    // Expansion/selection state, captured before the tree is rebuilt and
+    // re-applied afterwards (see DrawJsonTree(bPreserveExpansion = true)).
+    auto CaptureExpansionState() const -> TreeExpansionState;
+    void ApplyExpansionState(const TreeExpansionState& state);
+
+    auto CollectExpandedPaths() const -> std::vector<std::wstring>;
+    auto GetCurrentSelectedPath() const -> std::vector<std::wstring>;
+    void ExpandByPath(const std::vector<std::wstring>& path);
+    void SelectByPath(const std::vector<std::wstring>& path);
+
+    // Key of a node as used inside a node path: the raw key without the
+    // surrounding quotes ("name" -> name, [0] -> [0]).
+    auto GetPathKey(HTREEITEM hti) const -> std::wstring;
+
+    // Resolve a key path (relative to the tree root) back to a node.
+    // Returns nullptr when any level of the path cannot be found.
+    auto FindNodeByPath(const std::vector<std::wstring>& path) const -> HTREEITEM;
+
     auto GetTitleFileName() const -> std::wstring;
     void PrepareButtons();
     void SetIconAndTooltip(eButton ctrlType, const std::wstring& toolTip);

src/NppJsonViewer/NPPJSONViewer.vcxproj

@@ -197,6 +197,7 @@
     <ClCompile Include="SettingsDlg.cpp" />
     <ClCompile Include="ShortcutCommand.cpp" />
     <ClCompile Include="SliderCtrl.cpp" />
+    <ClCompile Include="TreeExpansion.cpp" />
     <ClCompile Include="TreeViewCtrl.cpp" />
     <ClCompile Include="npp_include\DockingFeature\StaticDialog.cpp" />
   </ItemGroup>
@@ -217,6 +218,7 @@
     <ClInclude Include="SliderCtrl.h" />
     <ClInclude Include="StopWatch.h" />
     <ClInclude Include="TrackingStream.h" />
+    <ClInclude Include="TreeExpansion.h" />
     <ClInclude Include="TreeHandler.h" />
     <ClInclude Include="TreeViewCtrl.h" />
     <ClInclude Include="npp_include\DockingFeature\Docking.h" />

src/NppJsonViewer/NPPJSONViewer.vcxproj.filters

@@ -57,6 +57,9 @@
     <ClCompile Include="SliderCtrl.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="TreeExpansion.cpp">
+      <Filter>Source Files</Filter>
+    </ClCompile>
     <ClCompile Include="TreeViewCtrl.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
@@ -110,6 +113,9 @@
     <ClInclude Include="TrackingStream.h">
       <Filter>Header Files</Filter>
     </ClInclude>
+    <ClInclude Include="TreeExpansion.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
     <ClInclude Include="TreeHandler.h">
       <Filter>Header Files</Filter>
     </ClInclude>

src/NppJsonViewer/TreeExpansion.cpp

@@ -0,0 +1,75 @@
+#include "TreeExpansion.h"
+
+#include <algorithm>
+
+auto TreeExpansionHelper::SplitPath(const std::wstring& path) -> std::vector<std::wstring>
+{
+    std::vector<std::wstring> keys;
+
+    if (path.empty())
+        return keys;
+
+    size_t start = 0;
+    while (start <= path.size())
+    {
+        auto pos = path.find(L'.', start);
+        if (pos == std::wstring::npos)
+        {
+            keys.emplace_back(path.substr(start));
+            break;
+        }
+
+        keys.emplace_back(path.substr(start, pos - start));
+        start = pos + 1;
+    }
+
+    // An empty trailing key (e.g. a path ending with a dot) is dropped
+    if (!keys.empty() && keys.back().empty())
+        keys.pop_back();
+
+    return keys;
+}
+
+auto TreeExpansionHelper::JoinPath(const std::vector<std::wstring>& parentKeys, const std::wstring& key) -> std::wstring
+{
+    std::wstring path;
+    for (const auto& part : parentKeys)
+    {
+        path += part;
+        path += L'.';
+    }
+    path += key;
+    return path;
+}
+
+auto TreeExpansionHelper::MatchExpansion(const TreeExpansionState& oldState,
+                                         const std::vector<std::wstring>& newPaths)
+    -> std::pair<std::vector<std::wstring>, std::vector<std::wstring>>
+{
+    std::vector<std::wstring> pathsToExpand;
+
+    for (const auto& path : newPaths)
+    {
+        auto find = oldState.expandedPaths.find(path);
+        if (find != oldState.expandedPaths.cend() && find->second)
+            pathsToExpand.push_back(path);
+    }
+
+    std::vector<std::wstring> pathToSelect;
+    if (!oldState.selectedPath.empty())
+    {
+        // Reconstruct the selected path in "joined" form and check existence
+        std::wstring joined;
+        for (const auto& key : oldState.selectedPath)
+        {
+            if (!joined.empty())
+                joined += L'.';
+            joined += key;
+        }
+
+        if (std::find(newPaths.cbegin(), newPaths.cend(), joined) != newPaths.cend())
+            pathToSelect = oldState.selectedPath;
+    }
+
+    return { pathsToExpand, pathToSelect };
+}