@@ -8,6 +8,7 @@ import * as RoslynProtocol from './roslynProtocol';
88import { RoslynLanguageServer } from './roslynLanguageServer' ;
99import { UriConverter } from './uriConverter' ;
1010import * as lsp from 'vscode-languageserver-protocol' ;
11+ import { IDisposable } from '@microsoft/servicehub-framework' ;
1112
1213export function registerSourceGeneratedFilesContentProvider (
1314 context : vscode . ExtensionContext ,
@@ -16,16 +17,91 @@ export function registerSourceGeneratedFilesContentProvider(
1617 context . subscriptions . push (
1718 vscode . workspace . registerTextDocumentContentProvider (
1819 'roslyn-source-generated' ,
19- new ( class implements vscode . TextDocumentContentProvider {
20- async provideTextDocumentContent ( uri : vscode . Uri , token : vscode . CancellationToken ) : Promise < string > {
21- const result = await languageServer . sendRequest (
22- RoslynProtocol . SourceGeneratorGetTextRequest . type ,
23- { textDocument : lsp . TextDocumentIdentifier . create ( UriConverter . serialize ( uri ) ) } ,
24- token
25- ) ;
26- return result . text ;
27- }
28- } ) ( )
20+ new RoslynSourceGeneratedContentProvider ( languageServer )
2921 )
3022 ) ;
3123}
24+
25+ class RoslynSourceGeneratedContentProvider implements vscode . TextDocumentContentProvider , IDisposable {
26+ private _onDidChangeEmitter : vscode . EventEmitter < vscode . Uri > = new vscode . EventEmitter < vscode . Uri > ( ) ;
27+
28+ // Stores all the source generated documents that we have opened so far and their up to date content.
29+ private _openedDocuments : Map < vscode . Uri , RoslynProtocol . SourceGeneratedDocumentText > = new Map ( ) ;
30+
31+ // Since we could potentially have multiple refresh notifications in flight at the same time,
32+ // we use a simple queue to ensure that updates to our state map only happen serially.
33+ private _updateQueue ?: Promise < RoslynProtocol . SourceGeneratedDocumentText > ;
34+
35+ private _cancellationSource = new vscode . CancellationTokenSource ( ) ;
36+
37+ constructor ( private languageServer : RoslynLanguageServer ) {
38+ languageServer . registerOnNotification (
39+ RoslynProtocol . RefreshSourceGeneratedDocumentNotification . method ,
40+ async ( ) => {
41+ this . _openedDocuments . forEach ( async ( _ , key ) => {
42+ await this . enqueueDocumentUpdateAsync ( key , this . _cancellationSource . token ) ;
43+ this . _onDidChangeEmitter . fire ( key ) ;
44+ } ) ;
45+ }
46+ ) ;
47+ vscode . workspace . onDidCloseTextDocument ( ( document ) => {
48+ const openedDoc = this . _openedDocuments . get ( document . uri ) ;
49+ if ( openedDoc !== undefined ) {
50+ this . _openedDocuments . delete ( document . uri ) ;
51+ }
52+ } ) ;
53+ }
54+
55+ public onDidChange : vscode . Event < vscode . Uri > = this . _onDidChangeEmitter . event ;
56+
57+ async provideTextDocumentContent ( uri : vscode . Uri , token : vscode . CancellationToken ) : Promise < string > {
58+ let content = this . _openedDocuments . get ( uri ) ;
59+
60+ if ( ! content ) {
61+ // We're being asked about this document for the first time, so we need to fetch it from the server.
62+ content = await this . enqueueDocumentUpdateAsync ( uri , token ) ;
63+ }
64+
65+ return content . text ?? vscode . l10n . t ( 'Generated document not found' ) ;
66+ }
67+
68+ private async enqueueDocumentUpdateAsync (
69+ uri : vscode . Uri ,
70+ token : vscode . CancellationToken
71+ ) : Promise < RoslynProtocol . SourceGeneratedDocumentText > {
72+ if ( ! this . _updateQueue ) {
73+ this . _updateQueue = this . updateDocumentAsync ( uri , token ) ;
74+ } else {
75+ this . _updateQueue = this . _updateQueue . then ( async ( ) => await this . updateDocumentAsync ( uri , token ) ) ;
76+ }
77+
78+ return await this . _updateQueue ;
79+ }
80+
81+ private async updateDocumentAsync (
82+ uri : vscode . Uri ,
83+ token : vscode . CancellationToken
84+ ) : Promise < RoslynProtocol . SourceGeneratedDocumentText > {
85+ const currentContent = this . _openedDocuments . get ( uri ) ;
86+ const newContent = await this . languageServer . sendRequest (
87+ RoslynProtocol . SourceGeneratorGetTextRequest . type ,
88+ {
89+ textDocument : lsp . TextDocumentIdentifier . create ( UriConverter . serialize ( uri ) ) ,
90+ resultId : currentContent ?. resultId ,
91+ } ,
92+ token
93+ ) ;
94+
95+ // If we had no content before, or the resultId has changed, update the content
96+ if ( ! currentContent || newContent . resultId !== currentContent ?. resultId ) {
97+ this . _openedDocuments . set ( uri , newContent ) ;
98+ return newContent ;
99+ }
100+
101+ return currentContent ;
102+ }
103+
104+ dispose ( ) : void {
105+ this . _cancellationSource . cancel ( ) ;
106+ }
107+ }
0 commit comments