Make frozen doc state non-nullable (#77896)

Author: CyrusNajmabadi

src/VisualStudio/Core/Test.Next/Remote/SerializationValidator.cs

@@ -239,13 +239,8 @@ private static void AssertDocumentChecksumCollectionEqual(
     internal void SolutionCompilationStateEqual(SolutionCompilationStateChecksums solutionObject1, SolutionCompilationStateChecksums solutionObject2)
     {
         Assert.Equal(solutionObject1.Checksum, solutionObject2.Checksum);
-        Assert.Equal(solutionObject1.FrozenSourceGeneratedDocumentIdentities.HasValue, solutionObject2.FrozenSourceGeneratedDocumentIdentities.HasValue);
-        if (solutionObject1.FrozenSourceGeneratedDocumentIdentities.HasValue)
-            AssertChecksumCollectionEqual(solutionObject1.FrozenSourceGeneratedDocumentIdentities.Value, solutionObject2.FrozenSourceGeneratedDocumentIdentities!.Value);
-
-        Assert.Equal(solutionObject1.FrozenSourceGeneratedDocuments.HasValue, solutionObject2.FrozenSourceGeneratedDocuments.HasValue);
-        if (solutionObject1.FrozenSourceGeneratedDocuments.HasValue)
-            AssertDocumentChecksumCollectionEqual(solutionObject1.FrozenSourceGeneratedDocuments.Value, solutionObject2.FrozenSourceGeneratedDocuments!.Value);
+        AssertChecksumCollectionEqual(solutionObject1.FrozenSourceGeneratedDocumentIdentities, solutionObject2.FrozenSourceGeneratedDocumentIdentities);
+        AssertDocumentChecksumCollectionEqual(solutionObject1.FrozenSourceGeneratedDocuments, solutionObject2.FrozenSourceGeneratedDocuments);
     }
 
     internal void SolutionStateEqual(SolutionStateChecksums solutionObject1, SolutionStateChecksums solutionObject2)

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

@@ -42,7 +42,7 @@ internal sealed partial class SolutionCompilationState
     public SolutionState SolutionState { get; }
 
     public bool PartialSemanticsEnabled { get; }
-    public TextDocumentStates<SourceGeneratedDocumentState>? FrozenSourceGeneratedDocumentStates { get; }
+    public TextDocumentStates<SourceGeneratedDocumentState> FrozenSourceGeneratedDocumentStates { get; }
 
     // Values for all these are created on demand.
     private ImmutableSegmentedDictionary<ProjectId, ICompilationTracker> _projectIdToTrackerMap;
@@ -62,7 +62,7 @@ private SolutionCompilationState(
         bool partialSemanticsEnabled,
         ImmutableSegmentedDictionary<ProjectId, ICompilationTracker> projectIdToTrackerMap,
         SourceGeneratorExecutionVersionMap sourceGeneratorExecutionVersionMap,
-        TextDocumentStates<SourceGeneratedDocumentState>? frozenSourceGeneratedDocumentStates,
+        TextDocumentStates<SourceGeneratedDocumentState> frozenSourceGeneratedDocumentStates,
         AsyncLazy<SolutionCompilationState>? cachedFrozenSnapshot = null)
     {
         SolutionState = solution;
@@ -94,7 +94,7 @@ public SolutionCompilationState(
               partialSemanticsEnabled,
               projectIdToTrackerMap: ImmutableSegmentedDictionary<ProjectId, ICompilationTracker>.Empty,
               sourceGeneratorExecutionVersionMap: SourceGeneratorExecutionVersionMap.Empty,
-              frozenSourceGeneratedDocumentStates: null)
+              frozenSourceGeneratedDocumentStates: TextDocumentStates<SourceGeneratedDocumentState>.Empty)
     {
     }
 
@@ -117,7 +117,7 @@ private SolutionCompilationState Branch(
         SolutionState newSolutionState,
         ImmutableSegmentedDictionary<ProjectId, ICompilationTracker>? projectIdToTrackerMap = null,
         SourceGeneratorExecutionVersionMap? sourceGeneratorExecutionVersionMap = null,
-        Optional<TextDocumentStates<SourceGeneratedDocumentState>?> frozenSourceGeneratedDocumentStates = default,
+        Optional<TextDocumentStates<SourceGeneratedDocumentState>> frozenSourceGeneratedDocumentStates = default,
         AsyncLazy<SolutionCompilationState>? cachedFrozenSnapshot = null)
     {
         projectIdToTrackerMap ??= _projectIdToTrackerMap;
@@ -1277,7 +1277,7 @@ public ValueTask<ImmutableArray<Diagnostic>> GetSourceGeneratorDiagnosticsAsync(
     public SolutionCompilationState WithoutFrozenSourceGeneratedDocuments()
     {
         // If there's nothing frozen, there's nothing to do.
-        if (FrozenSourceGeneratedDocumentStates == null)
+        if (FrozenSourceGeneratedDocumentStates.IsEmpty)
             return this;
 
         var projectIdsToUnfreeze = FrozenSourceGeneratedDocumentStates.States.Values
@@ -1312,7 +1312,7 @@ public SolutionCompilationState WithoutFrozenSourceGeneratedDocuments()
         return this.Branch(
             this.SolutionState,
             projectIdToTrackerMap: newTrackerMap,
-            frozenSourceGeneratedDocumentStates: null);
+            frozenSourceGeneratedDocumentStates: TextDocumentStates<SourceGeneratedDocumentState>.Empty);
     }
 
     /// <summary>
@@ -1328,7 +1328,7 @@ public SolutionCompilationState WithFrozenSourceGeneratedDocuments(
         // and so those are always forks off the main solution. There's also a bit of a design question -- does calling this a second time
         // leave the existing frozen documents in place, or replace them? It depends on the need, but until then we'll cross that bridge
         // if/when we need it.
-        Contract.ThrowIfFalse(FrozenSourceGeneratedDocumentStates == null, $"We shouldn't be calling {nameof(WithFrozenSourceGeneratedDocuments)} on a solution with frozen source generated documents.");
+        Contract.ThrowIfTrue(FrozenSourceGeneratedDocumentStates.Count > 0, $"We shouldn't be calling {nameof(WithFrozenSourceGeneratedDocuments)} on a solution with frozen source generated documents.");
 
         if (documents.IsEmpty)
             return this;

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

@@ -123,20 +123,13 @@ public async Task<Checksum> GetChecksumAsync(ProjectId projectId, CancellationTo
                     projectCone = stateChecksums.ProjectCone;
                 }
 
-                ChecksumCollection? frozenSourceGeneratedDocumentIdentities = null;
-                DocumentChecksumsAndIds? frozenSourceGeneratedDocumentTexts = null;
-                ImmutableArray<DateTime> frozenSourceGeneratedDocumentGenerationDateTimes = default;
+                var serializer = this.SolutionState.Services.GetRequiredService<ISerializerService>();
+                var identityChecksums = FrozenSourceGeneratedDocumentStates.SelectAsArray(
+                    static (s, arg) => arg.serializer.CreateChecksum(s.Identity, cancellationToken: arg.cancellationToken), (serializer, cancellationToken));
 
-                if (FrozenSourceGeneratedDocumentStates != null)
-                {
-                    var serializer = this.SolutionState.Services.GetRequiredService<ISerializerService>();
-                    var identityChecksums = FrozenSourceGeneratedDocumentStates.SelectAsArray(
-                        static (s, arg) => arg.serializer.CreateChecksum(s.Identity, cancellationToken: arg.cancellationToken), (serializer, cancellationToken));
-
-                    frozenSourceGeneratedDocumentTexts = await FrozenSourceGeneratedDocumentStates.GetDocumentChecksumsAndIdsAsync(cancellationToken).ConfigureAwait(false);
-                    frozenSourceGeneratedDocumentIdentities = new ChecksumCollection(identityChecksums);
-                    frozenSourceGeneratedDocumentGenerationDateTimes = FrozenSourceGeneratedDocumentStates.SelectAsArray(d => d.GenerationDateTime);
-                }
+                var frozenSourceGeneratedDocumentTexts = await FrozenSourceGeneratedDocumentStates.GetDocumentChecksumsAndIdsAsync(cancellationToken).ConfigureAwait(false);
+                var frozenSourceGeneratedDocumentIdentities = new ChecksumCollection(identityChecksums);
+                var frozenSourceGeneratedDocumentGenerationDateTimes = FrozenSourceGeneratedDocumentStates.SelectAsArray(d => d.GenerationDateTime);
 
                 // Ensure we only send the execution map over for projects in the project cone.
                 var versionMapChecksum = this.GetFilteredSourceGenerationExecutionMap(projectCone).GetChecksum();

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

@@ -22,15 +22,10 @@ public SolutionCompilationStateChecksums(
         Checksum solutionState,
         Checksum sourceGeneratorExecutionVersionMap,
         // These arrays are all the same length if present, and reference the same documents in the same order.
-        DocumentChecksumsAndIds? frozenSourceGeneratedDocuments,
-        ChecksumCollection? frozenSourceGeneratedDocumentIdentities,
+        DocumentChecksumsAndIds frozenSourceGeneratedDocuments,
+        ChecksumCollection frozenSourceGeneratedDocumentIdentities,
         ImmutableArray<DateTime> frozenSourceGeneratedDocumentGenerationDateTimes)
     {
-        // For the frozen source generated document info, we expect two either have both checksum collections or neither, and they
-        // should both be the same length as there is a 1:1 correspondence between them.
-        Contract.ThrowIfFalse(frozenSourceGeneratedDocumentIdentities.HasValue == frozenSourceGeneratedDocuments.HasValue);
-        Contract.ThrowIfFalse(frozenSourceGeneratedDocumentIdentities?.Count == frozenSourceGeneratedDocuments?.Length);
-
         SolutionState = solutionState;
         SourceGeneratorExecutionVersionMap = sourceGeneratorExecutionVersionMap;
         FrozenSourceGeneratedDocuments = frozenSourceGeneratedDocuments;
@@ -42,8 +37,8 @@ public SolutionCompilationStateChecksums(
         Checksum = Checksum.Create(
             SolutionState,
             SourceGeneratorExecutionVersionMap,
-            FrozenSourceGeneratedDocumentIdentities?.Checksum ?? Checksum.Null,
-            frozenSourceGeneratedDocuments?.Checksum ?? Checksum.Null);
+            FrozenSourceGeneratedDocumentIdentities.Checksum,
+            frozenSourceGeneratedDocuments.Checksum);
     }
 
     public Checksum Checksum { get; }
@@ -53,8 +48,8 @@ public SolutionCompilationStateChecksums(
     /// <summary>
     /// Checksums of the SourceTexts of the frozen documents directly.  Not checksums of their DocumentStates.
     /// </summary>
-    public DocumentChecksumsAndIds? FrozenSourceGeneratedDocuments { get; }
-    public ChecksumCollection? FrozenSourceGeneratedDocumentIdentities { get; }
+    public DocumentChecksumsAndIds FrozenSourceGeneratedDocuments { get; }
+    public ChecksumCollection FrozenSourceGeneratedDocumentIdentities { get; }
 
     // note: intentionally not part of the identity contract of this type.
     public ImmutableArray<DateTime> FrozenSourceGeneratedDocumentGenerationDateTimes { get; }
@@ -64,8 +59,8 @@ public void AddAllTo(HashSet<Checksum> checksums)
         checksums.AddIfNotNullChecksum(this.Checksum);
         checksums.AddIfNotNullChecksum(this.SolutionState);
         checksums.AddIfNotNullChecksum(this.SourceGeneratorExecutionVersionMap);
-        this.FrozenSourceGeneratedDocumentIdentities?.AddAllTo(checksums);
-        this.FrozenSourceGeneratedDocuments?.AddAllTo(checksums);
+        this.FrozenSourceGeneratedDocumentIdentities.AddAllTo(checksums);
+        this.FrozenSourceGeneratedDocuments.AddAllTo(checksums);
     }
 
     public void Serialize(ObjectWriter writer)
@@ -76,13 +71,9 @@ public void Serialize(ObjectWriter writer)
         this.SourceGeneratorExecutionVersionMap.WriteTo(writer);
 
         // Write out a boolean to know whether we'll have this extra information
-        writer.WriteBoolean(this.FrozenSourceGeneratedDocumentIdentities.HasValue);
-        if (FrozenSourceGeneratedDocumentIdentities.HasValue)
-        {
-            this.FrozenSourceGeneratedDocuments!.Value.WriteTo(writer);
-            this.FrozenSourceGeneratedDocumentIdentities.Value.WriteTo(writer);
-            writer.WriteArray(this.FrozenSourceGeneratedDocumentGenerationDateTimes, static (w, d) => w.WriteInt64(d.Ticks));
-        }
+        this.FrozenSourceGeneratedDocuments.WriteTo(writer);
+        this.FrozenSourceGeneratedDocumentIdentities.WriteTo(writer);
+        writer.WriteArray(this.FrozenSourceGeneratedDocumentGenerationDateTimes, static (w, d) => w.WriteInt64(d.Ticks));
     }
 
     public static SolutionCompilationStateChecksums Deserialize(ObjectReader reader)
@@ -91,17 +82,9 @@ public static SolutionCompilationStateChecksums Deserialize(ObjectReader reader)
         var solutionState = Checksum.ReadFrom(reader);
         var sourceGeneratorExecutionVersionMap = Checksum.ReadFrom(reader);
 
-        var hasFrozenSourceGeneratedDocuments = reader.ReadBoolean();
-        DocumentChecksumsAndIds? frozenSourceGeneratedDocumentTexts = null;
-        ChecksumCollection? frozenSourceGeneratedDocumentIdentities = null;
-        ImmutableArray<DateTime> frozenSourceGeneratedDocumentGenerationDateTimes = default;
-
-        if (hasFrozenSourceGeneratedDocuments)
-        {
-            frozenSourceGeneratedDocumentTexts = DocumentChecksumsAndIds.ReadFrom(reader);
-            frozenSourceGeneratedDocumentIdentities = ChecksumCollection.ReadFrom(reader);
-            frozenSourceGeneratedDocumentGenerationDateTimes = reader.ReadArray(r => new DateTime(r.ReadInt64()));
-        }
+        var frozenSourceGeneratedDocumentTexts = DocumentChecksumsAndIds.ReadFrom(reader);
+        var frozenSourceGeneratedDocumentIdentities = ChecksumCollection.ReadFrom(reader);
+        var frozenSourceGeneratedDocumentGenerationDateTimes = reader.ReadArray(r => new DateTime(r.ReadInt64()));
 
         var result = new SolutionCompilationStateChecksums(
             solutionState: solutionState,
@@ -138,50 +121,44 @@ public async Task FindAsync<TArg>(
                 onAssetFound(this.SourceGeneratorExecutionVersionMap, filteredExecutionMap, arg);
             }
 
-            if (compilationState.FrozenSourceGeneratedDocumentStates != null)
+            // This could either be the checksum for the text (which we'll use our regular helper for first)...
+            if (assetPath.IncludeSolutionFrozenSourceGeneratedDocumentText)
             {
-                Contract.ThrowIfFalse(FrozenSourceGeneratedDocumentIdentities.HasValue);
-                Contract.ThrowIfFalse(FrozenSourceGeneratedDocuments.HasValue);
-
-                // This could either be the checksum for the text (which we'll use our regular helper for first)...
-                if (assetPath.IncludeSolutionFrozenSourceGeneratedDocumentText)
-                {
-                    await ChecksumCollection.FindAsync(
-                        new AssetPath(AssetPathKind.DocumentText, assetPath.ProjectId, assetPath.DocumentId),
-                        compilationState.FrozenSourceGeneratedDocumentStates, searchingChecksumsLeft, onAssetFound, arg, cancellationToken).ConfigureAwait(false);
-                }
+                await ChecksumCollection.FindAsync(
+                    new AssetPath(AssetPathKind.DocumentText, assetPath.ProjectId, assetPath.DocumentId),
+                    compilationState.FrozenSourceGeneratedDocumentStates, searchingChecksumsLeft, onAssetFound, arg, cancellationToken).ConfigureAwait(false);
+            }
 
-                // ... or one of the identities. In this case, we'll use the fact that there's a 1:1 correspondence between the
-                // two collections we hold onto.
-                if (assetPath.IncludeSolutionFrozenSourceGeneratedDocumentIdentities)
+            // ... or one of the identities. In this case, we'll use the fact that there's a 1:1 correspondence between the
+            // two collections we hold onto.
+            if (assetPath.IncludeSolutionFrozenSourceGeneratedDocumentIdentities)
+            {
+                var documentId = assetPath.DocumentId;
+                if (documentId != null)
                 {
-                    var documentId = assetPath.DocumentId;
-                    if (documentId != null)
+                    // If the caller is asking for a specific document, we can just look it up directly.
+                    var index = FrozenSourceGeneratedDocuments.Ids.IndexOf(documentId);
+                    if (index >= 0)
                     {
-                        // If the caller is asking for a specific document, we can just look it up directly.
-                        var index = FrozenSourceGeneratedDocuments.Value.Ids.IndexOf(documentId);
-                        if (index >= 0)
+                        var identityChecksum = FrozenSourceGeneratedDocumentIdentities.Children[index];
+                        if (searchingChecksumsLeft.Remove(identityChecksum))
                         {
-                            var identityChecksum = FrozenSourceGeneratedDocumentIdentities.Value.Children[index];
-                            if (searchingChecksumsLeft.Remove(identityChecksum))
-                            {
-                                Contract.ThrowIfFalse(compilationState.FrozenSourceGeneratedDocumentStates.TryGetState(documentId, out var state));
-                                onAssetFound(identityChecksum, state.Identity, arg);
-                            }
+                            Contract.ThrowIfFalse(compilationState.FrozenSourceGeneratedDocumentStates.TryGetState(documentId, out var state));
+                            onAssetFound(identityChecksum, state.Identity, arg);
                         }
                     }
-                    else
+                }
+                else
+                {
+                    // Otherwise, we'll have to search through all of them.
+                    for (var i = 0; i < FrozenSourceGeneratedDocumentIdentities.Count; i++)
                     {
-                        // Otherwise, we'll have to search through all of them.
-                        for (var i = 0; i < FrozenSourceGeneratedDocumentIdentities.Value.Count; i++)
+                        var identityChecksum = FrozenSourceGeneratedDocumentIdentities[0];
+                        if (searchingChecksumsLeft.Remove(identityChecksum))
                         {
-                            var identityChecksum = FrozenSourceGeneratedDocumentIdentities.Value[0];
-                            if (searchingChecksumsLeft.Remove(identityChecksum))
-                            {
-                                var id = FrozenSourceGeneratedDocuments.Value.Ids[i];
-                                Contract.ThrowIfFalse(compilationState.FrozenSourceGeneratedDocumentStates.TryGetState(id, out var state));
-                                onAssetFound(identityChecksum, state.Identity, arg);
-                            }
+                            var id = FrozenSourceGeneratedDocuments.Ids[i];
+                            Contract.ThrowIfFalse(compilationState.FrozenSourceGeneratedDocumentStates.TryGetState(id, out var state));
+                            onAssetFound(identityChecksum, state.Identity, arg);
                         }
                     }
                 }

src/Workspaces/Remote/ServiceHub/Host/RemoteWorkspace.SolutionCreator.cs

@@ -82,12 +82,9 @@ public async Task<Solution> CreateSolutionAsync(Checksum newSolutionChecksum, Ca
                         AssetPathKind.SolutionFallbackAnalyzerOptions, newSolutionChecksums.FallbackAnalyzerOptions, cancellationToken).ConfigureAwait(false));
                 }
 
-                if (newSolutionCompilationChecksums.FrozenSourceGeneratedDocumentIdentities.HasValue &&
-                    newSolutionCompilationChecksums.FrozenSourceGeneratedDocuments.HasValue &&
-                    !newSolutionCompilationChecksums.FrozenSourceGeneratedDocumentGenerationDateTimes.IsDefault)
                 {
-                    var newSolutionFrozenSourceGeneratedDocumentIdentities = newSolutionCompilationChecksums.FrozenSourceGeneratedDocumentIdentities.Value;
-                    var newSolutionFrozenSourceGeneratedDocuments = newSolutionCompilationChecksums.FrozenSourceGeneratedDocuments.Value;
+                    var newSolutionFrozenSourceGeneratedDocumentIdentities = newSolutionCompilationChecksums.FrozenSourceGeneratedDocumentIdentities;
+                    var newSolutionFrozenSourceGeneratedDocuments = newSolutionCompilationChecksums.FrozenSourceGeneratedDocuments;
                     var count = newSolutionFrozenSourceGeneratedDocuments.Ids.Length;
 
                     var frozenDocuments = new FixedSizeArrayBuilder<(SourceGeneratedDocumentIdentity identity, DateTime generationDateTime, SourceText text)>(count);