From 484bf6e8b14eeeaab6b771b87d36dd369b68b634 Mon Sep 17 00:00:00 2001
From: David Wengier <david.wengier@microsoft.com>
Date: Fri, 12 Jan 2024 11:05:13 +1100
Subject: [PATCH] PR Feedback

---
 .../Cohost/CohostProjectSnapshot.cs                       | 8 ++++++--
 .../Cohost/DocumentSnapshotFactory.cs                     | 4 ++--
 .../Cohost/ProjectSnapshotFactory.cs                      | 4 ++--
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Cohost/CohostProjectSnapshot.cs b/src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Cohost/CohostProjectSnapshot.cs
index c7e1abafee..d7cbd77f73 100644
--- a/src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Cohost/CohostProjectSnapshot.cs
+++ b/src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Cohost/CohostProjectSnapshot.cs
@@ -33,6 +33,7 @@ internal class CohostProjectSnapshot : IProjectSnapshot
     private readonly Lazy<RazorConfiguration> _lazyConfiguration;
     private readonly Lazy<RazorProjectEngine> _lazyProjectEngine;
     private readonly AsyncLazy<ImmutableArray<TagHelperDescriptor>> _tagHelpersLazy;
+    private readonly Lazy<ProjectWorkspaceState> _projectWorkspaceStateLazy;
     private readonly Lazy<ImmutableDictionary<string, ImmutableArray<string>>> _importsToRelatedDocumentsLazy;
 
     public CohostProjectSnapshot(Project project, DocumentSnapshotFactory documentSnapshotFactory, ITelemetryReporter telemetryReporter, JoinableTaskContext joinableTaskContext)
@@ -40,7 +41,7 @@ internal class CohostProjectSnapshot : IProjectSnapshot
         _project = project;
         _documentSnapshotFactory = documentSnapshotFactory;
         _telemetryReporter = telemetryReporter;
-        _projectKey = ProjectKey.From(_project)!.Value;
+        _projectKey = ProjectKey.From(_project).AssumeNotNull();
 
         _lazyConfiguration = new Lazy<RazorConfiguration>(CreateRazorConfiguration);
         _lazyProjectEngine = new Lazy<RazorProjectEngine>(() => DefaultProjectEngineFactory.Create(
@@ -61,6 +62,8 @@ internal class CohostProjectSnapshot : IProjectSnapshot
             return resolver.GetTagHelpersAsync(_project, GetProjectEngine(), CancellationToken.None).AsTask();
         }, joinableTaskContext.Factory);
 
+        _projectWorkspaceStateLazy = new Lazy<ProjectWorkspaceState>(() => ProjectWorkspaceState.Create(TagHelpers, CSharpLanguageVersion));
+
         _importsToRelatedDocumentsLazy = new Lazy<ImmutableDictionary<string, ImmutableArray<string>>>(() =>
         {
             var importsToRelatedDocuments = ImmutableDictionary.Create<string, ImmutableArray<string>>(FilePathNormalizer.Comparer);
@@ -97,7 +100,7 @@ internal class CohostProjectSnapshot : IProjectSnapshot
 
     public ImmutableArray<TagHelperDescriptor> TagHelpers => _tagHelpersLazy.GetValue();
 
-    public ProjectWorkspaceState ProjectWorkspaceState => ProjectWorkspaceState.Create(TagHelpers, CSharpLanguageVersion);
+    public ProjectWorkspaceState ProjectWorkspaceState => _projectWorkspaceStateLazy.Value;
 
     public IDocumentSnapshot? GetDocument(string filePath)
     {
@@ -122,6 +125,7 @@ internal class CohostProjectSnapshot : IProjectSnapshot
         }
 
         using var _ = ArrayBuilderPool<IDocumentSnapshot>.GetPooledObject(out var builder);
+        builder.SetCapacityIfLarger(relatedDocuments.Length);
 
         foreach (var relatedDocumentFilePath in relatedDocuments)
         {
diff --git a/src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Cohost/DocumentSnapshotFactory.cs b/src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Cohost/DocumentSnapshotFactory.cs
index 98f9e33d55..0d7c4cb768 100644
--- a/src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Cohost/DocumentSnapshotFactory.cs
+++ b/src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Cohost/DocumentSnapshotFactory.cs
@@ -13,9 +13,9 @@ namespace Microsoft.AspNetCore.Razor.LanguageServer.Cohost;
 [method: ImportingConstructor]
 internal class DocumentSnapshotFactory(Lazy<ProjectSnapshotFactory> projectSnapshotFactory)
 {
-    private readonly Lazy<ProjectSnapshotFactory> _projectSnapshotFactory = projectSnapshotFactory;
+    private static readonly ConditionalWeakTable<TextDocument, IDocumentSnapshot> _documentSnapshots = new();
 
-    private readonly ConditionalWeakTable<TextDocument, IDocumentSnapshot> _documentSnapshots = new();
+    private readonly Lazy<ProjectSnapshotFactory> _projectSnapshotFactory = projectSnapshotFactory;
 
     public IDocumentSnapshot GetOrCreate(TextDocument textDocument)
     {
diff --git a/src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Cohost/ProjectSnapshotFactory.cs b/src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Cohost/ProjectSnapshotFactory.cs
index ccc25eb589..dc877d7a38 100644
--- a/src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Cohost/ProjectSnapshotFactory.cs
+++ b/src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Cohost/ProjectSnapshotFactory.cs
@@ -14,12 +14,12 @@ namespace Microsoft.AspNetCore.Razor.LanguageServer.Cohost;
 [method: ImportingConstructor]
 internal class ProjectSnapshotFactory(DocumentSnapshotFactory documentSnapshotFactory, ITelemetryReporter telemetryReporter, JoinableTaskContext joinableTaskContext)
 {
+    private static readonly ConditionalWeakTable<Project, IProjectSnapshot> _projectSnapshots = new();
+
     private readonly DocumentSnapshotFactory _documentSnapshotFactory = documentSnapshotFactory;
     private readonly ITelemetryReporter _telemetryReporter = telemetryReporter;
     private readonly JoinableTaskContext _joinableTaskContext = joinableTaskContext;
 
-    private readonly ConditionalWeakTable<Project, IProjectSnapshot> _projectSnapshots = new();
-
     public IProjectSnapshot GetOrCreate(Project project)
     {
         if (!_projectSnapshots.TryGetValue(project, out var projectSnapshot))