From 73690a41f6e3f7deb95bbc500310f0508ca1e46e Mon Sep 17 00:00:00 2001
From: Eric Amodio <eamodio@gmail.com>
Date: Sat, 20 Oct 2018 23:57:57 -0400
Subject: [PATCH] Adds node events for expand/collapse

---
 src/views/nodes/viewNode.ts | 23 ++++++++++++++++++++---
 src/views/viewBase.ts       | 23 ++++++++++++++++++++++-
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/src/views/nodes/viewNode.ts b/src/views/nodes/viewNode.ts
index d8be742..726f0c6 100644
--- a/src/views/nodes/viewNode.ts
+++ b/src/views/nodes/viewNode.ts
@@ -1,8 +1,8 @@
 'use strict';
-import { Command, Disposable, Event, TreeItem, TreeViewVisibilityChangeEvent } from 'vscode';
+import { Command, Disposable, Event, TreeItem, TreeItemCollapsibleState, TreeViewVisibilityChangeEvent } from 'vscode';
 import { GitUri } from '../../git/gitService';
 import { debug, logName } from '../../system';
-import { RefreshReason, View } from '../viewBase';
+import { RefreshReason, TreeViewNodeStateChangeEvent, View } from '../viewBase';
 
 export enum ResourceType {
     ActiveFileHistory = 'gitlens:active:history-file',
@@ -122,7 +122,10 @@ export abstract class SubscribeableViewNode<TView extends View> extends ViewNode
     ) {
         super(uri, parent);
 
-        const disposables = [this.view.onDidChangeVisibility(this.onVisibilityChanged, this)];
+        const disposables = [
+            this.view.onDidChangeVisibility(this.onVisibilityChanged, this),
+            this.view.onDidChangeNodeState(this.onNodeStateChanged, this)
+        ];
 
         if (supportsAutoRefresh(this.view)) {
             disposables.push(this.view.onDidChangeAutoRefresh(this.onAutoRefreshChanged, this));
@@ -180,6 +183,20 @@ export abstract class SubscribeableViewNode<TView extends View> extends ViewNode
         this.onVisibilityChanged({ visible: this.view.visible });
     }
 
+    protected onParentStateChanged(state: TreeItemCollapsibleState) {}
+    protected onStateChanged(state: TreeItemCollapsibleState) {}
+
+    protected _state: TreeItemCollapsibleState | undefined;
+    protected onNodeStateChanged(e: TreeViewNodeStateChangeEvent<ViewNode>) {
+        if (e.element === this) {
+            this._state = e.state;
+            this.onStateChanged(e.state);
+        }
+        else if (e.element === this._parent) {
+            this.onParentStateChanged(e.state);
+        }
+    }
+
     @debug()
     protected onVisibilityChanged(e: TreeViewVisibilityChangeEvent) {
         void this.ensureSubscription();
diff --git a/src/views/viewBase.ts b/src/views/viewBase.ts
index 5489e22..3f19724 100644
--- a/src/views/viewBase.ts
+++ b/src/views/viewBase.ts
@@ -6,7 +6,9 @@ import {
     EventEmitter,
     TreeDataProvider,
     TreeItem,
+    TreeItemCollapsibleState,
     TreeView,
+    TreeViewExpansionEvent,
     TreeViewVisibilityChangeEvent,
     window
 } from 'vscode';
@@ -30,6 +32,10 @@ export enum RefreshReason {
 
 export type View = RepositoriesView | FileHistoryView | LineHistoryView | ResultsView;
 
+export interface TreeViewNodeStateChangeEvent<T> extends TreeViewExpansionEvent<T> {
+    state: TreeItemCollapsibleState;
+}
+
 export abstract class ViewBase<TRoot extends ViewNode> implements TreeDataProvider<ViewNode>, Disposable {
     protected _onDidChangeTreeData = new EventEmitter<ViewNode>();
     public get onDidChangeTreeData(): Event<ViewNode> {
@@ -41,6 +47,11 @@ export abstract class ViewBase<TRoot extends ViewNode> implements TreeDataProvid
         return this._onDidChangeVisibility.event;
     }
 
+    private _onDidChangeNodeState = new EventEmitter<TreeViewNodeStateChangeEvent<ViewNode>>();
+    public get onDidChangeNodeState(): Event<TreeViewNodeStateChangeEvent<ViewNode>> {
+        return this._onDidChangeNodeState.event;
+    }
+
     protected _disposable: Disposable | undefined;
     protected _root: TRoot | undefined;
     protected _tree: TreeView<ViewNode> | undefined;
@@ -77,7 +88,9 @@ export abstract class ViewBase<TRoot extends ViewNode> implements TreeDataProvid
         });
         this._disposable = Disposable.from(
             this._tree,
-            this._tree.onDidChangeVisibility(this.onVisibilityChanged, this)
+            this._tree.onDidChangeVisibility(this.onVisibilityChanged, this),
+            this._tree.onDidCollapseElement(this.onElementCollapsed, this),
+            this._tree.onDidExpandElement(this.onElementExpanded, this)
         );
     }
 
@@ -99,6 +112,14 @@ export abstract class ViewBase<TRoot extends ViewNode> implements TreeDataProvid
         return node.getTreeItem();
     }
 
+    protected onElementCollapsed(e: TreeViewExpansionEvent<ViewNode>) {
+        this._onDidChangeNodeState.fire({ ...e, state: TreeItemCollapsibleState.Collapsed });
+    }
+
+    protected onElementExpanded(e: TreeViewExpansionEvent<ViewNode>) {
+        this._onDidChangeNodeState.fire({ ...e, state: TreeItemCollapsibleState.Expanded });
+    }
+
     protected onVisibilityChanged(e: TreeViewVisibilityChangeEvent) {
         this._onDidChangeVisibility.fire(e);
     }