From 4f306c91c8842f0ced5e34d10df80cbab459b5e5 Mon Sep 17 00:00:00 2001
From: Sandeep Somavarapu <sasomava@microsoft.com>
Date: Mon, 12 Feb 2018 17:02:42 +0100
Subject: [PATCH] Setter and getter for data provider in tree viewer

---
 .../electron-browser/mainThreadTreeViews.ts   |  2 +-
 .../browser/parts/views/customView.ts         | 21 ++++++-------------
 src/vs/workbench/common/views.ts              |  4 +---
 3 files changed, 8 insertions(+), 19 deletions(-)

diff --git a/src/vs/workbench/api/electron-browser/mainThreadTreeViews.ts b/src/vs/workbench/api/electron-browser/mainThreadTreeViews.ts
index 50f205d1be5..6b327816fcb 100644
--- a/src/vs/workbench/api/electron-browser/mainThreadTreeViews.ts
+++ b/src/vs/workbench/api/electron-browser/mainThreadTreeViews.ts
@@ -31,7 +31,7 @@ export class MainThreadTreeViews extends Disposable implements MainThreadTreeVie
 	$registerTreeViewDataProvider(treeViewId: string): void {
 		const dataProvider = this._register(new TreeViewDataProvider(treeViewId, this._proxy, this.messageService));
 		this._dataProviders.set(treeViewId, dataProvider);
-		this.viewsService.registerTreeViewDataProvider(treeViewId, dataProvider);
+		this.viewsService.getTreeViewer(treeViewId).dataProvider = dataProvider;
 	}
 
 	$refresh(treeViewId: string, itemsToRefresh: { [treeItemHandle: string]: ITreeItem }): void {
diff --git a/src/vs/workbench/browser/parts/views/customView.ts b/src/vs/workbench/browser/parts/views/customView.ts
index f276597e0e2..fd26b0a3372 100644
--- a/src/vs/workbench/browser/parts/views/customView.ts
+++ b/src/vs/workbench/browser/parts/views/customView.ts
@@ -6,7 +6,7 @@
 import 'vs/css!./media/views';
 import Event, { Emitter } from 'vs/base/common/event';
 import * as errors from 'vs/base/common/errors';
-import { IDisposable, Disposable } from 'vs/base/common/lifecycle';
+import { IDisposable, Disposable, dispose } from 'vs/base/common/lifecycle';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { TPromise } from 'vs/base/common/winjs.base';
 import * as DOM from 'vs/base/browser/dom';
@@ -52,14 +52,6 @@ export class CustomViewsService extends Disposable implements ICustomViewsServic
 		return this.viewers.get(id);
 	}
 
-	registerTreeViewDataProvider(id: string, dataProvider: ITreeViewDataProvider): void {
-		const treeViewer = <CustomTreeViewer>this.getTreeViewer(id);
-		if (treeViewer) {
-			treeViewer.setDataProvider(dataProvider);
-			dataProvider.onDispose(() => treeViewer.setDataProvider(null));
-		}
-	}
-
 	private createViewers(viewDescriptors: IViewDescriptor[]): void {
 		for (const viewDescriptor of viewDescriptors) {
 			if ((<ICustomViewDescriptor>viewDescriptor).treeView) {
@@ -104,7 +96,7 @@ class CustomTreeViewer extends Disposable implements ITreeViewer {
 	private refreshing = 0;
 
 	private _dataProvider: ITreeViewDataProvider;
-	private dataProviderElementChangeListener: IDisposable;
+	private dataProviderDisposables: IDisposable[] = [];
 
 	constructor(
 		private id: string,
@@ -123,10 +115,8 @@ class CustomTreeViewer extends Disposable implements ITreeViewer {
 		return this._dataProvider;
 	}
 
-	setDataProvider(dataProvider: ITreeViewDataProvider) {
-		if (this.dataProviderElementChangeListener) {
-			this.dataProviderElementChangeListener.dispose();
-		}
+	set dataProvider(dataProvider: ITreeViewDataProvider) {
+		dispose(this.dataProviderDisposables);
 		if (dataProvider) {
 			const customTreeView: CustomTreeViewer = this;
 			this._dataProvider = new class implements ITreeViewDataProvider {
@@ -146,7 +136,8 @@ class CustomTreeViewer extends Disposable implements ITreeViewer {
 					});
 				}
 			};
-			this.dataProviderElementChangeListener = this._register(dataProvider.onDidChange(elements => this.refresh(elements)));
+			this._register(dataProvider.onDidChange(elements => this.refresh(elements), this, this.dataProviderDisposables));
+			this._register(dataProvider.onDispose(() => this.dataProvider = null, this, this.dataProviderDisposables));
 		} else {
 			this._dataProvider = null;
 		}
diff --git a/src/vs/workbench/common/views.ts b/src/vs/workbench/common/views.ts
index cab9d6035ca..402d35a7536 100644
--- a/src/vs/workbench/common/views.ts
+++ b/src/vs/workbench/common/views.ts
@@ -158,7 +158,7 @@ export interface IViewsViewlet extends IViewlet {
 
 export interface ITreeViewer extends IDisposable {
 
-	readonly dataProvider: ITreeViewDataProvider;
+	dataProvider: ITreeViewDataProvider;
 
 	refresh(treeItems?: ITreeItem[]): TPromise<void>;
 
@@ -185,8 +185,6 @@ export interface ICustomViewsService {
 	_serviceBrand: any;
 
 	getTreeViewer(id: string): ITreeViewer;
-
-	registerTreeViewDataProvider(id: string, ITreeViewDataProvider): void;
 }
 
 export type TreeViewItemHandleArg = {