Use the same code document everywhere

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

@@ -266,7 +266,7 @@ internal class FormattingContext : IDisposable
 
         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

@@ -41,7 +41,7 @@ internal class RazorFormattingService : IRazorFormattingService
         FormattingOptions options,
         CancellationToken cancellationToken)
     {
-        var codeDocument = await GetCodeDocumentAsync(documentContext).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:
@@ -86,34 +86,6 @@ internal class RazorFormattingService : IRazorFormattingService
         return GetMinimalEdits(originalText, filteredEdits);
     }
 
-    private static Task<RazorCodeDocument> GetCodeDocumentAsync(VersionedDocumentContext documentContext)
-    {
-        var snapshot = documentContext.Snapshot;
-        var forceRuntimeCodeGeneration = documentContext.Project.Configuration.LanguageServerFlags?.ForceRuntimeCodeGeneration ?? false;
-        if (!forceRuntimeCodeGeneration)
-        {
-            return snapshot.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(documentContext);
-    }
-
-    private static async Task<RazorCodeDocument> GetDesignTimeDocumentAsync(VersionedDocumentContext documentContext)
-    {
-        var snapshot = documentContext.Snapshot;
-        var project = documentContext.Project;
-        var tagHelpers = await project.GetTagHelpersAsync(CancellationToken.None).ConfigureAwait(false);
-        var projectEngine = project.GetProjectEngine();
-        var imports = await DocumentState.GetImportsAsync(snapshot, projectEngine).ConfigureAwait(false);
-        return await DocumentState.GenerateCodeDocumentAsync(tagHelpers, project.GetProjectEngine(), snapshot, imports, false).ConfigureAwait(false);
-    }
-
     private static TextEdit[] GetMinimalEdits(SourceText originalText, IEnumerable<TextEdit> filteredEdits)
     {
         // Make sure the edits actually change something, or its not worth responding

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

@@ -59,4 +59,30 @@ internal static class IDocumentSnapshotExtensions
         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);
+    }
 }