Merge pull request #73000 from CyrusNajmabadi/fixedArray

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

@@ -533,12 +533,11 @@ public static ChecksumCollection GetOrCreateChecksumCollection<TReference>(
             references,
             static (references, tuple) =>
             {
-                var checksums = new Checksum[references.Count];
-                var index = 0;
+                var checksums = new FixedSizeArrayBuilder<Checksum>(references.Count);
                 foreach (var reference in references)
-                    checksums[index++] = tuple.serializer.CreateChecksum(reference, tuple.cancellationToken);
+                    checksums.Add(tuple.serializer.CreateChecksum(reference, tuple.cancellationToken));
 
-                return new ChecksumCollection(ImmutableCollectionsMarshal.AsImmutableArray(checksums));
+                return new ChecksumCollection(checksums.MoveToImmutable());
             },
             (serializer, cancellationToken));
     }

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

@@ -135,12 +135,11 @@ public ImmutableArray<TValue> SelectAsArray<TValue>(Func<TState, TValue> selecto
 
     public ImmutableArray<TValue> SelectAsArray<TValue, TArg>(Func<TState, TArg, TValue> selector, TArg arg)
     {
-        var result = new TValue[_map.Count];
-        var index = 0;
+        var result = new FixedSizeArrayBuilder<TValue>(_map.Count);
         foreach (var (_, state) in _map)
-            result[index++] = selector(state, arg);
+            result.Add(selector(state, arg));
 
-        return ImmutableCollectionsMarshal.AsImmutableArray(result);
+        return result.MoveToImmutable();
     }
 
     public TextDocumentStates<TState> AddRange(ImmutableArray<TState> states)
@@ -290,24 +289,22 @@ public int Compare(DocumentId? x, DocumentId? y)
 
     public async ValueTask<DocumentChecksumsAndIds> GetDocumentChecksumsAndIdsAsync(CancellationToken cancellationToken)
     {
-        var attributeChecksums = new Checksum[_map.Count];
-        var textChecksums = new Checksum[_map.Count];
-        var documentIds = new DocumentId[_map.Count];
+        var attributeChecksums = new FixedSizeArrayBuilder<Checksum>(_map.Count);
+        var textChecksums = new FixedSizeArrayBuilder<Checksum>(_map.Count);
+        var documentIds = new FixedSizeArrayBuilder<DocumentId>(_map.Count);
 
-        var index = 0;
         foreach (var (documentId, state) in _map)
         {
             var stateChecksums = await state.GetStateChecksumsAsync(cancellationToken).ConfigureAwait(false);
-            attributeChecksums[index] = stateChecksums.Info;
-            textChecksums[index] = stateChecksums.Text;
-            documentIds[index] = documentId;
-            index++;
+            attributeChecksums.Add(stateChecksums.Info);
+            textChecksums.Add(stateChecksums.Text);
+            documentIds.Add(documentId);
         }
 
         return new(
-            new ChecksumCollection(ImmutableCollectionsMarshal.AsImmutableArray(attributeChecksums)),
-            new ChecksumCollection(ImmutableCollectionsMarshal.AsImmutableArray(textChecksums)),
-            ImmutableCollectionsMarshal.AsImmutableArray(documentIds));
+            new ChecksumCollection(attributeChecksums.MoveToImmutable()),
+            new ChecksumCollection(textChecksums.MoveToImmutable()),
+            documentIds.MoveToImmutable());
     }
 
     public void AddDocumentIdsWithFilePath(ref TemporaryArray<DocumentId> temporaryArray, string filePath)

src/Workspaces/Remote/Core/AbstractAssetProvider.cs

@@ -95,16 +95,15 @@ await analyzerConfigDocumentInfosTask.ConfigureAwait(false),
 
         async Task<ImmutableArray<DocumentInfo>> CreateDocumentInfosAsync(DocumentChecksumsAndIds checksumsAndIds)
         {
-            var documentInfos = new DocumentInfo[checksumsAndIds.Length];
+            var documentInfos = new FixedSizeArrayBuilder<DocumentInfo>(checksumsAndIds.Length);
 
-            var index = 0;
             foreach (var (attributeChecksum, textChecksum, documentId) in checksumsAndIds)
             {
                 cancellationToken.ThrowIfCancellationRequested();
-                documentInfos[index++] = await CreateDocumentInfoAsync(documentId, attributeChecksum, textChecksum, cancellationToken).ConfigureAwait(false);
+                documentInfos.Add(await CreateDocumentInfoAsync(documentId, attributeChecksum, textChecksum, cancellationToken).ConfigureAwait(false));
             }
 
-            return ImmutableCollectionsMarshal.AsImmutableArray(documentInfos);
+            return documentInfos.MoveToImmutable();
         }
     }
 
@@ -195,16 +194,19 @@ public static async Task<ImmutableArray<T>> GetAssetsArrayAsync<T>(
         this AbstractAssetProvider assetProvider, AssetPath assetPath, ChecksumCollection checksums, CancellationToken cancellationToken) where T : class
     {
         using var _1 = PooledHashSet<Checksum>.GetInstance(out var checksumSet);
+#if NET
+        checksumSet.EnsureCapacity(checksums.Children.Length);
+#endif
         checksumSet.AddAll(checksums.Children);
 
-        var builder = ImmutableArray.CreateBuilder<T>(checksumSet.Count);
+        using var _ = ArrayBuilder<T>.GetInstance(checksumSet.Count, out var builder);
 
         await assetProvider.GetAssetHelper<T>().GetAssetsAsync(
             assetPath, checksumSet,
             static (checksum, asset, builder) => builder.Add(asset),
             builder,
             cancellationToken).ConfigureAwait(false);
 
-        return builder.MoveToImmutable();
+        return builder.ToImmutableAndClear();
     }
 }

src/Workspaces/Remote/ServiceHub/Services/SourceGeneration/RemoteSourceGenerationService.cs

@@ -38,15 +38,14 @@ protected override IRemoteSourceGenerationService CreateService(in ServiceConstr
             var project = solution.GetRequiredProject(projectId);
             var documentStates = await solution.CompilationState.GetSourceGeneratedDocumentStatesAsync(project.State, cancellationToken).ConfigureAwait(false);
 
-            var result = new (SourceGeneratedDocumentIdentity documentIdentity, SourceGeneratedDocumentContentIdentity contentIdentity, DateTime generationDateTime)[documentStates.Ids.Count];
-            var index = 0;
+            var result = new FixedSizeArrayBuilder<(SourceGeneratedDocumentIdentity documentIdentity, SourceGeneratedDocumentContentIdentity contentIdentity, DateTime generationDateTime)>(documentStates.Ids.Count);
             foreach (var (id, state) in documentStates.States)
             {
                 Contract.ThrowIfFalse(id.IsSourceGenerated);
-                result[index++] = (state.Identity, state.GetContentIdentity(), state.GenerationDateTime);
+                result.Add((state.Identity, state.GetContentIdentity(), state.GenerationDateTime));
             }
 
-            return ImmutableCollectionsMarshal.AsImmutableArray(result);
+            return result.MoveToImmutable();
         }, cancellationToken);
     }
 
@@ -58,17 +57,16 @@ public ValueTask<ImmutableArray<string>> GetContentsAsync(
             var project = solution.GetRequiredProject(projectId);
             var documentStates = await solution.CompilationState.GetSourceGeneratedDocumentStatesAsync(project.State, cancellationToken).ConfigureAwait(false);
 
-            var result = new string[documentIds.Length];
-            var index = 0;
+            var result = new FixedSizeArrayBuilder<string>(documentIds.Length);
             foreach (var id in documentIds)
             {
                 Contract.ThrowIfFalse(id.IsSourceGenerated);
                 var state = documentStates.GetRequiredState(id);
                 var text = await state.GetTextAsync(cancellationToken).ConfigureAwait(false);
-                result[index++] = text.ToString();
+                result.Add(text.ToString());
             }
 
-            return ImmutableCollectionsMarshal.AsImmutableArray(result);
+            return result.MoveToImmutable();
         }, cancellationToken);
     }
 

src/Workspaces/SharedUtilitiesAndExtensions/Compiler/Core/CompilerExtensions.projitems

@@ -497,6 +497,7 @@
     <Compile Include="$(MSBuildThisFileDirectory)Utilities\ComparerWithState.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Utilities\Contract.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Utilities\Contract.InterpolatedStringHandlers.cs" />
+    <Compile Include="$(MSBuildThisFileDirectory)Utilities\FixedSizeArrayBuilder.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Utilities\ISpeculationAnalyzer.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Utilities\OptionalExtensions.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Utilities\PathMetadataUtilities.cs" />

src/Workspaces/SharedUtilitiesAndExtensions/Compiler/Core/Utilities/FixedSizeArrayBuilder.cs

@@ -0,0 +1,42 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Collections.Immutable;
+using System.Runtime.InteropServices;
+using Microsoft.CodeAnalysis.PooledObjects;
+using Roslyn.Utilities;
+
+/// <summary>
+/// A bare-bones, pooled builder, focused on the case of producing <see cref="ImmutableArray{T}"/>s where the final
+/// array size is known at construction time.  In the golden path, where all the expected items are added to the
+/// builder, and <see cref="MoveToImmutable"/> is called, this type is entirely garbage free.  In the non-golden path
+/// (usually encountered when a cancellation token interrupts getting the final array), this will leak the intermediary
+/// array created to store the results.
+/// </summary>
+[NonCopyable]
+internal struct FixedSizeArrayBuilder<T>(int capacity)
+{
+    private T[] _values = new T[capacity];
+    private int _index;
+
+    public void Add(T value)
+        => _values[_index++] = value;
+
+    /// <summary>
+    /// Moves the underlying buffer out of control of this type, into the returned <see cref="ImmutableArray{T}"/>. It
+    /// is an error for a client of this type to specify a capacity and then attempt to call <see
+    /// cref="MoveToImmutable"/> without that number of elements actually having been added to the builder.  This will
+    /// throw if attempted.  This <see cref="FixedSizeArrayBuilder{T}"/> is effectively unusable once this is called.
+    /// The internal buffer will reset to an empty array, meaning no more items could ever be added to it.
+    /// </summary>
+    public ImmutableArray<T> MoveToImmutable()
+    {
+        Contract.ThrowIfTrue(_index != _values.Length);
+        var result = ImmutableCollectionsMarshal.AsImmutableArray(_values);
+        _values = Array.Empty<T>();
+        _index = 0;
+        return result;
+    }
+}