testing: fix improve test decoration performance (#173082)

Fixes #170493. Goes from a 1600ms hang to an 8ms computation.

src/vs/workbench/contrib/testing/browser/testingDecorations.ts

@@ -71,6 +71,68 @@ interface ITestDecoration extends IPublicTestDecoration {
 	click(e: IEditorMouseEvent): boolean;
 }
 
+/** Value for saved decorations, providing fast accessors for the hot 'syncDecorations' path */
+class CachedDecorations {
+	private readonly runByIdKey = new Map<string, RunTestDecoration>();
+	private readonly messages = new Map<ITestMessage, TestMessageDecoration>();
+
+	public get size() {
+		return this.runByIdKey.size + this.messages.size;
+	}
+
+	/** Gets a test run decoration that contains exactly the given test IDs */
+	public getForExactTests(testIds: string[]) {
+		const key = testIds.sort().join('\0\0');
+		return this.runByIdKey.get(key);
+	}
+
+	/** Gets the decoration that corresponds to the given test message */
+	public getMessage(message: ITestMessage) {
+		return this.messages.get(message);
+	}
+
+	/** Removes the decoration for the given test messsage */
+	public removeMessage(message: ITestMessage) {
+		this.messages.delete(message);
+	}
+
+	/** Adds a new test message decoration */
+	public addMessage(d: TestMessageDecoration) {
+		this.messages.set(d.testMessage, d);
+	}
+
+	/** Adds a new test run decroation */
+	public addTest(d: RunTestDecoration) {
+		const key = d.testIds.sort().join('\0\0');
+		this.runByIdKey.set(key, d);
+	}
+
+	/** Finds an extension by VS Code event ID */
+	public getById(decorationId: string) {
+		for (const d of this.runByIdKey.values()) {
+			if (d.id === decorationId) {
+				return d;
+			}
+		}
+		for (const d of this.messages.values()) {
+			if (d.id === decorationId) {
+				return d;
+			}
+		}
+		return undefined;
+	}
+
+	/** Iterate over all decorations */
+	*[Symbol.iterator](): IterableIterator<ITestDecoration> {
+		for (const d of this.runByIdKey.values()) {
+			yield d;
+		}
+		for (const d of this.messages.values()) {
+			yield d;
+		}
+	}
+}
+
 export class TestingDecorationService extends Disposable implements ITestingDecorationsService {
 	declare public _serviceBrand: undefined;
 
@@ -81,8 +143,7 @@ export class TestingDecorationService extends Disposable implements ITestingDeco
 		rangeUpdateVersionId?: number;
 		/** Counter for the results rendered in the document */
 		generation: number;
-		value: readonly ITestDecoration[];
-		map: ReadonlyMap<string, ITestDecoration>;
+		value: CachedDecorations;
 	}>();
 
 	/**
@@ -153,15 +214,15 @@ export class TestingDecorationService extends Disposable implements ITestingDeco
 	}
 
 	/** @inheritdoc */
-	public syncDecorations(resource: URI): ReadonlyMap<string, ITestDecoration> {
+	public syncDecorations(resource: URI): CachedDecorations {
 		const model = this.modelService.getModel(resource);
 		if (!model) {
-			return new Map();
+			return new CachedDecorations();
 		}
 
 		const cached = this.decorationCache.get(resource);
 		if (cached && cached.generation === this.generation && (cached.rangeUpdateVersionId === undefined || cached.rangeUpdateVersionId !== model.getVersionId())) {
-			return cached.map;
+			return cached.value;
 		}
 
 		return this.applyDecorations(model);
@@ -174,7 +235,7 @@ export class TestingDecorationService extends Disposable implements ITestingDeco
 			return undefined;
 		}
 
-		const decoration = Iterable.find(this.syncDecorations(resource).values(), v => v instanceof RunTestDecoration && v.isForTest(testId));
+		const decoration = Iterable.find(this.syncDecorations(resource), v => v instanceof RunTestDecoration && v.isForTest(testId));
 		if (!decoration) {
 			return undefined;
 		}
@@ -195,10 +256,10 @@ export class TestingDecorationService extends Disposable implements ITestingDeco
 		const uriStr = model.uri.toString();
 		const cached = this.decorationCache.get(model.uri);
 		const testRangesUpdated = cached?.rangeUpdateVersionId === model.getVersionId();
-		const lastDecorations = cached?.value ?? [];
+		const lastDecorations = cached?.value ?? new CachedDecorations();
 
-		const map = model.changeDecorations(accessor => {
-			const newDecorations: ITestDecoration[] = [];
+		const newDecorations = model.changeDecorations(accessor => {
+			const newDecorations = new CachedDecorations();
 			const runDecorations = new TestDecorations<{ line: number; id: ''; test: IncrementalTestCollectionItem; resultItem: TestResultItem | undefined }>();
 			for (const test of this.testService.collection.getNodeByUrl(model.uri)) {
 				if (!test.item.range) {
@@ -212,7 +273,7 @@ export class TestingDecorationService extends Disposable implements ITestingDeco
 
 			for (const [line, tests] of runDecorations.lines()) {
 				const multi = tests.length > 1;
-				let existing = lastDecorations.find(d => d instanceof RunTestDecoration && d.exactlyContainsTests(tests)) as RunTestDecoration | undefined;
+				let existing = lastDecorations.getForExactTests(tests.map(t => t.test.item.extId));
 
 				// see comment in the constructor for what's going on here
 				if (existing && testRangesUpdated && model.getDecorationRange(existing.id)?.startLineNumber !== line) {
@@ -223,9 +284,9 @@ export class TestingDecorationService extends Disposable implements ITestingDeco
 					if (existing.replaceOptions(tests, gutterEnabled)) {
 						accessor.changeDecorationOptions(existing.id, existing.editorDecoration.options);
 					}
-					newDecorations.push(existing);
+					newDecorations.addTest(existing);
 				} else {
-					newDecorations.push(multi
+					newDecorations.addTest(multi
 						? this.instantiationService.createInstance(MultiRunTestDecoration, tests, gutterEnabled, model)
 						: this.instantiationService.createInstance(RunSingleTestDecoration, tests[0].test, tests[0].resultItem, model, gutterEnabled));
 				}
@@ -236,14 +297,13 @@ export class TestingDecorationService extends Disposable implements ITestingDeco
 				for (const task of lastResult.tasks) {
 					for (const m of task.otherMessages) {
 						if (!this.invalidatedMessages.has(m) && m.location?.uri.toString() === uriStr) {
-							const decoration = lastDecorations.find(l => l instanceof TestMessageDecoration && l.testMessage === m)
-								|| this.instantiationService.createInstance(TestMessageDecoration, m, undefined, model);
-							newDecorations.push(decoration);
+							const decoration = lastDecorations.getMessage(m) || this.instantiationService.createInstance(TestMessageDecoration, m, undefined, model);
+							newDecorations.addMessage(decoration);
 						}
 					}
 				}
 
-				const messageLines = new Map</* line number */ number, /* index in newDecorations */ number>();
+				const messageLines = new Map</* line number */ number, /* last test message */ ITestMessage>();
 				for (const test of lastResult.tests) {
 					for (let taskId = 0; taskId < test.tasks.length; taskId++) {
 						const state = test.tasks[taskId];
@@ -256,29 +316,20 @@ export class TestingDecorationService extends Disposable implements ITestingDeco
 							// Only add one message per line number. Overlapping messages
 							// don't appear well, and the peek will show all of them (#134129)
 							const line = m.location.range.startLineNumber;
-							let index: number;
 							if (messageLines.has(line)) {
-								index = messageLines.get(line)!;
-							} else {
-								index = newDecorations.length;
-								messageLines.set(line, index);
+								newDecorations.removeMessage(messageLines.get(line)!);
 							}
 
-							const previous = lastDecorations.find(l => l instanceof TestMessageDecoration && l.testMessage === m);
-							if (previous) {
-								newDecorations[index] = previous;
-								continue;
-							}
-
-							const messageUri = buildTestUri({
+							const decoration = lastDecorations.getMessage(m) || this.instantiationService.createInstance(TestMessageDecoration, m, buildTestUri({
 								type: TestUriType.ResultActualOutput,
 								messageIndex: i,
 								taskIndex: taskId,
 								resultId: lastResult.id,
 								testExtId: test.item.extId,
-							});
+							}), model);
 
-							newDecorations.push(this.instantiationService.createInstance(TestMessageDecoration, m, messageUri, model));
+							newDecorations.addMessage(decoration);
+							messageLines.set(line, decoration.testMessage);
 						}
 					}
 				}
@@ -299,18 +350,16 @@ export class TestingDecorationService extends Disposable implements ITestingDeco
 				}
 			}
 
-			const map = new Map(newDecorations.map(d => [d.id, d]));
 			this.decorationCache.set(model.uri, {
 				generation: this.generation,
 				rangeUpdateVersionId: cached?.rangeUpdateVersionId,
 				value: newDecorations,
-				map,
 			});
 
-			return map;
+			return newDecorations;
 		});
 
-		return map || new Map();
+		return newDecorations || lastDecorations;
 	}
 }
 
@@ -347,9 +396,13 @@ export class TestingDecorations extends Disposable implements IEditorContributio
 		this._register(this.editor.onMouseDown(e => {
 			if (e.target.position && this.currentUri) {
 				const modelDecorations = editor.getModel()?.getDecorationsInRange(Range.fromPositions(e.target.position)) ?? [];
+				if (!modelDecorations.length) {
+					return;
+				}
+
+				const cache = decorations.syncDecorations(this.currentUri);
 				for (const { id } of modelDecorations) {
-					const cache = decorations.syncDecorations(this.currentUri);
-					if ((cache.get(id) as ITestDecoration | undefined)?.click(e)) {
+					if ((cache.getById(id) as ITestDecoration | undefined)?.click(e)) {
 						e.event.stopPropagation();
 						return;
 					}
@@ -371,7 +424,7 @@ export class TestingDecorations extends Disposable implements IEditorContributio
 				for (const change of e.changes) {
 					const modelDecorations = model.getLinesDecorations(change.range.startLineNumber, change.range.endLineNumber);
 					for (const { id } of modelDecorations) {
-						const decoration = currentDecorations.get(id);
+						const decoration = currentDecorations.getById(id);
 						if (decoration instanceof TestMessageDecoration) {
 							decorations.invalidateResultMessage(decoration.testMessage);
 						}
@@ -618,7 +671,12 @@ abstract class RunTestDecoration {
 		return this.editorDecoration.range.startLineNumber;
 	}
 
+	public get testIds() {
+		return this.tests.map(t => t.test.item.extId);
+	}
+
 	public editorDecoration: IModelDeltaDecoration;
+	public displayedStates: readonly (TestResultState | undefined)[];
 
 	constructor(
 		protected tests: readonly {
@@ -636,6 +694,7 @@ abstract class RunTestDecoration {
 		@IContextKeyService protected readonly contextKeyService: IContextKeyService,
 		@IMenuService protected readonly menuService: IMenuService,
 	) {
+		this.displayedStates = tests.map(t => t.resultItem?.computedState);
 		this.editorDecoration = createRunTestDecoration(tests.map(t => t.test), tests.map(t => t.resultItem), visible);
 		this.editorDecoration.options.glyphMarginHoverMessage = new MarkdownString().appendText(this.getGutterLabel());
 	}
@@ -667,27 +726,6 @@ abstract class RunTestDecoration {
 		return true;
 	}
 
-	public exactlyContainsTests(tests: readonly { test: IncrementalTestCollectionItem }[]): boolean {
-		if (tests.length !== this.tests.length) {
-			return false;
-		}
-		if (tests.length === 1) {
-			return tests[0].test.item.extId === this.tests[0].test.item.extId;
-		}
-
-		const ownTests = new Set();
-		for (const t of this.tests) {
-			ownTests.add(t.test.item.extId);
-		}
-		for (const t of tests) {
-			if (!ownTests.delete(t.test.item.extId)) {
-				return false;
-			}
-		}
-
-		return true;
-	}
-
 	/**
 	 * Updates the decoration to match the new set of tests.
 	 * @returns true if options were changed, false otherwise
@@ -696,13 +734,13 @@ abstract class RunTestDecoration {
 		test: IncrementalTestCollectionItem;
 		resultItem: TestResultItem | undefined;
 	}[], visible: boolean): boolean {
-		if (visible === this.visible
-			&& equals(this.tests.map(t => t.test.item.extId), newTests.map(t => t.test.item.extId))
-			&& this.tests.map(t => t.resultItem?.computedState) === newTests.map(t => t.resultItem?.computedState)) {
+		const displayedStates = newTests.map(t => t.resultItem?.computedState);
+		if (visible === this.visible && equals(this.displayedStates, displayedStates)) {
 			return false;
 		}
 
 		this.tests = newTests;
+		this.displayedStates = displayedStates;
 		this.visible = visible;
 		this.editorDecoration.options = createRunTestDecoration(newTests.map(t => t.test), newTests.map(t => t.resultItem), visible).options;
 		this.editorDecoration.options.glyphMarginHoverMessage = new MarkdownString().appendText(this.getGutterLabel());
@@ -836,14 +874,6 @@ abstract class RunTestDecoration {
 }
 
 class MultiRunTestDecoration extends RunTestDecoration implements ITestDecoration {
-	public get testIds() {
-		return this.tests.map(t => t.test.item.extId);
-	}
-
-	public get displayedStates() {
-		return this.tests.map(t => t.resultItem?.computedState);
-	}
-
 	protected override getContextMenuActions() {
 		const allActions: IAction[] = [];
 		if (this.tests.some(({ test }) => this.testProfileService.capabilitiesForTest(test) & TestRunProfileBitset.Run)) {

src/vs/workbench/contrib/testing/common/testingDecorations.ts

@@ -30,7 +30,10 @@ export interface ITestingDecorationsService {
 	 * Ensures decorations in the given document URI are up to date,
 	 * and returns them.
 	 */
-	syncDecorations(resource: URI): ReadonlyMap<string, ITestDecoration>;
+	syncDecorations(resource: URI): Iterable<ITestDecoration> & {
+		readonly size: number;
+		getById(decorationId: string): ITestDecoration | undefined;
+	};
 
 	/**
 	 * Gets the range where a test ID is displayed, in the given URI.