Always use designtime document for formatting with FUSE

src/Compiler/Microsoft.CodeAnalysis.Razor.Compiler/src/Language/RazorProjectEngine.cs

@@ -52,7 +52,7 @@ public RazorCodeDocument Process(RazorProjectItem projectItem)
 
     public RazorCodeDocument Process(
         RazorSourceDocument source,
-        string fileKind,
+        string? fileKind,
         ImmutableArray<RazorSourceDocument> importSources,
         IReadOnlyList<TagHelperDescriptor> tagHelpers)
     {
@@ -109,9 +109,7 @@ public RazorCodeDocument ProcessDesignTime(RazorProjectItem projectItem)
             throw new ArgumentNullException(nameof(projectItem));
         }
 
-        var codeDocument = Configuration.LanguageServerFlags?.ForceRuntimeCodeGeneration == true
-            ? CreateCodeDocumentCore(projectItem)
-            : CreateCodeDocumentDesignTimeCore(projectItem);
+        var codeDocument = CreateCodeDocumentDesignTimeCore(projectItem);
         ProcessCore(codeDocument);
         return codeDocument;
     }
@@ -127,9 +125,7 @@ public RazorCodeDocument ProcessDesignTime(
             throw new ArgumentNullException(nameof(source));
         }
 
-        var codeDocument = Configuration.LanguageServerFlags?.ForceRuntimeCodeGeneration == true
-            ? CreateCodeDocumentCore(source, fileKind, importSources, tagHelpers, configureParser: null, configureCodeGeneration: null)
-            : CreateCodeDocumentDesignTimeCore(source, fileKind, importSources, tagHelpers, configureParser: null, configureCodeGeneration: null);
+        var codeDocument = CreateCodeDocumentDesignTimeCore(source, fileKind, importSources, tagHelpers, configureParser: null, configureCodeGeneration: null);
         ProcessCore(codeDocument);
         return codeDocument;
     }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Formatting/FormattingContext.cs

@@ -266,7 +266,7 @@ public async Task<FormattingContext> WithTextAsync(SourceText changedText)
 
         var changedSnapshot = OriginalSnapshot.WithText(changedText);
 
-        var codeDocument = await changedSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await changedSnapshot.GetFormatterCodeDocumentAsync().ConfigureAwait(false);
 
         DEBUG_ValidateComponents(CodeDocument, codeDocument);
 

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Formatting/RazorFormattingService.cs

@@ -8,6 +8,7 @@
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.TextDifferencing;
+using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Razor.Protocol;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
@@ -40,7 +41,7 @@ public async Task<TextEdit[]> FormatAsync(
         FormattingOptions options,
         CancellationToken cancellationToken)
     {
-        var codeDocument = await documentContext.Snapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await documentContext.Snapshot.GetFormatterCodeDocumentAsync().ConfigureAwait(false);
 
         // Range formatting happens on every paste, and if there are Razor diagnostics in the file
         // that can make some very bad results. eg, given:

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Extensions/IDocumentSnapshotExtensions.cs

@@ -59,4 +59,30 @@ public static bool IsPathCandidateForComponent(this IDocumentSnapshot documentSn
         var fileName = Path.GetFileNameWithoutExtension(documentSnapshot.FilePath);
         return fileName.AsSpan().Equals(path.Span, FilePathComparison.Instance);
     }
+
+    public static Task<RazorCodeDocument> GetFormatterCodeDocumentAsync(this IDocumentSnapshot documentSnapshot)
+    {
+        var forceRuntimeCodeGeneration = documentSnapshot.Project.Configuration.LanguageServerFlags?.ForceRuntimeCodeGeneration ?? false;
+        if (!forceRuntimeCodeGeneration)
+        {
+            return documentSnapshot.GetGeneratedOutputAsync();
+        }
+
+        // if forceRuntimeCodeGeneration is on, GetGeneratedOutputAsync will get runtime code. As of now
+        // the formatting service doesn't expect the form of code generated to be what the compiler does with
+        // runtime. For now force usage of design time and avoid the cache. There may be a slight perf hit
+        // but either the user is typing (which will invalidate the cache) or the user is manually attempting to
+        // format. We expect formatting to invalidate the cache if it changes things and consider this an
+        // acceptable overhead for now.
+        return GetDesignTimeDocumentAsync(documentSnapshot);
+    }
+
+    private static async Task<RazorCodeDocument> GetDesignTimeDocumentAsync(IDocumentSnapshot documentSnapshot)
+    {
+        var project = documentSnapshot.Project;
+        var tagHelpers = await project.GetTagHelpersAsync(CancellationToken.None).ConfigureAwait(false);
+        var projectEngine = project.GetProjectEngine();
+        var imports = await DocumentState.GetImportsAsync(documentSnapshot, projectEngine).ConfigureAwait(false);
+        return await DocumentState.GenerateCodeDocumentAsync(tagHelpers, project.GetProjectEngine(), documentSnapshot, imports, false).ConfigureAwait(false);
+    }
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/ProjectSystem/DocumentState.ComputedStateTracker.cs

@@ -0,0 +1,252 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System;
+using System.Diagnostics.CodeAnalysis;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor.Language;
+
+namespace Microsoft.CodeAnalysis.Razor.ProjectSystem;
+
+internal partial class DocumentState
+{
+    private class ComputedStateTracker
+    {
+        private readonly object _lock;
+
+        private ComputedStateTracker? _older;
+
+        // We utilize a WeakReference here to avoid bloating committed memory. If pieces request document output inbetween GC collections
+        // then we will provide the weak referenced task; otherwise we require any state requests to be re-computed.
+        private WeakReference<Task<(RazorCodeDocument, VersionStamp)>>? _taskUnsafeReference;
+
+        private ComputedOutput? _computedOutput;
+
+        public ComputedStateTracker(DocumentState state, ComputedStateTracker? older = null)
+        {
+            _lock = state._lock;
+            _older = older;
+        }
+
+        public bool IsResultAvailable
+        {
+            get
+            {
+                if (_computedOutput?.TryGetCachedOutput(out _, out _) == true)
+                {
+                    return true;
+                }
+
+                if (_taskUnsafeReference is null)
+                {
+                    return false;
+                }
+
+                if (_taskUnsafeReference.TryGetTarget(out var taskUnsafe))
+                {
+                    return taskUnsafe.IsCompleted;
+                }
+
+                return false;
+            }
+        }
+
+        public async Task<(RazorCodeDocument, VersionStamp)> GetGeneratedOutputAndVersionAsync(ProjectSnapshot project, IDocumentSnapshot document)
+        {
+            if (_computedOutput?.TryGetCachedOutput(out var cachedCodeDocument, out var cachedInputVersion) == true)
+            {
+                return (cachedCodeDocument, cachedInputVersion);
+            }
+
+            var (codeDocument, inputVersion) = await GetMemoizedGeneratedOutputAndVersionAsync(project, document).ConfigureAwait(false);
+
+            _computedOutput = new ComputedOutput(codeDocument, inputVersion);
+            return (codeDocument, inputVersion);
+        }
+
+        private Task<(RazorCodeDocument, VersionStamp)> GetMemoizedGeneratedOutputAndVersionAsync(ProjectSnapshot project, IDocumentSnapshot document)
+        {
+            if (project is null)
+            {
+                throw new ArgumentNullException(nameof(project));
+            }
+
+            if (document is null)
+            {
+                throw new ArgumentNullException(nameof(document));
+            }
+
+            if (_taskUnsafeReference is null ||
+                !_taskUnsafeReference.TryGetTarget(out var taskUnsafe))
+            {
+                TaskCompletionSource<(RazorCodeDocument, VersionStamp)>? tcs = null;
+
+                lock (_lock)
+                {
+                    if (_taskUnsafeReference is null ||
+                        !_taskUnsafeReference.TryGetTarget(out taskUnsafe))
+                    {
+                        // So this is a bit confusing. Instead of directly calling the Razor parser inside of this lock we create an indirect TaskCompletionSource
+                        // to represent when it completes. The reason behind this is that there are several scenarios in which the Razor parser will run synchronously
+                        // (mostly all in VS) resulting in this lock being held for significantly longer than expected. To avoid threads queuing up repeatedly on the
+                        // above lock and blocking we can allow those threads to await asynchronously for the completion of the original parse.
+
+                        tcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
+                        taskUnsafe = tcs.Task;
+                        _taskUnsafeReference = new WeakReference<Task<(RazorCodeDocument, VersionStamp)>>(taskUnsafe);
+                    }
+                }
+
+                if (tcs is null)
+                {
+                    // There's no task completion source created meaning a value was retrieved from cache, just return it.
+                    return taskUnsafe;
+                }
+
+                // Typically in VS scenarios this will run synchronously because all resources are readily available.
+                var outputTask = ComputeGeneratedOutputAndVersionAsync(project, document);
+                if (outputTask.IsCompleted)
+                {
+                    // Compiling ran synchronously, lets just immediately propagate to the TCS
+                    PropagateToTaskCompletionSource(outputTask, tcs);
+                }
+                else
+                {
+                    // Task didn't run synchronously (most likely outside of VS), lets allocate a bit more but utilize ContinueWith
+                    // to properly connect the output task and TCS
+                    _ = outputTask.ContinueWith(
+                        static (task, state) =>
+                        {
+                            Assumes.NotNull(state);
+                            var tcs = (TaskCompletionSource<(RazorCodeDocument, VersionStamp)>)state;
+
+                            PropagateToTaskCompletionSource(task, tcs);
+                        },
+                        tcs,
+                        CancellationToken.None,
+                        TaskContinuationOptions.ExecuteSynchronously,
+                        TaskScheduler.Default);
+                }
+            }
+
+            return taskUnsafe;
+
+            static void PropagateToTaskCompletionSource(
+                Task<(RazorCodeDocument, VersionStamp)> targetTask,
+                TaskCompletionSource<(RazorCodeDocument, VersionStamp)> tcs)
+            {
+                if (targetTask.Status == TaskStatus.RanToCompletion)
+                {
+#pragma warning disable VSTHRD103 // Call async methods when in an async method
+                    tcs.SetResult(targetTask.Result);
+#pragma warning restore VSTHRD103 // Call async methods when in an async method
+                }
+                else if (targetTask.Status == TaskStatus.Canceled)
+                {
+                    tcs.SetCanceled();
+                }
+                else if (targetTask.Status == TaskStatus.Faulted)
+                {
+                    // Faulted tasks area always aggregate exceptions so we need to extract the "true" exception if it's available:
+                    Assumes.NotNull(targetTask.Exception);
+                    var exception = targetTask.Exception.InnerException ?? targetTask.Exception;
+                    tcs.SetException(exception);
+                }
+            }
+        }
+
+        private async Task<(RazorCodeDocument, VersionStamp)> ComputeGeneratedOutputAndVersionAsync(ProjectSnapshot project, IDocumentSnapshot document)
+        {
+            // We only need to produce the generated code if any of our inputs is newer than the
+            // previously cached output.
+            //
+            // First find the versions that are the inputs:
+            // - The project + computed state
+            // - The imports
+            // - This document
+            //
+            // All of these things are cached, so no work is wasted if we do need to generate the code.
+            var configurationVersion = project.State.ConfigurationVersion;
+            var projectWorkspaceStateVersion = project.State.ProjectWorkspaceStateVersion;
+            var documentCollectionVersion = project.State.DocumentCollectionVersion;
+            var imports = await GetImportsAsync(document, project.GetProjectEngine()).ConfigureAwait(false);
+            var documentVersion = await document.GetTextVersionAsync().ConfigureAwait(false);
+
+            // OK now that have the previous output and all of the versions, we can see if anything
+            // has changed that would require regenerating the code.
+            var inputVersion = documentVersion;
+            if (inputVersion.GetNewerVersion(configurationVersion) == configurationVersion)
+            {
+                inputVersion = configurationVersion;
+            }
+
+            if (inputVersion.GetNewerVersion(projectWorkspaceStateVersion) == projectWorkspaceStateVersion)
+            {
+                inputVersion = projectWorkspaceStateVersion;
+            }
+
+            if (inputVersion.GetNewerVersion(documentCollectionVersion) == documentCollectionVersion)
+            {
+                inputVersion = documentCollectionVersion;
+            }
+
+            foreach (var import in imports)
+            {
+                var importVersion = import.Version;
+                if (inputVersion.GetNewerVersion(importVersion) == importVersion)
+                {
+                    inputVersion = importVersion;
+                }
+            }
+
+            if (_older?._taskUnsafeReference != null &&
+                _older._taskUnsafeReference.TryGetTarget(out var taskUnsafe))
+            {
+                var (olderOutput, olderInputVersion) = await taskUnsafe.ConfigureAwait(false);
+                if (inputVersion.GetNewerVersion(olderInputVersion) == olderInputVersion)
+                {
+                    // Nothing has changed, we can use the cached result.
+                    lock (_lock)
+                    {
+                        _taskUnsafeReference = _older._taskUnsafeReference;
+                        _older = null;
+                        return (olderOutput, olderInputVersion);
+                    }
+                }
+            }
+
+            var tagHelpers = await project.GetTagHelpersAsync(CancellationToken.None).ConfigureAwait(false);
+            var forceRuntimeCodeGeneration = project.Configuration.LanguageServerFlags?.ForceRuntimeCodeGeneration ?? false;
+            var codeDocument = await GenerateCodeDocumentAsync(tagHelpers, project.GetProjectEngine(), document, imports, forceRuntimeCodeGeneration).ConfigureAwait(false);
+            return (codeDocument, inputVersion);
+        }
+
+        private class ComputedOutput
+        {
+            private readonly VersionStamp _inputVersion;
+            private readonly WeakReference<RazorCodeDocument> _codeDocumentReference;
+
+            public ComputedOutput(RazorCodeDocument codeDocument, VersionStamp inputVersion)
+            {
+                _codeDocumentReference = new WeakReference<RazorCodeDocument>(codeDocument);
+                _inputVersion = inputVersion;
+            }
+
+            public bool TryGetCachedOutput([NotNullWhen(true)] out RazorCodeDocument? codeDocument, out VersionStamp inputVersion)
+            {
+                // The goal here is to capture a weak reference to the code document so if there's ever a sub-system that's still utilizing it
+                // our computed output maintains its cache.
+
+                if (_codeDocumentReference.TryGetTarget(out codeDocument))
+                {
+                    inputVersion = _inputVersion;
+                    return true;
+                }
+
+                inputVersion = default;
+                return false;
+            }
+        }
+    }
+}

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/ProjectSystem/DocumentState.ImportItem.cs

@@ -0,0 +1,12 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+namespace Microsoft.CodeAnalysis.Razor.ProjectSystem;
+
+internal partial class DocumentState
+{
+    internal record struct ImportItem(string? FilePath, VersionStamp Version, IDocumentSnapshot Document)
+    {
+        public readonly string? FileKind => Document.FileKind;
+    }
+}