Support live diagnostics in source generated files

dotnet · Dec 16, 2021 · 1fa1970 · 1fa1970
1 parent 9e35a98
commit 1fa1970
Show file tree

Hide file tree

Showing 12 changed files with 148 additions and 7 deletions.
diff --git a/...Features/Core/Portable/Diagnostics/EngineV2/DiagnosticIncrementalAnalyzer.ProjectState.cs b/...Features/Core/Portable/Diagnostics/EngineV2/DiagnosticIncrementalAnalyzer.ProjectState.cs
@@ -199,7 +199,8 @@ public async ValueTask SaveToInMemoryStorageAsync(Project project, DiagnosticAna
                 var serializerVersion = result.Version;
                 foreach (var documentId in result.DocumentIds)
                 {
-                    var document = project.GetTextDocument(documentId);
+                    var document = project.GetTextDocument(documentId)
+                        ?? project.TryGetSourceGeneratedDocumentForAlreadyGeneratedId(documentId);
                     if (document == null)
                     {
                         // it can happen with build synchronization since, in build case, 

diff --git a/...atures/Core/Portable/Diagnostics/EngineV2/DiagnosticIncrementalAnalyzer_GetDiagnostics.cs b/...atures/Core/Portable/Diagnostics/EngineV2/DiagnosticIncrementalAnalyzer_GetDiagnostics.cs
@@ -204,7 +204,8 @@ private static async Task<ImmutableArray<DiagnosticData>> GetProjectStateDiagnos
                 if (documentId != null)
                 {
                     // file doesn't exist in current solution
-                    var document = project.Solution.GetDocument(documentId);
+                    var document = project.Solution.GetDocument(documentId)
+                        ?? project.TryGetSourceGeneratedDocumentForAlreadyGeneratedId(documentId);
                     if (document == null)
                     {
                         return ImmutableArray<DiagnosticData>.Empty;

diff --git a/...s/Core/Portable/Diagnostics/EngineV2/DiagnosticIncrementalAnalyzer_IncrementalAnalyzer.cs b/...s/Core/Portable/Diagnostics/EngineV2/DiagnosticIncrementalAnalyzer_IncrementalAnalyzer.cs
@@ -535,7 +535,8 @@ private void RaiseProjectDiagnosticsCreated(Project project, StateSet stateSet,
 
             foreach (var documentId in newAnalysisResult.DocumentIds)
             {
-                var document = project.GetTextDocument(documentId);
+                var document = project.GetTextDocument(documentId)
+                    ?? project.TryGetSourceGeneratedDocumentForAlreadyGeneratedId(documentId);
                 if (document == null)
                 {
                     // it can happen with build synchronization since, in build case, 

diff --git a/src/Features/Core/Portable/SolutionCrawler/WorkCoordinator.NormalPriorityProcessor.cs b/src/Features/Core/Portable/SolutionCrawler/WorkCoordinator.NormalPriorityProcessor.cs
@@ -334,7 +334,7 @@ private async Task ProcessDocumentAsync(ImmutableArray<IIncrementalAnalyzer> ana
                         {
                             using (Logger.LogBlock(FunctionId.WorkCoordinator_ProcessDocumentAsync, w => w.ToString(), workItem, cancellationToken))
                             {
-                                var textDocument = solution.GetTextDocument(documentId);
+                                var textDocument = solution.GetTextDocument(documentId) ?? await solution.GetSourceGeneratedDocumentAsync(documentId, cancellationToken).ConfigureAwait(false);
 
                                 if (textDocument != null)
                                 {

diff --git a/src/Features/Core/Portable/SolutionCrawler/WorkCoordinator.cs b/src/Features/Core/Portable/SolutionCrawler/WorkCoordinator.cs
@@ -76,6 +76,8 @@ public WorkCoordinator(
                     _registration.Workspace.WorkspaceChanged += OnWorkspaceChanged;
                     _registration.Workspace.DocumentOpened += OnDocumentOpened;
                     _registration.Workspace.DocumentClosed += OnDocumentClosed;
+                    _registration.Workspace.SourceGeneratedDocumentOpened += OnDocumentOpened;
+                    _registration.Workspace.SourceGeneratedDocumentClosed += OnDocumentClosed;
                 }
 
                 // subscribe to option changed event after all required fields are set
@@ -107,6 +109,8 @@ public void Shutdown(bool blockingShutdown)
                 _registration.Workspace.WorkspaceChanged -= OnWorkspaceChanged;
                 _registration.Workspace.DocumentOpened -= OnDocumentOpened;
                 _registration.Workspace.DocumentClosed -= OnDocumentClosed;
+                _registration.Workspace.SourceGeneratedDocumentOpened -= OnDocumentOpened;
+                _registration.Workspace.SourceGeneratedDocumentClosed -= OnDocumentClosed;
 
                 // cancel any pending blocks
                 _shutdownNotificationSource.Cancel();
@@ -151,12 +155,16 @@ private void OnOptionChanged(object? sender, OptionChangedEventArgs e)
                         _registration.Workspace.WorkspaceChanged += OnWorkspaceChanged;
                         _registration.Workspace.DocumentOpened += OnDocumentOpened;
                         _registration.Workspace.DocumentClosed += OnDocumentClosed;
+                        _registration.Workspace.SourceGeneratedDocumentOpened += OnDocumentOpened;
+                        _registration.Workspace.SourceGeneratedDocumentClosed += OnDocumentClosed;
                     }
                     else
                     {
                         _registration.Workspace.WorkspaceChanged -= OnWorkspaceChanged;
                         _registration.Workspace.DocumentOpened -= OnDocumentOpened;
                         _registration.Workspace.DocumentClosed -= OnDocumentClosed;
+                        _registration.Workspace.SourceGeneratedDocumentOpened -= OnDocumentOpened;
+                        _registration.Workspace.SourceGeneratedDocumentClosed -= OnDocumentClosed;
                     }
 
                     SolutionCrawlerLogger.LogOptionChanged(CorrelationId, value);

diff --git a/src/Features/Core/Portable/Workspace/BackgroundCompiler.cs b/src/Features/Core/Portable/Workspace/BackgroundCompiler.cs
@@ -39,6 +39,8 @@ public BackgroundCompiler(Workspace workspace)
             _workspace.WorkspaceChanged += OnWorkspaceChanged;
             _workspace.DocumentOpened += OnDocumentOpened;
             _workspace.DocumentClosed += OnDocumentClosed;
+            _workspace.SourceGeneratedDocumentOpened += OnDocumentOpened;
+            _workspace.SourceGeneratedDocumentClosed += OnDocumentClosed;
         }
 
         public void Dispose()
@@ -49,6 +51,8 @@ public void Dispose()
 
                 _workspace.DocumentClosed -= OnDocumentClosed;
                 _workspace.DocumentOpened -= OnDocumentOpened;
+                _workspace.SourceGeneratedDocumentClosed -= OnDocumentClosed;
+                _workspace.SourceGeneratedDocumentOpened -= OnDocumentOpened;
                 _workspace.WorkspaceChanged -= OnWorkspaceChanged;
 
                 _workspace = null!;

diff --git a/src/Features/Core/Portable/Workspace/BackgroundParser.cs b/src/Features/Core/Portable/Workspace/BackgroundParser.cs
@@ -49,6 +49,8 @@ public BackgroundParser(Workspace workspace)
 
             workspace.DocumentOpened += OnDocumentOpened;
             workspace.DocumentClosed += OnDocumentClosed;
+            workspace.SourceGeneratedDocumentOpened += OnDocumentOpened;
+            workspace.SourceGeneratedDocumentClosed += OnDocumentClosed;
         }
 
         private void OnActiveDocumentChanged(object sender, DocumentId activeDocumentId)

diff --git a/src/Workspaces/Core/Portable/Diagnostics/DiagnosticAnalysisResult.cs b/src/Workspaces/Core/Portable/Diagnostics/DiagnosticAnalysisResult.cs
@@ -320,7 +320,8 @@ private static void VerifyDocumentMap(Project project, ImmutableDictionary<Docum
         {
             foreach (var documentId in map.Keys)
             {
-                Debug.Assert(project.GetTextDocument(documentId)?.SupportsDiagnostics() == true);
+                var textDocument = project.GetTextDocument(documentId) ?? project.TryGetSourceGeneratedDocumentForAlreadyGeneratedId(documentId);
+                Debug.Assert(textDocument?.SupportsDiagnostics() == true);
             }
         }
     }

diff --git a/src/Workspaces/Core/Portable/Diagnostics/Extensions.cs b/src/Workspaces/Core/Portable/Diagnostics/Extensions.cs
@@ -59,7 +59,8 @@ public static async ValueTask<Location> ConvertLocationAsync(
                 return Location.None;
             }
 
-            var textDocument = project.GetTextDocument(dataLocation.DocumentId);
+            var textDocument = project.GetTextDocument(dataLocation.DocumentId)
+                ?? await project.GetSourceGeneratedDocumentAsync(dataLocation.DocumentId, cancellationToken).ConfigureAwait(false);
             if (textDocument == null)
             {
                 return Location.None;

diff --git a/src/Workspaces/Core/Portable/Workspace/Solution/SourceGeneratedDocumentState.cs b/src/Workspaces/Core/Portable/Workspace/Solution/SourceGeneratedDocumentState.cs
@@ -36,7 +36,7 @@ public static SourceGeneratedDocumentState Create(
                 documentIdentity,
                 languageServices,
                 solutionServices,
-                documentServiceProvider: null,
+                documentServiceProvider: SourceGeneratedTextDocumentServiceProvider.Instance,
                 new DocumentInfo.DocumentAttributes(
                     documentIdentity.DocumentId,
                     name: documentIdentity.HintName,
@@ -89,5 +89,41 @@ public SourceGeneratedDocumentState WithUpdatedGeneratedContent(SourceText sourc
                 this.LanguageServices,
                 this.solutionServices);
         }
+
+        internal sealed class SourceGeneratedTextDocumentServiceProvider : IDocumentServiceProvider
+        {
+            public static readonly SourceGeneratedTextDocumentServiceProvider Instance = new();
+
+            private SourceGeneratedTextDocumentServiceProvider()
+            {
+            }
+
+            public TService? GetService<TService>()
+                where TService : class, IDocumentService
+            {
+                // right now, it doesn't implement much services but we expect it to implements all 
+                // document services in future so that we can remove all if branches in feature code
+                // but just delegate work to default document services.
+                if (SourceGeneratedDocumentOperationService.Instance is TService documentOperationService)
+                {
+                    return documentOperationService;
+                }
+
+                if (DocumentPropertiesService.Default is TService documentPropertiesService)
+                {
+                    return documentPropertiesService;
+                }
+
+                return null;
+            }
+
+            private class SourceGeneratedDocumentOperationService : IDocumentOperationService
+            {
+                public static readonly SourceGeneratedDocumentOperationService Instance = new();
+
+                public bool CanApplyChange => false;
+                public bool SupportDiagnostics => true;
+            }
+        }
     }
 }
diff --git a/src/Workspaces/Core/Portable/Workspace/Workspace_Editor.cs b/src/Workspaces/Core/Portable/Workspace/Workspace_Editor.cs
@@ -428,6 +428,15 @@ internal void OnSourceGeneratedDocumentOpened(
                 _openSourceGeneratedDocumentIdentities.Add(documentId, documentIdentity);
 
                 UpdateCurrentContextMapping_NoLock(textContainer, documentId, isCurrentContext: true);
+
+                // Fire and forget that the workspace is changing.
+                _ = RaiseSourceGeneratedDocumentOpenedAsync(this, CurrentSolution, documentId);
+
+                static async Task RaiseSourceGeneratedDocumentOpenedAsync(Workspace workspace, Solution currentSolution, DocumentId documentId)
+                {
+                    var document = await currentSolution.GetSourceGeneratedDocumentAsync(documentId, CancellationToken.None).ConfigureAwait(false);
+                    await workspace.RaiseSourceGeneratedDocumentOpenedEventAsync(document).ConfigureAwait(false);
+                }
             }
 
             this.RegisterText(textContainer);
@@ -441,6 +450,15 @@ internal void OnSourceGeneratedDocumentClosed(DocumentId documentId)
 
                 Contract.ThrowIfFalse(_openSourceGeneratedDocumentIdentities.Remove(documentId));
                 ClearOpenDocument(documentId);
+
+                // Fire and forget that the workspace is changing.
+                _ = RaiseSourceGeneratedDocumentClosedAsync(this, CurrentSolution, documentId);
+
+                static async Task RaiseSourceGeneratedDocumentClosedAsync(Workspace workspace, Solution currentSolution, DocumentId documentId)
+                {
+                    var document = await currentSolution.GetSourceGeneratedDocumentAsync(documentId, CancellationToken.None).ConfigureAwait(false);
+                    await workspace.RaiseSourceGeneratedDocumentClosedEventAsync(document).ConfigureAwait(false);
+                }
             }
         }
 

diff --git a/src/Workspaces/Core/Portable/Workspace/Workspace_Events.cs b/src/Workspaces/Core/Portable/Workspace/Workspace_Events.cs
@@ -22,6 +22,8 @@ public abstract partial class Workspace
         private const string WorkspaceFailedEventName = "WorkspaceFailed";
         private const string DocumentOpenedEventName = "DocumentOpened";
         private const string DocumentClosedEventName = "DocumentClosed";
+        private const string SourceGeneratedDocumentOpenedEventName = "SourceGeneratedDocumentOpened";
+        private const string SourceGeneratedDocumentClosedEventName = "SourceGeneratedDocumentClosed";
         private const string DocumentActiveContextChangedName = "DocumentActiveContextChanged";
 
         /// <summary>
@@ -168,6 +170,72 @@ protected Task RaiseDocumentClosedEventAsync(Document document)
             }
         }
 
+        /// <summary>
+        /// An event that is fired when a documents is opened in the editor.
+        /// </summary>
+        internal event EventHandler<DocumentEventArgs> SourceGeneratedDocumentOpened
+        {
+            add
+            {
+                _eventMap.AddEventHandler(SourceGeneratedDocumentOpenedEventName, value);
+            }
+
+            remove
+            {
+                _eventMap.RemoveEventHandler(SourceGeneratedDocumentOpenedEventName, value);
+            }
+        }
+
+        private protected Task RaiseSourceGeneratedDocumentOpenedEventAsync(Document document)
+        {
+            var ev = GetEventHandlers<DocumentEventArgs>(SourceGeneratedDocumentOpenedEventName);
+            if (ev.HasHandlers && document != null)
+            {
+                return this.ScheduleTask(() =>
+                {
+                    var args = new DocumentEventArgs(document);
+                    ev.RaiseEvent(handler => handler(this, args));
+                }, SourceGeneratedDocumentOpenedEventName);
+            }
+            else
+            {
+                return Task.CompletedTask;
+            }
+        }
+
+        /// <summary>
+        /// An event that is fired when a document is closed in the editor.
+        /// </summary>
+        internal event EventHandler<DocumentEventArgs> SourceGeneratedDocumentClosed
+        {
+            add
+            {
+                _eventMap.AddEventHandler(SourceGeneratedDocumentClosedEventName, value);
+            }
+
+            remove
+            {
+                _eventMap.RemoveEventHandler(SourceGeneratedDocumentClosedEventName, value);
+            }
+        }
+
+        private protected Task RaiseSourceGeneratedDocumentClosedEventAsync(Document document)
+        {
+            var ev = GetEventHandlers<DocumentEventArgs>(SourceGeneratedDocumentClosedEventName);
+            if (ev.HasHandlers && document != null)
+            {
+                return this.ScheduleTask(() =>
+                {
+                    var args = new DocumentEventArgs(document);
+                    ev.RaiseEvent(handler => handler(this, args));
+                }, SourceGeneratedDocumentClosedEventName);
+            }
+            else
+            {
+                return Task.CompletedTask;
+            }
+        }
+
         /// <summary>
         /// An event that is fired when the active context document associated with a buffer 
         /// changes.