Custom Editors: pass original edit objects back to extensions

For #88719

With this change, instead of passing custom editor edit json back and forth with the extension host, we keep the original edit objects on the extension host. This means that we can pass extensions back the exact same edit object they first hand to us. It also means that edits no longer need to be json serializable.

src/vs/vscode.proposed.d.ts

@@ -1181,7 +1181,7 @@ declare module 'vscode' {
 	 * Defines the editing functionality of a webview editor. This allows the webview editor to hook into standard
 	 * editor events such as `undo` or `save`.
 	 *
-	 * @param EditType Type of edits. Edit objects must be json serializable.
+	 * @param EditType Type of edits.
 	 */
 	interface WebviewCustomEditorEditingDelegate<EditType> {
 		/**

src/vs/workbench/api/browser/mainThreadWebview.ts

@@ -306,6 +306,8 @@ export class MainThreadWebviews extends Disposable implements extHostProtocol.Ma
 
 		provider.dispose();
 		this._editorProviders.delete(viewType);
+
+		this._customEditorService.models.disposeAllModelsForView(viewType);
 	}
 
 	private async retainCustomEditorModel(webviewInput: WebviewInput, resource: URI, viewType: string, capabilities: readonly extHostProtocol.WebviewEditorCapabilities[]) {
@@ -323,14 +325,17 @@ export class MainThreadWebviews extends Disposable implements extHostProtocol.Ma
 		const capabilitiesSet = new Set(capabilities);
 		const isEditable = capabilitiesSet.has(extHostProtocol.WebviewEditorCapabilities.Editable);
 		if (isEditable) {
-			model.onUndo(edits => {
-				this._proxy.$undoEdits(resource, viewType, edits.map(x => x.data));
+			model.onUndo(e => {
+				this._proxy.$undoEdits(resource, viewType, e.edits);
+			});
+
+			model.onDisposeEdits(e => {
+				this._proxy.$disposeEdits(e.edits);
 			});
 
-			model.onApplyEdit(edits => {
-				const editsToApply = edits.filter(x => x.source !== model).map(x => x.data);
-				if (editsToApply.length) {
-					this._proxy.$applyEdits(resource, viewType, editsToApply);
+			model.onApplyEdit(e => {
+				if (e.trigger !== model) {
+					this._proxy.$applyEdits(resource, viewType, e.edits);
 				}
 			});
 
@@ -369,13 +374,13 @@ export class MainThreadWebviews extends Disposable implements extHostProtocol.Ma
 		}
 	}
 
-	public $onEdit(resource: UriComponents, viewType: string, editData: any): void {
+	public $onEdit(resource: UriComponents, viewType: string, editId: number): void {
 		const model = this._customEditorService.models.get(URI.revive(resource), viewType);
 		if (!model) {
 			throw new Error('Could not find model for webview editor');
 		}
 
-		model.pushEdit({ source: model, data: editData });
+		model.pushEdit(editId, model);
 	}
 
 	private hookupWebviewEventDelegate(handle: extHostProtocol.WebviewPanelHandle, input: WebviewInput) {

src/vs/workbench/api/common/cache.ts

+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+export class Cache<T> {
+
+	private static readonly enableDebugLogging = false;
+
+	private readonly _data = new Map<number, readonly T[]>();
+	private _idPool = 1;
+
+	constructor(
+		private readonly id: string
+	) { }
+
+	add(item: readonly T[]): number {
+		const id = this._idPool++;
+		this._data.set(id, item);
+		this.logDebugInfo();
+		return id;
+	}
+
+	get(pid: number, id: number): T | undefined {
+		return this._data.has(pid) ? this._data.get(pid)![id] : undefined;
+	}
+
+	delete(id: number) {
+		this._data.delete(id);
+		this.logDebugInfo();
+	}
+
+	private logDebugInfo() {
+		if (!Cache.enableDebugLogging) {
+			return;
+		}
+		console.log(`${this.id} cache size — ${this._data.size}`);
+	}
+}

src/vs/workbench/api/common/extHost.protocol.ts

@@ -584,7 +584,7 @@ export interface MainThreadWebviewsShape extends IDisposable {
 	$registerEditorProvider(extension: WebviewExtensionDescription, viewType: string, options: modes.IWebviewPanelOptions, capabilities: readonly WebviewEditorCapabilities[]): void;
 	$unregisterEditorProvider(viewType: string): void;
 
-	$onEdit(resource: UriComponents, viewType: string, editJson: any): void;
+	$onEdit(resource: UriComponents, viewType: string, editId: number): void;
 }
 
 export interface WebviewPanelViewStateData {
@@ -604,8 +604,9 @@ export interface ExtHostWebviewsShape {
 	$deserializeWebviewPanel(newWebviewHandle: WebviewPanelHandle, viewType: string, title: string, state: any, position: EditorViewColumn, options: modes.IWebviewOptions & modes.IWebviewPanelOptions): Promise<void>;
 	$resolveWebviewEditor(resource: UriComponents, newWebviewHandle: WebviewPanelHandle, viewType: string, title: string, position: EditorViewColumn, options: modes.IWebviewOptions & modes.IWebviewPanelOptions): Promise<void>;
 
-	$undoEdits(resource: UriComponents, viewType: string, edits: readonly any[]): void;
-	$applyEdits(resource: UriComponents, viewType: string, edits: readonly any[]): void;
+	$undoEdits(resource: UriComponents, viewType: string, editIds: readonly number[]): void;
+	$applyEdits(resource: UriComponents, viewType: string, editIds: readonly number[]): void;
+	$disposeEdits(editIds: readonly number[]): void;
 
 	$onSave(resource: UriComponents, viewType: string): Promise<void>;
 	$onSaveAs(resource: UriComponents, viewType: string, targetResource: UriComponents): Promise<void>;

src/vs/workbench/api/common/extHostLanguageFeatures.ts

@@ -30,6 +30,7 @@ import { VSBuffer } from 'vs/base/common/buffer';
 import { encodeSemanticTokensDto } from 'vs/workbench/api/common/shared/semanticTokens';
 import { IdGenerator } from 'vs/base/common/idGenerator';
 import { IExtHostApiDeprecationService } from 'vs/workbench/api/common/extHostApiDeprecationService';
+import { Cache } from './cache';
 
 // --- adapter
 
@@ -1064,40 +1065,6 @@ class SignatureHelpAdapter {
 	}
 }
 
-class Cache<T> {
-	private static readonly enableDebugLogging = false;
-
-	private readonly _data = new Map<number, readonly T[]>();
-	private _idPool = 1;
-
-	constructor(
-		private readonly id: string
-	) { }
-
-	add(item: readonly T[]): number {
-		const id = this._idPool++;
-		this._data.set(id, item);
-		this.logDebugInfo();
-		return id;
-	}
-
-	get(pid: number, id: number): T | undefined {
-		return this._data.has(pid) ? this._data.get(pid)![id] : undefined;
-	}
-
-	delete(id: number) {
-		this._data.delete(id);
-		this.logDebugInfo();
-	}
-
-	private logDebugInfo() {
-		if (!Cache.enableDebugLogging) {
-			return;
-		}
-		console.log(`${this.id} cache size — ${this._data.size}`);
-	}
-}
-
 class LinkProviderAdapter {
 
 	private _cache = new Cache<vscode.DocumentLink>('DocumentLink');

src/vs/workbench/api/common/extHostWebview.ts

@@ -15,6 +15,7 @@ import { IExtHostWorkspace } from 'vs/workbench/api/common/extHostWorkspace';
 import { EditorViewColumn } from 'vs/workbench/api/common/shared/editor';
 import { asWebviewUri, WebviewInitData } from 'vs/workbench/api/common/shared/webview';
 import type * as vscode from 'vscode';
+import { Cache } from './cache';
 import { ExtHostWebviewsShape, IMainContext, MainContext, MainThreadWebviewsShape, WebviewEditorCapabilities, WebviewPanelHandle, WebviewPanelViewStateData } from './extHost.protocol';
 import { Disposable as VSCodeDisposable } from './extHostTypes';
 
@@ -251,8 +252,18 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 
 	private readonly _proxy: MainThreadWebviewsShape;
 	private readonly _webviewPanels = new Map<WebviewPanelHandle, ExtHostWebviewEditor>();
-	private readonly _serializers = new Map<string, { readonly serializer: vscode.WebviewPanelSerializer, readonly extension: IExtensionDescription }>();
-	private readonly _editorProviders = new Map<string, { readonly provider: vscode.WebviewCustomEditorProvider, readonly extension: IExtensionDescription }>();
+
+	private readonly _serializers = new Map<string, {
+		readonly serializer: vscode.WebviewPanelSerializer;
+		readonly extension: IExtensionDescription;
+	}>();
+
+	private readonly _editorProviders = new Map<string, {
+		readonly provider: vscode.WebviewCustomEditorProvider;
+		readonly extension: IExtensionDescription;
+	}>();
+
+	private readonly _edits = new Cache<unknown>('edits');
 
 	constructor(
 		mainContext: IMainContext,
@@ -312,11 +323,14 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 		if (this._editorProviders.has(viewType)) {
 			throw new Error(`Editor provider for '${viewType}' already registered`);
 		}
-
 		this._editorProviders.set(viewType, { extension, provider, });
+
 		this._proxy.$registerEditorProvider({ id: extension.identifier, location: extension.extensionLocation }, viewType, options || {}, this.getCapabilites(provider));
+
+		// Hook up events
 		provider?.editingDelegate?.onEdit(({ edit, resource }) => {
-			this._proxy.$onEdit(resource, viewType, edit);
+			const id = this._edits.add([edit]);
+			this._proxy.$onEdit(resource, viewType, id);
 		});
 
 		return new VSCodeDisposable(() => {
@@ -426,14 +440,32 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 		await provider.resolveWebviewEditor(revivedResource, revivedPanel);
 	}
 
-	$undoEdits(resource: UriComponents, viewType: string, edits: readonly any[]): void {
+	$undoEdits(resourceComponents: UriComponents, viewType: string, editIds: readonly number[]): void {
 		const provider = this.getEditorProvider(viewType);
-		provider?.editingDelegate?.undoEdits(URI.revive(resource), edits);
+		if (!provider?.editingDelegate) {
+			return;
+		}
+
+		const resource = URI.revive(resourceComponents);
+		const edits = editIds.map(id => this._edits.get(id, 0));
+		provider.editingDelegate.undoEdits(resource, edits);
 	}
 
-	$applyEdits(resource: UriComponents, viewType: string, edits: readonly any[]): void {
+	$applyEdits(resourceComponents: UriComponents, viewType: string, editIds: readonly number[]): void {
 		const provider = this.getEditorProvider(viewType);
-		provider?.editingDelegate?.applyEdits(URI.revive(resource), edits);
+		if (!provider?.editingDelegate) {
+			return;
+		}
+
+		const resource = URI.revive(resourceComponents);
+		const edits = editIds.map(id => this._edits.get(id, 0));
+		provider.editingDelegate.applyEdits(resource, edits);
+	}
+
+	$disposeEdits(editIds: readonly number[]): void {
+		for (const edit of editIds) {
+			this._edits.delete(edit);
+		}
 	}
 
 	async $onSave(resource: UriComponents, viewType: string): Promise<void> {

src/vs/workbench/contrib/customEditor/common/customEditor.ts

@@ -42,7 +42,7 @@ export interface ICustomEditorService {
 	promptOpenWith(resource: URI, options?: ITextEditorOptions, group?: IEditorGroup): Promise<IEditor | undefined>;
 }
 
-export type CustomEditorEdit = { source?: any, data: any };
+export type CustomEditorEdit = number;
 
 export interface ICustomEditorModelManager {
 	get(resource: URI, viewType: string): ICustomEditorModel | undefined;
@@ -50,6 +50,8 @@ export interface ICustomEditorModelManager {
 	resolve(resource: URI, viewType: string): Promise<ICustomEditorModel>;
 
 	disposeModel(model: ICustomEditorModel): void;
+
+	disposeAllModelsForView(viewType: string): void;
 }
 
 export interface CustomEditorSaveEvent {
@@ -64,13 +66,15 @@ export interface CustomEditorSaveAsEvent {
 }
 
 export interface ICustomEditorModel extends IWorkingCopy {
-	readonly onUndo: Event<readonly CustomEditorEdit[]>;
-	readonly onApplyEdit: Event<readonly CustomEditorEdit[]>;
+	readonly viewType: string;
+
+	readonly onUndo: Event<{ edits: readonly CustomEditorEdit[], trigger: any | undefined }>;
+	readonly onApplyEdit: Event<{ edits: readonly CustomEditorEdit[], trigger: any | undefined }>;
+	readonly onDisposeEdits: Event<{ edits: readonly CustomEditorEdit[] }>;
+
 	readonly onWillSave: Event<CustomEditorSaveEvent>;
 	readonly onWillSaveAs: Event<CustomEditorSaveAsEvent>;
 
-	readonly currentEdits: readonly CustomEditorEdit[];
-
 	undo(): void;
 	redo(): void;
 	revert(options?: IRevertOptions): Promise<boolean>;
@@ -78,7 +82,7 @@ export interface ICustomEditorModel extends IWorkingCopy {
 	save(options?: ISaveOptions): Promise<boolean>;
 	saveAs(resource: URI, targetResource: URI, currentOptions?: ISaveOptions): Promise<boolean>;
 
-	pushEdit(edit: CustomEditorEdit): void;
+	pushEdit(edit: CustomEditorEdit, trigger: any): void;
 }
 
 export const enum CustomEditorPriority {

src/vs/workbench/contrib/customEditor/common/customEditorModel.ts

@@ -14,14 +14,20 @@ export class CustomEditorModel extends Disposable implements ICustomEditorModel
 
 	private _currentEditIndex: number = -1;
 	private _savePoint: number = -1;
-	private _edits: Array<any> = [];
+	private readonly _edits: Array<CustomEditorEdit> = [];
 
 	constructor(
+		public readonly viewType: string,
 		private readonly _resource: URI,
 	) {
 		super();
 	}
 
+	dispose() {
+		this._onDisposeEdits.fire({ edits: this._edits });
+		super.dispose();
+	}
+
 	//#region IWorkingCopy
 
 	public get resource() {
@@ -44,30 +50,43 @@ export class CustomEditorModel extends Disposable implements ICustomEditorModel
 
 	//#endregion
 
-	protected readonly _onUndo = this._register(new Emitter<readonly CustomEditorEdit[]>());
+	protected readonly _onUndo = this._register(new Emitter<{ edits: readonly CustomEditorEdit[], trigger: any | undefined }>());
 	readonly onUndo = this._onUndo.event;
 
-	protected readonly _onApplyEdit = this._register(new Emitter<readonly CustomEditorEdit[]>());
+	protected readonly _onApplyEdit = this._register(new Emitter<{ edits: readonly CustomEditorEdit[], trigger: any | undefined }>());
 	readonly onApplyEdit = this._onApplyEdit.event;
 
+	protected readonly _onDisposeEdits = this._register(new Emitter<{ edits: readonly CustomEditorEdit[] }>());
+	readonly onDisposeEdits = this._onDisposeEdits.event;
+
 	protected readonly _onWillSave = this._register(new Emitter<CustomEditorSaveEvent>());
 	readonly onWillSave = this._onWillSave.event;
 
 	protected readonly _onWillSaveAs = this._register(new Emitter<CustomEditorSaveAsEvent>());
 	readonly onWillSaveAs = this._onWillSaveAs.event;
 
-	get currentEdits(): readonly CustomEditorEdit[] {
-		return this._edits.slice(0, Math.max(0, this._currentEditIndex + 1));
-	}
+	public pushEdit(edit: CustomEditorEdit, trigger: any): void {
+		this.spliceEdits(edit);
 
-	public pushEdit(edit: CustomEditorEdit): void {
-		this._edits.splice(this._currentEditIndex + 1, this._edits.length - this._currentEditIndex, edit.data);
 		this._currentEditIndex = this._edits.length - 1;
 		this.updateDirty();
-		this._onApplyEdit.fire([edit]);
+		this._onApplyEdit.fire({ edits: [edit], trigger });
 		this.updateContentChanged();
 	}
 
+	private spliceEdits(editToInsert?: CustomEditorEdit) {
+		const start = this._currentEditIndex + 1;
+		const toRemove = this._edits.length - this._currentEditIndex;
+
+		const removedEdits = editToInsert
+			? this._edits.splice(start, toRemove, editToInsert)
+			: this._edits.splice(start, toRemove);
+
+		if (removedEdits.length) {
+			this._onDisposeEdits.fire({ edits: removedEdits });
+		}
+	}
+
 	private updateDirty() {
 		// TODO@matt this should to be more fine grained and avoid
 		// emitting events if there was no change actually
@@ -128,14 +147,15 @@ export class CustomEditorModel extends Disposable implements ICustomEditorModel
 
 		if (this._currentEditIndex >= this._savePoint) {
 			const editsToUndo = this._edits.slice(this._savePoint, this._currentEditIndex);
-			this._onUndo.fire(editsToUndo.reverse());
+			this._onUndo.fire({ edits: editsToUndo.reverse(), trigger: undefined });
 		} else if (this._currentEditIndex < this._savePoint) {
 			const editsToRedo = this._edits.slice(this._currentEditIndex, this._savePoint);
-			this._onApplyEdit.fire(editsToRedo);
+			this._onApplyEdit.fire({ edits: editsToRedo, trigger: undefined });
 		}
 
 		this._currentEditIndex = this._savePoint;
-		this._edits.splice(this._currentEditIndex + 1, this._edits.length - this._currentEditIndex);
+		this.spliceEdits();
+
 		this.updateDirty();
 		this.updateContentChanged();
 		return true;
@@ -149,7 +169,7 @@ export class CustomEditorModel extends Disposable implements ICustomEditorModel
 
 		const undoneEdit = this._edits[this._currentEditIndex];
 		--this._currentEditIndex;
-		this._onUndo.fire([{ data: undoneEdit }]);
+		this._onUndo.fire({ edits: [undoneEdit], trigger: undefined });
 
 		this.updateDirty();
 		this.updateContentChanged();
@@ -164,7 +184,7 @@ export class CustomEditorModel extends Disposable implements ICustomEditorModel
 		++this._currentEditIndex;
 		const redoneEdit = this._edits[this._currentEditIndex];
 
-		this._onApplyEdit.fire([{ data: redoneEdit }]);
+		this._onApplyEdit.fire({ edits: [redoneEdit], trigger: undefined });
 
 		this.updateDirty();
 		this.updateContentChanged();

src/vs/workbench/contrib/customEditor/common/customEditorModelManager.ts

@@ -27,7 +27,7 @@ export class CustomEditorModelManager implements ICustomEditorModelManager {
 			return existing;
 		}
 
-		const model = new CustomEditorModel(resource);
+		const model = new CustomEditorModel(viewType, resource);
 		const disposables = new DisposableStore();
 		disposables.add(this._workingCopyService.registerWorkingCopy(model));
 		this._models.set(this.key(resource, viewType), { model, disposables });
@@ -39,6 +39,7 @@ export class CustomEditorModelManager implements ICustomEditorModelManager {
 		this._models.forEach((value, key) => {
 			if (model === value.model) {
 				value.disposables.dispose();
+				value.model.dispose();
 				foundKey = key;
 			}
 		});
@@ -48,6 +49,14 @@ export class CustomEditorModelManager implements ICustomEditorModelManager {
 		return;
 	}
 
+	public disposeAllModelsForView(viewType: string): void {
+		this._models.forEach((value) => {
+			if (value.model.viewType === viewType) {
+				this.disposeModel(value.model);
+			}
+		});
+	}
+
 	private key(resource: URI, viewType: string): string {
 		return `${resource.toString()}@@@${viewType}`;
 	}