Rebase against the upstream be900a4

vscode-upstream-sha1: be900a4

code/src/vs/base/common/prefixTree.ts

@@ -11,26 +11,51 @@ const unset = Symbol('unset');
  */
 export class WellDefinedPrefixTree<V> {
 	private readonly root = new Node<V>();
+	private _size = 0;
+
+	public get size() {
+		return this._size;
+	}
 
 	/** Inserts a new value in the prefix tree. */
 	insert(key: Iterable<string>, value: V): void {
-		let node = this.root;
+		this.opNode(key, n => n.value = value);
+	}
+
+	/** Mutates a value in the prefix tree. */
+	mutate(key: Iterable<string>, mutate: (value?: V) => V): void {
+		this.opNode(key, n => n.value = mutate(n.value === unset ? undefined : n.value));
+	}
+
+	/** Deletes a node from the prefix tree, returning the value it contained. */
+	delete(key: Iterable<string>): V | undefined {
+		const path = [{ part: '', node: this.root }];
+		let i = 0;
 		for (const part of key) {
-			if (!node.children) {
-				const next = new Node<V>();
-				node.children = new Map([[part, next]]);
-				node = next;
-			} else if (!node.children.has(part)) {
-				const next = new Node<V>();
-				node.children.set(part, next);
-				node = next;
-			} else {
-				node = node.children.get(part)!;
+			const node = path[i].node.children?.get(part);
+			if (!node) {
+				return undefined; // node not in tree
 			}
 
+			path.push({ part, node });
+			i++;
+		}
+
+		const value = path[i].node.value;
+		if (value === unset) {
+			return; // not actually a real node
 		}
 
-		node.value = value;
+		this._size--;
+		for (; i > 0; i--) {
+			const parent = path[i - 1];
+			parent.node.children!.delete(path[i].part);
+			if (parent.node.children!.size > 0 || parent.node.value !== unset) {
+				break;
+			}
+		}
+
+		return value;
 	}
 
 	/** Gets a value from the tree. */
@@ -80,6 +105,61 @@ export class WellDefinedPrefixTree<V> {
 
 		return true;
 	}
+
+	/** Gets whether the tree has the given key. */
+	hasKey(key: Iterable<string>): boolean {
+		let node = this.root;
+		for (const segment of key) {
+			const next = node.children?.get(segment);
+			if (!next) {
+				return false;
+			}
+
+			node = next;
+		}
+
+		return node.value !== unset;
+	}
+
+	private opNode(key: Iterable<string>, fn: (node: Node<V>) => void): void {
+		let node = this.root;
+		for (const part of key) {
+			if (!node.children) {
+				const next = new Node<V>();
+				node.children = new Map([[part, next]]);
+				node = next;
+			} else if (!node.children.has(part)) {
+				const next = new Node<V>();
+				node.children.set(part, next);
+				node = next;
+			} else {
+				node = node.children.get(part)!;
+			}
+		}
+
+		if (node.value === unset) {
+			this._size++;
+		}
+
+		fn(node);
+	}
+
+	/** Returns an iterable of the tree values in no defined order. */
+	*values() {
+		const stack = [this.root];
+		while (stack.length > 0) {
+			const node = stack.pop()!;
+			if (node.value !== unset) {
+				yield node.value;
+			}
+
+			if (node.children) {
+				for (const child of node.children.values()) {
+					stack.push(child);
+				}
+			}
+		}
+	}
 }
 
 class Node<T> {

code/src/vs/base/test/common/prefixTree.test.ts

@@ -44,4 +44,94 @@ suite('WellDefinedPrefixTree', () => {
 		assert.strictEqual(tree.hasKeyOrChildren(['foo', 'bar']), true);
 		assert.strictEqual(tree.hasKeyOrChildren(['foo', 'bar', 'baz']), false);
 	});
+
+	test('hasKey', () => {
+		const key = ['foo', 'bar'];
+		tree.insert(key, 42);
+
+		assert.strictEqual(tree.hasKey(key), true);
+		assert.strictEqual(tree.hasKey(['foo']), false);
+		assert.strictEqual(tree.hasKey(['baz']), false);
+		assert.strictEqual(tree.hasKey(['foo', 'bar', 'baz']), false);
+	});
+
+	test('size', () => {
+		const key1 = ['foo', 'bar'];
+		const key2 = ['foo', 'baz'];
+		assert.strictEqual(tree.size, 0);
+		tree.insert(key1, 42);
+		assert.strictEqual(tree.size, 1);
+		tree.insert(key2, 43);
+		assert.strictEqual(tree.size, 2);
+		tree.insert(key2, 44);
+		assert.strictEqual(tree.size, 2);
+	});
+
+	test('mutate', () => {
+		const key1 = ['foo', 'bar'];
+		const key2 = ['foo', 'baz'];
+		tree.insert(key1, 42);
+		tree.insert(key2, 43);
+		tree.mutate(key1, (value) => {
+			assert.strictEqual(value, 42);
+			return 44;
+		});
+		assert.strictEqual(tree.find(key1), 44);
+		assert.strictEqual(tree.find(key2), 43);
+	});
+
+	test('delete', () => {
+		const key1 = ['foo', 'bar'];
+		const key2 = ['foo', 'baz'];
+		tree.insert(key1, 42);
+		tree.insert(key2, 43);
+		assert.strictEqual(tree.size, 2);
+
+		assert.strictEqual(tree.delete(key1), 42);
+		assert.strictEqual(tree.size, 1);
+		assert.strictEqual(tree.find(key1), undefined);
+		assert.strictEqual(tree.find(key2), 43);
+
+		assert.strictEqual(tree.delete(key2), 43);
+		assert.strictEqual(tree.size, 0);
+		assert.strictEqual(tree.find(key1), undefined);
+		assert.strictEqual(tree.find(key2), undefined);
+
+		tree.delete(key2);
+		assert.strictEqual(tree.size, 0);
+	});
+
+	test('delete child', () => {
+		const key1 = ['foo', 'bar'];
+		const key2 = ['foo', 'bar', 'baz'];
+		tree.insert(key1, 42);
+		tree.insert(key2, 43);
+		assert.strictEqual(tree.size, 2);
+
+		assert.strictEqual(tree.delete(key2), 43);
+		assert.strictEqual(tree.size, 1);
+		assert.strictEqual(tree.find(key1), 42);
+		assert.strictEqual(tree.find(key2), undefined);
+	});
+
+	test('delete noops if deleting parent', () => {
+		const key1 = ['foo', 'bar'];
+		const key2 = ['foo', 'bar', 'baz'];
+		tree.insert(key2, 43);
+		assert.strictEqual(tree.size, 1);
+
+		assert.strictEqual(tree.delete(key1), undefined);
+		assert.strictEqual(tree.size, 1);
+		assert.strictEqual(tree.find(key2), 43);
+		assert.strictEqual(tree.find(key1), undefined);
+	});
+
+	test('values', () => {
+		const key1 = ['foo', 'bar'];
+		const key2 = ['foo', 'baz'];
+		tree.insert(key1, 42);
+		tree.insert(key2, 43);
+
+		assert.deepStrictEqual([...tree.values()], [43, 42]);
+	});
 });

code/src/vs/platform/lifecycle/electron-main/lifecycleMainService.ts

@@ -68,6 +68,20 @@ export interface ShutdownEvent {
 	join(id: string, promise: Promise<void>): void;
 }
 
+export interface IRelaunchHandler {
+
+	/**
+	 * Allows a handler to deal with relaunching the application. The return
+	 * value indicates if the relaunch is handled or not.
+	 */
+	handleRelaunch(options?: IRelaunchOptions): boolean;
+}
+
+export interface IRelaunchOptions {
+	readonly addArgs?: string[];
+	readonly removeArgs?: string[];
+}
+
 export interface ILifecycleMainService {
 
 	readonly _serviceBrand: undefined;
@@ -130,7 +144,12 @@ export interface ILifecycleMainService {
 	/**
 	 * Restart the application with optional arguments (CLI). All lifecycle event handlers are triggered.
 	 */
-	relaunch(options?: { addArgs?: string[]; removeArgs?: string[] }): Promise<void>;
+	relaunch(options?: IRelaunchOptions): Promise<void>;
+
+	/**
+	 * Sets a custom handler for relaunching the application.
+	 */
+	setRelaunchHandler(handler: IRelaunchHandler): void;
 
 	/**
 	 * Shutdown the application normally. All lifecycle event handlers are triggered.
@@ -224,6 +243,8 @@ export class LifecycleMainService extends Disposable implements ILifecycleMainSe
 
 	private readonly phaseWhen = new Map<LifecycleMainPhase, Barrier>();
 
+	private relaunchHandler: IRelaunchHandler | undefined = undefined;
+
 	constructor(
 		@ILogService private readonly logService: ILogService,
 		@IStateService private readonly stateService: IStateService,
@@ -553,6 +574,29 @@ export class LifecycleMainService extends Disposable implements ILifecycleMainSe
 	}
 
 	quit(willRestart?: boolean): Promise<boolean /* veto */> {
+		return this.doQuit(willRestart).then(veto => {
+			if (!veto && willRestart) {
+				// Windows: we are about to restart and as such we need to restore the original
+				// current working directory we had on startup to get the exact same startup
+				// behaviour. As such, we briefly change back to that directory and then when
+				// Code starts it will set it back to the installation directory again.
+				try {
+					if (isWindows) {
+						const currentWorkingDir = cwd();
+						if (currentWorkingDir !== process.cwd()) {
+							process.chdir(currentWorkingDir);
+						}
+					}
+				} catch (err) {
+					this.logService.error(err);
+				}
+			}
+
+			return veto;
+		});
+	}
+
+	private doQuit(willRestart?: boolean): Promise<boolean /* veto */> {
 		this.trace(`Lifecycle#quit() - begin (willRestart: ${willRestart})`);
 
 		if (this.pendingQuitPromise) {
@@ -588,7 +632,11 @@ export class LifecycleMainService extends Disposable implements ILifecycleMainSe
 		}
 	}
 
-	async relaunch(options?: { addArgs?: string[]; removeArgs?: string[] }): Promise<void> {
+	setRelaunchHandler(handler: IRelaunchHandler): void {
+		this.relaunchHandler = handler;
+	}
+
+	async relaunch(options?: IRelaunchOptions): Promise<void> {
 		this.trace('Lifecycle#relaunch()');
 
 		const args = process.argv.slice(1);
@@ -606,24 +654,10 @@ export class LifecycleMainService extends Disposable implements ILifecycleMainSe
 		}
 
 		const quitListener = () => {
-			// Windows: we are about to restart and as such we need to restore the original
-			// current working directory we had on startup to get the exact same startup
-			// behaviour. As such, we briefly change back to that directory and then when
-			// Code starts it will set it back to the installation directory again.
-			try {
-				if (isWindows) {
-					const currentWorkingDir = cwd();
-					if (currentWorkingDir !== process.cwd()) {
-						process.chdir(currentWorkingDir);
-					}
-				}
-			} catch (err) {
-				this.logService.error(err);
+			if (!this.relaunchHandler?.handleRelaunch(options)) {
+				this.trace('Lifecycle#relaunch() - calling app.relaunch()');
+				app.relaunch({ args });
 			}
-
-			// relaunch after we are sure there is no veto
-			this.trace('Lifecycle#relaunch() - calling app.relaunch()');
-			app.relaunch({ args });
 		};
 		app.once('quit', quitListener);
 

code/src/vs/platform/native/electron-main/nativeHostMainService.ts

@@ -26,7 +26,7 @@ import { INativeOpenDialogOptions } from 'vs/platform/dialogs/common/dialogs';
 import { IDialogMainService } from 'vs/platform/dialogs/electron-main/dialogMainService';
 import { IEnvironmentMainService } from 'vs/platform/environment/electron-main/environmentMainService';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
-import { ILifecycleMainService } from 'vs/platform/lifecycle/electron-main/lifecycleMainService';
+import { ILifecycleMainService, IRelaunchOptions } from 'vs/platform/lifecycle/electron-main/lifecycleMainService';
 import { ILogService } from 'vs/platform/log/common/log';
 import { ICommonNativeHostService, IOSProperties, IOSStatistics } from 'vs/platform/native/common/native';
 import { IProductService } from 'vs/platform/product/common/productService';
@@ -650,7 +650,7 @@ export class NativeHostMainService extends Disposable implements INativeHostMain
 		window?.setReady();
 	}
 
-	async relaunch(windowId: number | undefined, options?: { addArgs?: string[]; removeArgs?: string[] }): Promise<void> {
+	async relaunch(windowId: number | undefined, options?: IRelaunchOptions): Promise<void> {
 		return this.lifecycleMainService.relaunch(options);
 	}
 

code/src/vs/platform/test/electron-main/workbenchTestServices.ts

@@ -6,7 +6,7 @@
 import { Promises } from 'vs/base/common/async';
 import { Event, Emitter } from 'vs/base/common/event';
 import { NativeParsedArgs } from 'vs/platform/environment/common/argv';
-import { ILifecycleMainService, LifecycleMainPhase, ShutdownEvent, ShutdownReason } from 'vs/platform/lifecycle/electron-main/lifecycleMainService';
+import { ILifecycleMainService, IRelaunchHandler, LifecycleMainPhase, ShutdownEvent, ShutdownReason } from 'vs/platform/lifecycle/electron-main/lifecycleMainService';
 import { IStateService } from 'vs/platform/state/node/state';
 import { ICodeWindow, UnloadReason } from 'vs/platform/window/electron-main/window';
 
@@ -43,6 +43,7 @@ export class TestLifecycleMainService implements ILifecycleMainService {
 	registerWindow(window: ICodeWindow): void { }
 	async reload(window: ICodeWindow, cli?: NativeParsedArgs): Promise<void> { }
 	async unload(window: ICodeWindow, reason: UnloadReason): Promise<boolean> { return true; }
+	setRelaunchHandler(handler: IRelaunchHandler): void { }
 	async relaunch(options?: { addArgs?: string[] | undefined; removeArgs?: string[] | undefined }): Promise<void> { }
 	async quit(willRestart?: boolean): Promise<boolean> { return true; }
 	async kill(code?: number): Promise<void> { }

code/src/vs/platform/update/electron-main/abstractUpdateService.ts

@@ -46,7 +46,7 @@ export abstract class AbstractUpdateService implements IUpdateService {
 	}
 
 	constructor(
-		@ILifecycleMainService private readonly lifecycleMainService: ILifecycleMainService,
+		@ILifecycleMainService protected readonly lifecycleMainService: ILifecycleMainService,
 		@IConfigurationService protected configurationService: IConfigurationService,
 		@IEnvironmentMainService private readonly environmentMainService: IEnvironmentMainService,
 		@IRequestService protected requestService: IRequestService,