Merge pull request #69747 from jasonmalinowski/add-additional-validat…

…ion-for-trees-being-in-sync

Add extra validation for Compilations having the expected trees

src/Features/LanguageServer/Protocol/Features/Options/WorkspaceConfigurationOptionsStorage.cs

@@ -16,7 +16,8 @@ public static WorkspaceConfigurationOptions GetWorkspaceConfigurationOptions(thi
             EnableOpeningSourceGeneratedFiles: globalOptions.GetOption(EnableOpeningSourceGeneratedFilesInWorkspace) ??
                                                globalOptions.GetOption(EnableOpeningSourceGeneratedFilesInWorkspaceFeatureFlag),
             DisableSharedSyntaxTrees: globalOptions.GetOption(DisableSharedSyntaxTrees),
-            DisableRecoverableText: globalOptions.GetOption(DisableRecoverableText));
+            DisableRecoverableText: globalOptions.GetOption(DisableRecoverableText),
+            ValidateCompilationTrackerStates: globalOptions.GetOption(ValidateCompilationTrackerStates));
 
     public static readonly Option2<StorageDatabase> Database = new(
         "dotnet_storage_database", WorkspaceConfigurationOptions.Default.CacheStorage, serializer: EditorConfigValueSerializer.CreateSerializerForEnum<StorageDatabase>());
@@ -30,6 +31,9 @@ public static WorkspaceConfigurationOptions GetWorkspaceConfigurationOptions(thi
     public static readonly Option2<bool> DisableRecoverableText = new(
         "dotnet_disable_recoverable_text", WorkspaceConfigurationOptions.Default.DisableRecoverableText);
 
+    public static readonly Option2<bool> ValidateCompilationTrackerStates = new Option2<bool>(
+        "dotnet_validate_compilation_tracker_states", WorkspaceConfigurationOptions.Default.ValidateCompilationTrackerStates);
+
     /// <summary>
     /// This option allows the user to enable this. We are putting this behind a feature flag for now since we could have extensions
     /// surprised by this and we want some time to work through those issues.

src/VisualStudio/Core/Def/Options/VisualStudioOptionStorage.cs

@@ -433,6 +433,7 @@ public bool TryFetch(LocalUserRegistryOptionPersister persister, OptionKey2 opti
         {"visual_studio_workspace_partial_load_mode", new FeatureFlagStorage(@"Roslyn.PartialLoadMode")},
         {"dotnet_disable_shared_syntax_trees", new FeatureFlagStorage(@"Roslyn.DisableSharedSyntaxTrees")},
         {"dotnet_disable_recoverable_text", new FeatureFlagStorage(@"Roslyn.DisableRecoverableText")},
+        {"dotnet_validate_compilation_tracker_states", new FeatureFlagStorage(@"Roslyn.ValidateCompilationTrackerStates")},
         {"dotnet_enable_diagnostics_in_source_generated_files", new RoamingProfileStorage("TextEditor.Roslyn.Specific.EnableDiagnosticsInSourceGeneratedFilesExperiment")},
         {"dotnet_enable_diagnostics_in_source_generated_files_feature_flag", new FeatureFlagStorage(@"Roslyn.EnableDiagnosticsInSourceGeneratedFiles")},
         {"dotnet_enable_opening_source_generated_files_in_workspace", new RoamingProfileStorage("TextEditor.Roslyn.Specific.EnableOpeningSourceGeneratedFilesInWorkspaceExperiment")},

src/Workspaces/Core/Portable/Workspace/IWorkspaceConfigurationService.cs

@@ -41,7 +41,14 @@ internal readonly record struct WorkspaceConfigurationOptions(
         [property: DataMember(Order = 0)] StorageDatabase CacheStorage = StorageDatabase.SQLite,
         [property: DataMember(Order = 1)] bool EnableOpeningSourceGeneratedFiles = false,
         [property: DataMember(Order = 2)] bool DisableSharedSyntaxTrees = false,
-        [property: DataMember(Order = 3)] bool DisableRecoverableText = false)
+        [property: DataMember(Order = 3)] bool DisableRecoverableText = false,
+        [property: DataMember(Order = 4)] bool ValidateCompilationTrackerStates =
+#if DEBUG // We will default this on in DEBUG builds
+            true
+#else
+            false
+#endif
+        )
     {
         public WorkspaceConfigurationOptions()
             : this(CacheStorage: StorageDatabase.SQLite)
@@ -52,7 +59,7 @@ public WorkspaceConfigurationOptions()
 
         /// <summary>
         /// These values are such that the correctness of remote services is not affected if these options are changed from defaults
-        /// to non-defauls while the services have already been executing.
+        /// to non-defaults while the services have already been executing.
         /// </summary>
         public static readonly WorkspaceConfigurationOptions RemoteDefault = new(
             CacheStorage: StorageDatabase.None,

src/Workspaces/Core/Portable/Workspace/Solution/SolutionState.CompilationTracker.cs

@@ -47,6 +47,12 @@ private partial class CompilationTracker : ICompilationTracker
 
             public SkeletonReferenceCache SkeletonReferenceCache { get; }
 
+            /// <summary>
+            /// Set via a feature flag to enable strict validation of the compilations that are produced, in that they match the original states. This validation is expensive, so we don't want it
+            /// running in normal production scenarios.
+            /// </summary>
+            private readonly bool _validateStates;
+
             private CompilationTracker(
                 ProjectState project,
                 CompilationTrackerState state,
@@ -57,6 +63,10 @@ private CompilationTracker(
                 this.ProjectState = project;
                 _stateDoNotAccessDirectly = state;
                 this.SkeletonReferenceCache = cachedSkeletonReferences;
+
+                _validateStates = project.LanguageServices.SolutionServices.GetRequiredService<IWorkspaceConfigurationService>().Options.ValidateCompilationTrackerStates;
+
+                ValidateState(state);
             }
 
             /// <summary>
@@ -72,7 +82,10 @@ private CompilationTrackerState ReadState()
                 => Volatile.Read(ref _stateDoNotAccessDirectly);
 
             private void WriteState(CompilationTrackerState state)
-                => Volatile.Write(ref _stateDoNotAccessDirectly, state);
+            {
+                Volatile.Write(ref _stateDoNotAccessDirectly, state);
+                ValidateState(state);
+            }
 
             public GeneratorDriver? GeneratorDriver
             {
@@ -1116,6 +1129,67 @@ private bool IsGeneratorRunResultToIgnore(GeneratorRunResult result)
             // END HACK HACK HACK HACK, or the setup of it at least; once this hack is removed the calls to IsGeneratorRunResultToIgnore
             // need to be cleaned up.
 
+            /// <summary>
+            /// Validates the compilation is consistent and we didn't have a bug in producing it. This only runs under a feature flag.
+            /// </summary>
+            private void ValidateState(CompilationTrackerState state)
+            {
+                if (!_validateStates)
+                    return;
+
+                if (state is FinalState finalState)
+                {
+                    ValidateCompilationTreesMatchesProjectState(finalState.FinalCompilationWithGeneratedDocuments, ProjectState, state.GeneratorInfo);
+                }
+                else if (state is InProgressState inProgressState)
+                {
+                    ValidateCompilationTreesMatchesProjectState(inProgressState.CompilationWithoutGeneratedDocuments!, inProgressState.IntermediateProjects[0].oldState, generatorInfo: null);
+
+                    if (inProgressState.CompilationWithGeneratedDocuments != null)
+                    {
+                        ValidateCompilationTreesMatchesProjectState(inProgressState.CompilationWithGeneratedDocuments, inProgressState.IntermediateProjects[0].oldState, inProgressState.GeneratorInfo);
+                    }
+                }
+            }
+
+            private static void ValidateCompilationTreesMatchesProjectState(Compilation compilation, ProjectState projectState, CompilationTrackerGeneratorInfo? generatorInfo)
+            {
+                // We'll do this all in a try/catch so it makes validations easy to do with Contract.ThrowIfFalse().
+                try
+                {
+                    // Assert that all the trees we expect to see are in the Compilation...
+                    var syntaxTreesInWorkspaceStates = new HashSet<SyntaxTree>(
+#if NET
+                        capacity: projectState.DocumentStates.Count + generatorInfo?.Documents.Count ?? 0
+#endif
+                        );
+
+                    foreach (var documentInProjectState in projectState.DocumentStates.States)
+                    {
+                        Contract.ThrowIfFalse(documentInProjectState.Value.TryGetSyntaxTree(out var tree), "We should have a tree since we have a compilation that should contain it.");
+                        syntaxTreesInWorkspaceStates.Add(tree);
+                        Contract.ThrowIfFalse(compilation.ContainsSyntaxTree(tree), "The tree in the ProjectState should have been in the compilation.");
+                    }
+
+                    if (generatorInfo != null)
+                    {
+                        foreach (var generatedDocument in generatorInfo.Value.Documents.States)
+                        {
+                            Contract.ThrowIfFalse(generatedDocument.Value.TryGetSyntaxTree(out var tree), "We should have a tree since we have a compilation that should contain it.");
+                            syntaxTreesInWorkspaceStates.Add(tree);
+                            Contract.ThrowIfFalse(compilation.ContainsSyntaxTree(tree), "The tree for the generated document should have been in the compilation.");
+                        }
+                    }
+
+                    // ...and that the reverse is true too.
+                    foreach (var tree in compilation.SyntaxTrees)
+                        Contract.ThrowIfFalse(syntaxTreesInWorkspaceStates.Contains(tree), "The tree in the Compilation should have been from the workspace.");
+                }
+                catch (Exception e) when (FatalError.ReportWithDumpAndCatch(e, ErrorSeverity.Critical))
+                {
+                }
+            }
+
             #region Versions and Checksums
 
             // Dependent Versions are stored on compilation tracker so they are more likely to survive when unrelated solution branching occurs.