GH-1643: keep expanded/collapsed state of structure tree nodes between refreshes

Author: martinlippert

vscode-extensions/vscode-spring-boot/lib/Main.ts

@@ -7,8 +7,9 @@ import {
     workspace,
     ExtensionContext,
     Uri,
-    lm
- } from 'vscode';
+    lm,
+    TreeItemCollapsibleState
+} from 'vscode';
 
 import * as commons from '@pivotal-tools/commons-vscode';
 import * as liveHoverUi from './live-hover-connect-ui';
@@ -186,7 +187,20 @@ export function activate(context: ExtensionContext): Thenable<ExtensionAPI> {
          */
         const structureManager = new StructureManager();
         const explorerTreeProvider = new ExplorerTreeProvider(structureManager);
-        context.subscriptions.push(window.createTreeView('explorer.spring', { treeDataProvider: explorerTreeProvider, showCollapseAll: true }));
+        const treeView = window.createTreeView('explorer.spring', { treeDataProvider: explorerTreeProvider, showCollapseAll: true });
+        
+        // Track expansion/collapse events to preserve state across refreshes
+        context.subscriptions.push(treeView.onDidExpandElement(e => {
+            const nodeId = e.element.getNodeId();
+            explorerTreeProvider.setExpansionState(nodeId, TreeItemCollapsibleState.Expanded);
+        }));
+        
+        context.subscriptions.push(treeView.onDidCollapseElement(e => {
+            const nodeId = e.element.getNodeId();
+            explorerTreeProvider.setExpansionState(nodeId, TreeItemCollapsibleState.Collapsed);
+        }));
+        
+        context.subscriptions.push(treeView);
         context.subscriptions.push(commands.registerCommand("vscode-spring-boot.structure.refresh", () => structureManager.refresh(true)));
         context.subscriptions.push(commands.registerCommand("vscode-spring-boot.structure.openReference", (node) => {
             if (node && node.getReferenceValue) {

vscode-extensions/vscode-spring-boot/lib/explorer/explorer-tree-provider.ts

@@ -1,20 +1,26 @@
-import { Event, EventEmitter, ProviderResult, TreeDataProvider, TreeItem } from "vscode";
+import { Event, EventEmitter, ProviderResult, TreeDataProvider, TreeItem, TreeItemCollapsibleState } from "vscode";
 import { StructureManager } from "./structure-tree-manager";
 import { SpringNode } from "./nodes";
 
 export class ExplorerTreeProvider implements TreeDataProvider<SpringNode> {
 
     private emitter: EventEmitter<undefined | SpringNode | SpringNode[]>;
     public readonly onDidChangeTreeData: Event<undefined | SpringNode | SpringNode[]>;
+    private expansionStates: Map<string, TreeItemCollapsibleState> = new Map();
 
     constructor(private manager: StructureManager) {
         this.emitter = new EventEmitter<undefined | SpringNode | SpringNode[]>();
         this.onDidChangeTreeData = this.emitter.event;
-        this.manager.onDidChange(e => this.emitter.fire(e));
+        this.manager.onDidChange(e => {
+            // Expansion states are tracked via onDidExpandElement/onDidCollapseElement events
+            this.emitter.fire(e);
+        });
     }
 
     getTreeItem(element: SpringNode): TreeItem | Thenable<TreeItem> {
-        return element.getTreeItem();
+        const nodeId = element.getNodeId();
+        const savedState = this.expansionStates.get(nodeId);
+        return element.getTreeItem(savedState);
     }
 
     getChildren(element?: SpringNode): ProviderResult<SpringNode[]> {
@@ -28,6 +34,15 @@ export class ExplorerTreeProvider implements TreeDataProvider<SpringNode> {
         return this.manager.rootElements;
     }
 
+
+    getExpansionState(nodeId: string): TreeItemCollapsibleState | undefined {
+        return this.expansionStates.get(nodeId);
+    }
+
+    setExpansionState(nodeId: string, state: TreeItemCollapsibleState): void {
+        this.expansionStates.set(nodeId, state);
+    }
+
     // getParent?(element: SpringNode): ProviderResult<SpringNode> {
     //     throw new Error("Method not implemented.");
     // }

vscode-extensions/vscode-spring-boot/lib/explorer/nodes.ts

@@ -4,20 +4,28 @@ import { LsStereoTypedNode } from "./structure-tree-manager";
 
 export class SpringNode {
     constructor(readonly children: SpringNode[]) {}
-    getTreeItem(): TreeItem {
-        return new TreeItem("<node>", this.computeState(TreeItemCollapsibleState.Expanded));
+    
+    getTreeItem(savedState?: TreeItemCollapsibleState): TreeItem {
+        const defaultState = savedState !== undefined ? savedState : TreeItemCollapsibleState.Collapsed;
+        return new TreeItem("<node>", this.computeState(defaultState));
     }
-    computeState(defaultState: TreeItemCollapsibleState.Collapsed | TreeItemCollapsibleState.Expanded): TreeItemCollapsibleState {
+    
+    computeState(defaultState: TreeItemCollapsibleState): TreeItemCollapsibleState {
         return Array.isArray(this.children) && this.children.length ? defaultState : TreeItemCollapsibleState.None;
     }
+    
+    getNodeId(): string {
+        return "<base-node>";
+    }
 }
 
 export class StereotypedNode extends SpringNode {
     constructor(private n: LsStereoTypedNode, children: SpringNode[]) {
         super(children);
     }
-    getTreeItem(): TreeItem {
-        const item = super.getTreeItem();
+    
+    getTreeItem(savedState?: TreeItemCollapsibleState): TreeItem {
+        const item = super.getTreeItem(savedState);
         item.label = this.n.attributes.text;
         item.iconPath = this.computeIcon();
 
@@ -39,6 +47,19 @@ export class StereotypedNode extends SpringNode {
         }
         return item;
     }
+    
+    getNodeId(): string {
+        // Create a unique identifier based on node attributes
+        // Use text, icon, and location as identifying factors
+        const textId = this.n.attributes.text || '';
+        const iconId = this.n.attributes.icon || '';
+        const locationId = this.n.attributes.location ? 
+            `${this.n.attributes.location.uri}:${this.n.attributes.location.range.start.line}:${this.n.attributes.location.range.start.character}` : '';
+        const referenceId = this.n.attributes.reference ? String(this.n.attributes.reference) : '';
+        
+        // Create a stable ID that can be used to match nodes across refreshes
+        return `${textId}|${iconId}|${locationId}|${referenceId}`.replace(/\|+$/, ''); // Remove trailing separators
+    }
 
     getReferenceValue(): any {
         return this.n.attributes.reference;