Merge pull request #22494 from jasonmalinowski/reduce-documentid-array-allocations

Reduce DocumentId array allocations

src/Workspaces/Core/Portable/Workspace/Solution/ProjectState.cs

@@ -20,10 +20,21 @@ internal partial class ProjectState
         private readonly ProjectInfo _projectInfo;
         private readonly HostLanguageServices _languageServices;
         private readonly SolutionServices _solutionServices;
-        private readonly ImmutableDictionary<DocumentId, DocumentState> _documentStates;
-        private readonly ImmutableDictionary<DocumentId, TextDocumentState> _additionalDocumentStates;
-        private readonly IReadOnlyList<DocumentId> _documentIds;
-        private readonly IReadOnlyList<DocumentId> _additionalDocumentIds;
+
+        /// <summary>
+        /// The documents in this project. They are sorted by <see cref="DocumentId.Id"/> to provide a stable sort for
+        /// <see cref="GetChecksumAsync(CancellationToken)"/>.
+        /// </summary>
+        private readonly ImmutableSortedDictionary<DocumentId, DocumentState> _documentStates;
+
+        /// <summary>
+        /// The additional documents in this project. They are sorted by <see cref="DocumentId.Id"/> to provide a stable sort for
+        /// <see cref="GetChecksumAsync(CancellationToken)"/>.
+        /// </summary>
+        private readonly ImmutableSortedDictionary<DocumentId, TextDocumentState> _additionalDocumentStates;
+
+        private readonly ImmutableList<DocumentId> _documentIds;
+        private readonly ImmutableList<DocumentId> _additionalDocumentIds;
         private readonly AsyncLazy<VersionStamp> _lazyLatestDocumentVersion;
         private readonly AsyncLazy<VersionStamp> _lazyLatestDocumentTopLevelChangeVersion;
 
@@ -37,17 +48,17 @@ internal partial class ProjectState
             ProjectInfo projectInfo,
             HostLanguageServices languageServices,
             SolutionServices solutionServices,
-            IEnumerable<DocumentId> documentIds,
-            IEnumerable<DocumentId> additionalDocumentIds,
-            ImmutableDictionary<DocumentId, DocumentState> documentStates,
-            ImmutableDictionary<DocumentId, TextDocumentState> additionalDocumentStates,
+            ImmutableList<DocumentId> documentIds,
+            ImmutableList<DocumentId> additionalDocumentIds,
+            ImmutableSortedDictionary<DocumentId, DocumentState> documentStates,
+            ImmutableSortedDictionary<DocumentId, TextDocumentState> additionalDocumentStates,
             AsyncLazy<VersionStamp> lazyLatestDocumentVersion,
             AsyncLazy<VersionStamp> lazyLatestDocumentTopLevelChangeVersion)
         {
             _solutionServices = solutionServices;
             _languageServices = languageServices;
-            _documentIds = documentIds.ToImmutableReadOnlyListOrEmpty();
-            _additionalDocumentIds = additionalDocumentIds.ToImmutableReadOnlyListOrEmpty();
+            _documentIds = documentIds;
+            _additionalDocumentIds = additionalDocumentIds;
             _documentStates = documentStates;
             _additionalDocumentStates = additionalDocumentStates;
             _lazyLatestDocumentVersion = lazyLatestDocumentVersion;
@@ -72,18 +83,18 @@ public ProjectState(ProjectInfo projectInfo, HostLanguageServices languageServic
 
             var projectInfoFixed = FixProjectInfo(projectInfo);
 
-            _documentIds = projectInfoFixed.Documents.Select(d => d.Id).ToImmutableArray();
-            _additionalDocumentIds = projectInfoFixed.AdditionalDocuments.Select(d => d.Id).ToImmutableArray();
+            _documentIds = projectInfoFixed.Documents.Select(d => d.Id).ToImmutableList();
+            _additionalDocumentIds = projectInfoFixed.AdditionalDocuments.Select(d => d.Id).ToImmutableList();
 
             var parseOptions = projectInfoFixed.ParseOptions;
-            var docStates = ImmutableDictionary.CreateRange<DocumentId, DocumentState>(
+            var docStates = ImmutableSortedDictionary.CreateRange(DocumentIdComparer.Instance,
                 projectInfoFixed.Documents.Select(d =>
                     new KeyValuePair<DocumentId, DocumentState>(d.Id,
                         CreateDocument(d, parseOptions, languageServices, solutionServices))));
 
             _documentStates = docStates;
 
-            var additionalDocStates = ImmutableDictionary.CreateRange<DocumentId, TextDocumentState>(
+            var additionalDocStates = ImmutableSortedDictionary.CreateRange(DocumentIdComparer.Instance,
                     projectInfoFixed.AdditionalDocuments.Select(d =>
                         new KeyValuePair<DocumentId, TextDocumentState>(d.Id, TextDocumentState.Create(d, solutionServices))));
 
@@ -129,7 +140,7 @@ private ProjectInfo FixProjectInfo(ProjectInfo projectInfo)
             return projectInfo;
         }
 
-        private static async Task<VersionStamp> ComputeLatestDocumentVersionAsync(ImmutableDictionary<DocumentId, DocumentState> documentStates, ImmutableDictionary<DocumentId, TextDocumentState> additionalDocumentStates, CancellationToken cancellationToken)
+        private static async Task<VersionStamp> ComputeLatestDocumentVersionAsync(IImmutableDictionary<DocumentId, DocumentState> documentStates, IImmutableDictionary<DocumentId, TextDocumentState> additionalDocumentStates, CancellationToken cancellationToken)
         {
             // this may produce a version that is out of sync with the actual Document versions.
             var latestVersion = VersionStamp.Default;
@@ -157,8 +168,8 @@ private static async Task<VersionStamp> ComputeLatestDocumentVersionAsync(Immuta
 
         private AsyncLazy<VersionStamp> CreateLazyLatestDocumentTopLevelChangeVersion(
             TextDocumentState newDocument,
-            ImmutableDictionary<DocumentId, DocumentState> newDocumentStates,
-            ImmutableDictionary<DocumentId, TextDocumentState> newAdditionalDocumentStates)
+            IImmutableDictionary<DocumentId, DocumentState> newDocumentStates,
+            IImmutableDictionary<DocumentId, TextDocumentState> newAdditionalDocumentStates)
         {
             if (_lazyLatestDocumentTopLevelChangeVersion.TryGetValue(out var oldVersion))
             {
@@ -176,7 +187,7 @@ private static async Task<VersionStamp> ComputeTopLevelChangeTextVersionAsync(Ve
             return newVersion.GetNewerVersion(oldVersion);
         }
 
-        private static async Task<VersionStamp> ComputeLatestDocumentTopLevelChangeVersionAsync(ImmutableDictionary<DocumentId, DocumentState> documentStates, ImmutableDictionary<DocumentId, TextDocumentState> additionalDocumentStates, CancellationToken cancellationToken)
+        private static async Task<VersionStamp> ComputeLatestDocumentTopLevelChangeVersionAsync(IImmutableDictionary<DocumentId, DocumentState> documentStates, IImmutableDictionary<DocumentId, TextDocumentState> additionalDocumentStates, CancellationToken cancellationToken)
         {
             // this may produce a version that is out of sync with the actual Document versions.
             var latestVersion = VersionStamp.Default;
@@ -312,10 +323,10 @@ public async Task<VersionStamp> GetSemanticVersionAsync(CancellationToken cancel
         public IReadOnlyList<DocumentId> AdditionalDocumentIds => _additionalDocumentIds;
 
         [DebuggerBrowsable(DebuggerBrowsableState.Collapsed)]
-        public ImmutableDictionary<DocumentId, DocumentState> DocumentStates => _documentStates;
+        public IImmutableDictionary<DocumentId, DocumentState> DocumentStates => _documentStates;
 
         [DebuggerBrowsable(DebuggerBrowsableState.Collapsed)]
-        public ImmutableDictionary<DocumentId, TextDocumentState> AdditionalDocumentStates => _additionalDocumentStates;
+        public IImmutableDictionary<DocumentId, TextDocumentState> AdditionalDocumentStates => _additionalDocumentStates;
 
         public bool ContainsDocument(DocumentId documentId)
         {
@@ -341,10 +352,10 @@ public TextDocumentState GetAdditionalDocumentState(DocumentId documentId)
 
         private ProjectState With(
             ProjectInfo projectInfo = null,
-            ImmutableArray<DocumentId> documentIds = default,
-            ImmutableArray<DocumentId> additionalDocumentIds = default,
-            ImmutableDictionary<DocumentId, DocumentState> documentStates = null,
-            ImmutableDictionary<DocumentId, TextDocumentState> additionalDocumentStates = null,
+            ImmutableList<DocumentId> documentIds = default,
+            ImmutableList<DocumentId> additionalDocumentIds = default,
+            ImmutableSortedDictionary<DocumentId, DocumentState> documentStates = null,
+            ImmutableSortedDictionary<DocumentId, TextDocumentState> additionalDocumentStates = null,
             AsyncLazy<VersionStamp> latestDocumentVersion = null,
             AsyncLazy<VersionStamp> latestDocumentTopLevelChangeVersion = null)
         {
@@ -352,8 +363,8 @@ public TextDocumentState GetAdditionalDocumentState(DocumentId documentId)
                 projectInfo ?? _projectInfo,
                 _languageServices,
                 _solutionServices,
-                documentIds.IsDefault ? _documentIds : documentIds,
-                additionalDocumentIds.IsDefault ? _additionalDocumentIds : additionalDocumentIds,
+                documentIds ?? _documentIds,
+                additionalDocumentIds ?? _additionalDocumentIds,
                 documentStates ?? _documentStates,
                 additionalDocumentStates ?? _additionalDocumentStates,
                 latestDocumentVersion ?? _lazyLatestDocumentVersion,
@@ -558,8 +569,8 @@ public ProjectState AddDocument(DocumentState document)
 
             return this.With(
                 projectInfo: this.ProjectInfo.WithVersion(this.Version.GetNewerVersion()),
-                documentIds: this.DocumentIds.ToImmutableArray().Add(document.Id),
-                documentStates: this.DocumentStates.Add(document.Id, document));
+                documentIds: _documentIds.Add(document.Id),
+                documentStates: _documentStates.Add(document.Id, document));
         }
 
         public ProjectState AddAdditionalDocument(TextDocumentState document)
@@ -568,8 +579,8 @@ public ProjectState AddAdditionalDocument(TextDocumentState document)
 
             return this.With(
                 projectInfo: this.ProjectInfo.WithVersion(this.Version.GetNewerVersion()),
-                additionalDocumentIds: this.AdditionalDocumentIds.ToImmutableArray().Add(document.Id),
-                additionalDocumentStates: this.AdditionalDocumentStates.Add(document.Id, document));
+                additionalDocumentIds: _additionalDocumentIds.Add(document.Id),
+                additionalDocumentStates: _additionalDocumentStates.Add(document.Id, document));
         }
 
         public ProjectState RemoveDocument(DocumentId documentId)
@@ -578,8 +589,8 @@ public ProjectState RemoveDocument(DocumentId documentId)
 
             return this.With(
                 projectInfo: this.ProjectInfo.WithVersion(this.Version.GetNewerVersion()),
-                documentIds: this.DocumentIds.ToImmutableArray().Remove(documentId),
-                documentStates: this.DocumentStates.Remove(documentId));
+                documentIds: _documentIds.Remove(documentId),
+                documentStates: _documentStates.Remove(documentId));
         }
 
         public ProjectState RemoveAdditionalDocument(DocumentId documentId)
@@ -588,16 +599,16 @@ public ProjectState RemoveAdditionalDocument(DocumentId documentId)
 
             return this.With(
                 projectInfo: this.ProjectInfo.WithVersion(this.Version.GetNewerVersion()),
-                additionalDocumentIds: this.AdditionalDocumentIds.ToImmutableArray().Remove(documentId),
-                additionalDocumentStates: this.AdditionalDocumentStates.Remove(documentId));
+                additionalDocumentIds: _additionalDocumentIds.Remove(documentId),
+                additionalDocumentStates: _additionalDocumentStates.Remove(documentId));
         }
 
         public ProjectState RemoveAllDocuments()
         {
             return this.With(
                 projectInfo: this.ProjectInfo.WithVersion(this.Version.GetNewerVersion()).WithDocuments(SpecializedCollections.EmptyEnumerable<DocumentInfo>()),
-                documentIds: ImmutableArray.Create<DocumentId>(),
-                documentStates: ImmutableDictionary<DocumentId, DocumentState>.Empty);
+                documentIds: ImmutableList<DocumentId>.Empty,
+                documentStates: ImmutableSortedDictionary.Create<DocumentId, DocumentState>(DocumentIdComparer.Instance));
         }
 
         public ProjectState UpdateDocument(DocumentState newDocument, bool textChanged, bool recalculateDependentVersions)
@@ -610,7 +621,7 @@ public ProjectState UpdateDocument(DocumentState newDocument, bool textChanged,
                 return this;
             }
 
-            var newDocumentStates = this.DocumentStates.SetItem(newDocument.Id, newDocument);
+            var newDocumentStates = _documentStates.SetItem(newDocument.Id, newDocument);
             GetLatestDependentVersions(
                 newDocumentStates, _additionalDocumentStates, oldDocument, newDocument, recalculateDependentVersions, textChanged,
                 out var dependentDocumentVersion, out var dependentSemanticVersion);
@@ -631,7 +642,7 @@ public ProjectState UpdateAdditionalDocument(TextDocumentState newDocument, bool
                 return this;
             }
 
-            var newDocumentStates = this.AdditionalDocumentStates.SetItem(newDocument.Id, newDocument);
+            var newDocumentStates = _additionalDocumentStates.SetItem(newDocument.Id, newDocument);
             GetLatestDependentVersions(
                 _documentStates, newDocumentStates, oldDocument, newDocument, recalculateDependentVersions, textChanged,
                 out var dependentDocumentVersion, out var dependentSemanticVersion);
@@ -643,8 +654,8 @@ public ProjectState UpdateAdditionalDocument(TextDocumentState newDocument, bool
         }
 
         private void GetLatestDependentVersions(
-            ImmutableDictionary<DocumentId, DocumentState> newDocumentStates,
-            ImmutableDictionary<DocumentId, TextDocumentState> newAdditionalDocumentStates,
+            IImmutableDictionary<DocumentId, DocumentState> newDocumentStates,
+            IImmutableDictionary<DocumentId, TextDocumentState> newAdditionalDocumentStates,
             TextDocumentState oldDocument, TextDocumentState newDocument,
             bool recalculateDependentVersions, bool textChanged,
             out AsyncLazy<VersionStamp> dependentDocumentVersion, out AsyncLazy<VersionStamp> dependentSemanticVersion)
@@ -680,5 +691,19 @@ public ProjectState UpdateAdditionalDocument(TextDocumentState newDocument, bool
                     CreateLazyLatestDocumentTopLevelChangeVersion(newDocument, newDocumentStates, newAdditionalDocumentStates) :
                     _lazyLatestDocumentTopLevelChangeVersion;
         }
+
+        private sealed class DocumentIdComparer : IComparer<DocumentId>
+        {
+            public static IComparer<DocumentId> Instance = new DocumentIdComparer();
+
+            private DocumentIdComparer()
+            {
+            }
+
+            public int Compare(DocumentId x, DocumentId y)
+            {
+                return x.Id.CompareTo(y.Id);
+            }
+        }
     }
 }

src/Workspaces/Core/Portable/Workspace/Solution/ProjectState_Checksum.cs

@@ -31,12 +31,10 @@ private async Task<ProjectStateChecksums> ComputeChecksumsAsync(CancellationToke
         {
             using (Logger.LogBlock(FunctionId.ProjectState_ComputeChecksumsAsync, FilePath, cancellationToken))
             {
-                // get states by id order to have deterministic checksum
-                var orderedDocumentIds = ChecksumCache.GetOrCreate(DocumentIds, _ => DocumentIds.OrderBy(id => id.Id).ToImmutableArray());
-                var documentChecksumsTasks = orderedDocumentIds.Select(id => DocumentStates[id].GetChecksumAsync(cancellationToken));
-
-                var orderedAdditionalDocumentIds = ChecksumCache.GetOrCreate(AdditionalDocumentIds, _ => AdditionalDocumentIds.OrderBy(id => id.Id).ToImmutableArray());
-                var additionalDocumentChecksumTasks = orderedAdditionalDocumentIds.Select(id => AdditionalDocumentStates[id].GetChecksumAsync(cancellationToken));
+                // Here, we use the _documentStates and _additionalDocumentStates and visit them in order; we ensure that those are
+                // sorted by ID so we have a consistent sort.
+                var documentChecksumsTasks = _documentStates.Select(pair => pair.Value.GetChecksumAsync(cancellationToken));
+                var additionalDocumentChecksumTasks = _additionalDocumentStates.Select(pair => pair.Value.GetChecksumAsync(cancellationToken));
 
                 var serializer = new Serializer(_solutionServices.Workspace);
 

src/Workspaces/Core/Portable/Workspace/Solution/StateChecksums.cs

@@ -174,7 +174,7 @@ public ProjectStateChecksums(params object[] children) : base(WellKnownSynchroni
         }
 
         private static void Find<T>(
-            ImmutableDictionary<DocumentId, T> values,
+            IImmutableDictionary<DocumentId, T> values,
             HashSet<Checksum> searchingChecksumsLeft,
             Dictionary<Checksum, object> result,
             CancellationToken cancellationToken) where T : TextDocumentState

src/Workspaces/Remote/Core/Services/SolutionCreator.cs

@@ -413,17 +413,17 @@ private async Task<TextDocument> UpdateDocumentInfoAsync(TextDocument document,
             return GetDocumentMapAsync(project, project.State.DocumentStates, documents);
         }
 
-        private async Task<Dictionary<DocumentId, DocumentStateChecksums>> GetDocumentMapAsync<T>(Project project, ImmutableDictionary<DocumentId, T> states, HashSet<Checksum> documents)
+        private async Task<Dictionary<DocumentId, DocumentStateChecksums>> GetDocumentMapAsync<T>(Project project, IImmutableDictionary<DocumentId, T> states, HashSet<Checksum> documents)
             where T : TextDocumentState
         {
             var map = new Dictionary<DocumentId, DocumentStateChecksums>();
 
             foreach (var kv in states)
             {
-                var doucmentChecksums = await kv.Value.GetStateChecksumsAsync(_cancellationToken).ConfigureAwait(false);
-                if (documents.Contains(doucmentChecksums.Checksum))
+                var documentChecksums = await kv.Value.GetStateChecksumsAsync(_cancellationToken).ConfigureAwait(false);
+                if (documents.Contains(documentChecksums.Checksum))
                 {
-                    map.Add(kv.Key, doucmentChecksums);
+                    map.Add(kv.Key, documentChecksums);
                 }
             }