Merge two dictionaries into one

Author: DustinCampbell

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Diagnostics/RazorDiagnosticsPublisher.TestAccessor.cs

@@ -44,19 +44,11 @@ public Task WaitForDiagnosticsToPublishAsync()
             return instance._workQueue.WaitUntilCurrentBatchCompletesAsync();
         }
 
-        public void SetPublishedCSharpDiagnostics(string filePath, ImmutableArray<Diagnostic> diagnostics)
+        public void SetPublishedDiagnostics(string filePath, RazorDiagnostic[] razorDiagnostics, Diagnostic[]? csharpDiagnostics)
         {
-            lock (instance._publishedDiagnosticsGate)
+            lock (instance._publishedDiagnostics)
             {
-                instance._publishedCSharpDiagnostics[filePath] = diagnostics;
-            }
-        }
-
-        public void SetPublishedRazorDiagnostics(string filePath, ImmutableArray<RazorDiagnostic> diagnostics)
-        {
-            lock (instance._publishedDiagnosticsGate)
-            {
-                instance._publishedRazorDiagnostics[filePath] = diagnostics;
+                instance._publishedDiagnostics[filePath] = new(razorDiagnostics, csharpDiagnostics);
             }
         }
 

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Diagnostics/RazorDiagnosticsPublisher.cs

@@ -24,6 +24,11 @@ namespace Microsoft.AspNetCore.Razor.LanguageServer.Diagnostics;
 
 internal partial class RazorDiagnosticsPublisher : IDocumentProcessedListener, IDisposable
 {
+    private readonly record struct PublishedDiagnostics(IReadOnlyList<RazorDiagnostic> Razor, Diagnostic[]? CSharp)
+    {
+        public int Count => Razor.Count + (CSharp?.Length ?? 0);
+    }
+
     private static readonly TimeSpan s_publishDelay = TimeSpan.FromSeconds(2);
     private static readonly TimeSpan s_clearClosedDocumentsDelay = TimeSpan.FromSeconds(5);
 
@@ -36,10 +41,7 @@ internal partial class RazorDiagnosticsPublisher : IDocumentProcessedListener, I
 
     private readonly CancellationTokenSource _disposeTokenSource;
     private readonly AsyncBatchingWorkQueue<IDocumentSnapshot> _workQueue;
-
-    private readonly object _publishedDiagnosticsGate = new();
-    private readonly Dictionary<string, IReadOnlyList<RazorDiagnostic>> _publishedRazorDiagnostics;
-    private readonly Dictionary<string, IReadOnlyList<Diagnostic>> _publishedCSharpDiagnostics;
+    private readonly Dictionary<string, PublishedDiagnostics> _publishedDiagnostics;
 
     private readonly object _documentClosedGate = new();
     private Task _clearClosedDocumentsTask = Task.CompletedTask;
@@ -77,8 +79,7 @@ protected RazorDiagnosticsPublisher(
         _disposeTokenSource = new();
         _workQueue = new AsyncBatchingWorkQueue<IDocumentSnapshot>(publishDelay, ProcessBatchAsync, _disposeTokenSource.Token);
 
-        _publishedRazorDiagnostics = new Dictionary<string, IReadOnlyList<RazorDiagnostic>>(FilePathComparer.Instance);
-        _publishedCSharpDiagnostics = new Dictionary<string, IReadOnlyList<Diagnostic>>(FilePathComparer.Instance);
+        _publishedDiagnostics = new Dictionary<string, PublishedDiagnostics>(FilePathComparer.Instance);
         _logger = loggerFactory.GetOrCreateLogger<RazorDiagnosticsPublisher>();
     }
 
@@ -114,22 +115,28 @@ private async Task PublishDiagnosticsAsync(IDocumentSnapshot document, Cancellat
         var csharpDiagnostics = await GetCSharpDiagnosticsAsync(document, token).ConfigureAwait(false);
         var razorDiagnostics = result.GetCSharpDocument().Diagnostics;
 
-        lock (_publishedDiagnosticsGate)
+        lock (_publishedDiagnostics)
         {
             var filePath = document.FilePath.AssumeNotNull();
 
-            if (_publishedRazorDiagnostics.TryGetValue(filePath, out var previousRazorDiagnostics) && razorDiagnostics.SequenceEqual(previousRazorDiagnostics)
-                && (csharpDiagnostics == null || (_publishedCSharpDiagnostics.TryGetValue(filePath, out var previousCsharpDiagnostics) && csharpDiagnostics.SequenceEqual(previousCsharpDiagnostics))))
+            // See if these are the same diagnostics as last time. If so, we don't need to publish.
+            if (_publishedDiagnostics.TryGetValue(filePath, out var previousDiagnostics))
             {
-                // Diagnostics are the same as last publish
-                return;
-            }
+                var sameRazorDiagnostics = razorDiagnostics.SequenceEqual(previousDiagnostics.Razor);
+                var sameCSharpDiagnostics = (csharpDiagnostics, previousDiagnostics.CSharp) switch
+                {
+                    (not null, not null) => csharpDiagnostics.SequenceEqual(previousDiagnostics.CSharp),
+                    (null, null) => true,
+                    _ => false,
+                };
 
-            _publishedRazorDiagnostics[filePath] = razorDiagnostics;
-            if (csharpDiagnostics != null)
-            {
-                _publishedCSharpDiagnostics[filePath] = csharpDiagnostics;
+                if (sameRazorDiagnostics && sameCSharpDiagnostics)
+                {
+                    return;
+                }
             }
+
+            _publishedDiagnostics[filePath] = new(razorDiagnostics, csharpDiagnostics);
         }
 
         if (!document.TryGetText(out var sourceText))
@@ -138,8 +145,11 @@ private async Task PublishDiagnosticsAsync(IDocumentSnapshot document, Cancellat
             return;
         }
 
-        var convertedDiagnostics = razorDiagnostics.Select(razorDiagnostic => RazorDiagnosticConverter.Convert(razorDiagnostic, sourceText, document));
-        var combinedDiagnostics = csharpDiagnostics == null ? convertedDiagnostics : convertedDiagnostics.Concat(csharpDiagnostics);
+        Diagnostic[] combinedDiagnostics = [
+            .. razorDiagnostics.Select(d => RazorDiagnosticConverter.Convert(d, sourceText, document)),
+            .. csharpDiagnostics ?? []
+        ];
+
         PublishDiagnosticsForFilePath(document.FilePath, combinedDiagnostics);
 
         if (_logger.IsEnabled(LogLevel.Trace))
@@ -220,51 +230,45 @@ async Task ClearClosedDocumentsAfterDelayAsync()
 
     private void ClearClosedDocuments()
     {
-        lock (_publishedDiagnosticsGate)
+        lock (_publishedDiagnostics)
         {
-            ClearClosedDocumentsPublishedDiagnostics(_publishedRazorDiagnostics);
-            ClearClosedDocumentsPublishedDiagnostics(_publishedCSharpDiagnostics);
+            using var documentsToRemove = new PooledArrayBuilder<(string filePath, bool publishToClear)>(capacity: _publishedDiagnostics.Count);
 
-            if (_publishedRazorDiagnostics.Count > 0 || _publishedCSharpDiagnostics.Count > 0)
+            foreach (var (filePath, diagnostics) in _publishedDiagnostics)
             {
-                // There's no way for us to know when a document is closed at this layer. Therefore, we need to poll every X seconds
-                // and check if the currently tracked documents are closed. In practice this work is super minimal.
-                StartDelayToClearDocuments();
+                if (!_projectManager.IsDocumentOpen(filePath))
+                {
+                    // If there were previously published diagnostics for this document, take note so
+                    // we can publish an empty set of diagnostics.
+                    documentsToRemove.Add((filePath, publishToClear: diagnostics.Count > 0));
+                }
             }
 
-            void ClearClosedDocumentsPublishedDiagnostics<T>(Dictionary<string, IReadOnlyList<T>> publishedDiagnostics) where T : class
+            if (documentsToRemove.Count == 0)
             {
-                using var documentsToRemove = new PooledArrayBuilder<(string key, bool publishEmptyDiagnostics)>(capacity: publishedDiagnostics.Count);
+                return;
+            }
 
-                foreach (var (key, value) in publishedDiagnostics)
-                {
-                    if (!_projectManager.IsDocumentOpen(key))
-                    {
-                        // If there were previously published diagnostics for this document, take note so
-                        // we can publish an empty set of diagnostics.
-                        documentsToRemove.Add((key, publishEmptyDiagnostics: value.Count > 0));
-                    }
-                }
+            foreach (var (filePath, publishToClear) in documentsToRemove)
+            {
+                _publishedDiagnostics.Remove(filePath);
 
-                if (documentsToRemove.Count == 0)
+                if (publishToClear)
                 {
-                    return;
+                    PublishDiagnosticsForFilePath(filePath, []);
                 }
+            }
 
-                foreach (var (key, publishEmptyDiagnostics) in documentsToRemove)
-                {
-                    publishedDiagnostics.Remove(key);
-
-                    if (publishEmptyDiagnostics)
-                    {
-                        PublishDiagnosticsForFilePath(key, []);
-                    }
-                }
+            if (_publishedDiagnostics.Count > 0)
+            {
+                // There's no way for us to know when a document is closed at this layer. Therefore, we need to poll every X seconds
+                // and check if the currently tracked documents are closed. In practice this work is super minimal.
+                StartDelayToClearDocuments();
             }
         }
     }
 
-    private void PublishDiagnosticsForFilePath(string filePath, IEnumerable<Diagnostic> diagnostics)
+    private void PublishDiagnosticsForFilePath(string filePath, Diagnostic[] diagnostics)
     {
         var uriBuilder = new UriBuilder()
         {
@@ -279,7 +283,7 @@ private void PublishDiagnosticsForFilePath(string filePath, IEnumerable<Diagnost
                 new PublishDiagnosticParams()
                 {
                     Uri = uriBuilder.Uri,
-                    Diagnostics = diagnostics.ToArray(),
+                    Diagnostics = diagnostics,
                 },
                 _disposeTokenSource.Token)
             .Forget();

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/Diagnostics/RazorDiagnosticsPublisherTest.cs

@@ -31,12 +31,12 @@ namespace Microsoft.AspNetCore.Razor.LanguageServer.Diagnostics;
 
 public class RazorDiagnosticsPublisherTest(ITestOutputHelper testOutput) : LanguageServerTestBase(testOutput)
 {
-    private static readonly ImmutableArray<RazorDiagnostic> s_singleRazorDiagnostic =
+    private static readonly RazorDiagnostic[] s_singleRazorDiagnostic =
     [
         RazorDiagnosticFactory.CreateDirective_BlockDirectiveCannotBeImported("test")
     ];
 
-    private static readonly ImmutableArray<Diagnostic> s_singleCSharpDiagnostic =
+    private static readonly Diagnostic[] s_singleCSharpDiagnostic =
     [
         new Diagnostic()
         {
@@ -288,7 +288,7 @@ public async Task PublishDiagnosticsAsync_NewRazorDiagnosticsGetPublished()
 
         using var publisher = new TestRazorDiagnosticsPublisher(_projectManager, clientConnectionMock.Object, TestLanguageServerFeatureOptions.Instance, translateDiagnosticsService, documentContextFactory, LoggerFactory);
         var publisherAccessor = publisher.GetTestAccessor();
-        publisherAccessor.SetPublishedRazorDiagnostics(processedOpenDocument.FilePath, []);
+        publisherAccessor.SetPublishedDiagnostics(processedOpenDocument.FilePath, razorDiagnostics: [], csharpDiagnostics: null);
 
         // Act
         await publisherAccessor.PublishDiagnosticsAsync(processedOpenDocument, DisposalToken);
@@ -384,7 +384,7 @@ public async Task PublishDiagnosticsAsync_NoopsIfRazorDiagnosticsAreSameAsPrevio
 
         using var publisher = new TestRazorDiagnosticsPublisher(_projectManager, clientConnectionMock.Object, TestLanguageServerFeatureOptions.Instance, translateDiagnosticsService, documentContextFactory, LoggerFactory);
         var publisherAccessor = publisher.GetTestAccessor();
-        publisherAccessor.SetPublishedRazorDiagnostics(processedOpenDocument.FilePath, s_singleRazorDiagnostic);
+        publisherAccessor.SetPublishedDiagnostics(processedOpenDocument.FilePath, s_singleRazorDiagnostic, csharpDiagnostics: null);
 
         // Act & Assert
         await publisherAccessor.PublishDiagnosticsAsync(processedOpenDocument, DisposalToken);
@@ -467,8 +467,7 @@ public void ClearClosedDocuments_ClearsDiagnosticsForClosedDocument()
         using var publisher = new TestRazorDiagnosticsPublisher(_projectManager, clientConnectionMock.Object, TestLanguageServerFeatureOptions.Instance, translateDiagnosticsService, documentContextFactory, LoggerFactory);
         var publisherAccessor = publisher.GetTestAccessor();
         Assert.NotNull(_closedDocument.FilePath);
-        publisherAccessor.SetPublishedRazorDiagnostics(_closedDocument.FilePath, s_singleRazorDiagnostic);
-        publisherAccessor.SetPublishedCSharpDiagnostics(_closedDocument.FilePath, s_singleCSharpDiagnostic);
+        publisherAccessor.SetPublishedDiagnostics(_closedDocument.FilePath, s_singleRazorDiagnostic, s_singleCSharpDiagnostic);
 
         // Act
         publisherAccessor.ClearClosedDocuments();
@@ -488,8 +487,7 @@ public void ClearClosedDocuments_NoopsIfDocumentIsStillOpen()
         using var publisher = new TestRazorDiagnosticsPublisher(_projectManager, clientConnectionMock.Object, TestLanguageServerFeatureOptions.Instance, translateDiagnosticsService, documentContextFactory, LoggerFactory);
         var publisherAccessor = publisher.GetTestAccessor();
         Assert.NotNull(_openedDocument.FilePath);
-        publisherAccessor.SetPublishedRazorDiagnostics(_openedDocument.FilePath, s_singleRazorDiagnostic);
-        publisherAccessor.SetPublishedCSharpDiagnostics(_openedDocument.FilePath, s_singleCSharpDiagnostic);
+        publisherAccessor.SetPublishedDiagnostics(_openedDocument.FilePath, s_singleRazorDiagnostic, s_singleCSharpDiagnostic);
 
         // Act & Assert
         publisherAccessor.ClearClosedDocuments();
@@ -506,8 +504,7 @@ public void ClearClosedDocuments_NoopsIfDocumentIsClosedButNoDiagnostics()
         using var publisher = new TestRazorDiagnosticsPublisher(_projectManager, clientConnectionMock.Object, TestLanguageServerFeatureOptions.Instance, translateDiagnosticsService, documentContextFactory, LoggerFactory);
         var publisherAccessor = publisher.GetTestAccessor();
         Assert.NotNull(_closedDocument.FilePath);
-        publisherAccessor.SetPublishedRazorDiagnostics(_closedDocument.FilePath, []);
-        publisherAccessor.SetPublishedCSharpDiagnostics(_closedDocument.FilePath, []);
+        publisherAccessor.SetPublishedDiagnostics(_closedDocument.FilePath, razorDiagnostics: [], csharpDiagnostics: []);
 
         // Act & Assert
         publisherAccessor.ClearClosedDocuments();
@@ -525,10 +522,8 @@ public void ClearClosedDocuments_RestartsTimerIfDocumentsStillOpen()
         var publisherAccessor = publisher.GetTestAccessor();
         Assert.NotNull(_closedDocument.FilePath);
         Assert.NotNull(_openedDocument.FilePath);
-        publisherAccessor.SetPublishedRazorDiagnostics(_closedDocument.FilePath, []);
-        publisherAccessor.SetPublishedCSharpDiagnostics(_closedDocument.FilePath, []);
-        publisherAccessor.SetPublishedRazorDiagnostics(_openedDocument.FilePath, []);
-        publisherAccessor.SetPublishedCSharpDiagnostics(_openedDocument.FilePath, []);
+        publisherAccessor.SetPublishedDiagnostics(_closedDocument.FilePath, razorDiagnostics: [], csharpDiagnostics: []);
+        publisherAccessor.SetPublishedDiagnostics(_openedDocument.FilePath, razorDiagnostics: [], csharpDiagnostics: []);
 
         // Act
         publisherAccessor.ClearClosedDocuments();
@@ -537,7 +532,7 @@ public void ClearClosedDocuments_RestartsTimerIfDocumentsStillOpen()
         Assert.True(publisherAccessor.IsWaitingToClearClosedDocuments);
     }
 
-    private static RazorCodeDocument CreateCodeDocument(ImmutableArray<RazorDiagnostic> diagnostics)
+    private static RazorCodeDocument CreateCodeDocument(RazorDiagnostic[] diagnostics)
     {
         var codeDocument = TestRazorCodeDocument.Create("hello");
         var razorCSharpDocument = RazorCSharpDocument.Create(codeDocument, "hello", RazorCodeGenerationOptions.CreateDefault(), diagnostics);