Finish off the CohostDocumentSnapshot by sharing reusing code from DocumentState

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Cohost/CohostDocumentSnapshot.cs

@@ -1,7 +1,6 @@
 // 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.Collections.Immutable;
 using System.Diagnostics.CodeAnalysis;
 using System.Threading.Tasks;
@@ -17,6 +16,8 @@ internal class CohostDocumentSnapshot(TextDocument textDocument, IProjectSnapsho
     private readonly TextDocument _textDocument = textDocument;
     private readonly IProjectSnapshot _projectSnapshot = projectSnapshot;
 
+    private RazorCodeDocument? _codeDocument;
+
     public string? FileKind => FileKinds.GetFileKindFromFilePath(FilePath);
 
     public string? FilePath => _textDocument.FilePath;
@@ -37,16 +38,36 @@ internal class CohostDocumentSnapshot(TextDocument textDocument, IProjectSnapsho
 
     public ImmutableArray<IDocumentSnapshot> GetImports()
     {
-        throw new NotImplementedException();
+        return DocumentState.GetImportsCore(Project, FilePath.AssumeNotNull(), FileKind.AssumeNotNull());
     }
 
-    public Task<RazorCodeDocument> GetGeneratedOutputAsync()
+    public async Task<RazorCodeDocument> GetGeneratedOutputAsync()
     {
-        throw new NotImplementedException();
+        // TODO: We don't need to worry about locking if we get called from the didOpen/didChange LSP requests, as CLaSP
+        //       takes care of that for us, and blocks requests until those are complete. If that doesn't end up happening,
+        //       then a locking mechanism here would prevent concurrent compilations.
+        if (_codeDocument is not null)
+        {
+            return _codeDocument;
+        }
+
+        // The non-cohosted DocumentSnapshot implementation uses DocumentState to get the generated output, and we could do that too
+        // but most of that code is optimized around caching pre-computed results when things change that don't affect the compilation.
+        // We can't do that here because we are using Roslyn's project snapshots, which don't contain the info that Razor needs. We could
+        // in future provide a side-car mechanism so we can cache things, but still take advantage of snapshots etc. but the working
+        // assumption for this code is that the source generator will be used, and it will do all of that, so this implementation is naive
+        // and simply compiles when asked, and if a new document snapshot comes in, we compile again. This is presumably worse for perf
+        // but since we don't expect users to ever use cohosting without source generators, it's fine for now.
+
+        var imports = await DocumentState.ComputedStateTracker.GetImportsAsync(this).ConfigureAwait(false);
+        _codeDocument = await DocumentState.ComputedStateTracker.GenerateCodeDocumentAsync(Project, this, imports).ConfigureAwait(false);
+
+        return _codeDocument;
     }
 
     public bool TryGetGeneratedOutput([NotNullWhen(true)] out RazorCodeDocument? result)
     {
-        throw new NotImplementedException();
+        result = _codeDocument;
+        return result is not null;
     }
 }

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

@@ -82,7 +82,7 @@ internal class DocumentState
 
     public ImmutableArray<IDocumentSnapshot> GetImports(ProjectSnapshot project)
     {
-        return GetImportsCore(project);
+        return GetImportsCore(project, HostDocument.FilePath, HostDocument.FileKind);
     }
 
     public async Task<SourceText> GetTextAsync()
@@ -226,10 +226,11 @@ internal class DocumentState
         return new DocumentState(HostDocument, null, null, loader);
     }
 
-    private ImmutableArray<IDocumentSnapshot> GetImportsCore(ProjectSnapshot project)
+    // Internal, because we are temporarily sharing code with CohostDocumentSnapshot
+    internal static ImmutableArray<IDocumentSnapshot> GetImportsCore(IProjectSnapshot project, string filePath, string fileKind)
     {
         var projectEngine = project.GetProjectEngine();
-        var projectItem = projectEngine.FileSystem.GetItem(HostDocument.FilePath, HostDocument.FileKind);
+        var projectItem = projectEngine.FileSystem.GetItem(filePath, fileKind);
 
         using var _1 = ListPool<RazorProjectItem>.GetPooledObject(out var importItems);
 
@@ -265,8 +266,8 @@ internal class DocumentState
         return imports.ToImmutable();
     }
 
-    // See design notes on ProjectState.ComputedStateTracker.
-    private class ComputedStateTracker
+    // Internal, because we are temporarily sharing code with CohostDocumentSnapshot
+    internal class ComputedStateTracker
     {
         private readonly object _lock;
 
@@ -471,6 +472,12 @@ internal class DocumentState
                 }
             }
 
+            var codeDocument = await GenerateCodeDocumentAsync(project, document, imports).ConfigureAwait(false);
+            return (codeDocument, inputVersion);
+        }
+
+        internal static async Task<RazorCodeDocument> GenerateCodeDocumentAsync(IProjectSnapshot project, IDocumentSnapshot document, ImmutableArray<ImportItem> imports)
+        {
             // OK we have to generate the code.
             using var importSources = new PooledArrayBuilder<RazorSourceDocument>(imports.Length);
             var projectEngine = project.GetProjectEngine();
@@ -484,8 +491,7 @@ internal class DocumentState
             var projectItem = document.FilePath is null ? null : projectEngine.FileSystem.GetItem(document.FilePath, document.FileKind);
             var documentSource = await GetRazorSourceDocumentAsync(document, projectItem).ConfigureAwait(false);
 
-            var codeDocument = projectEngine.ProcessDesignTime(documentSource, fileKind: document.FileKind, importSources.DrainToImmutable(), project.TagHelpers);
-            return (codeDocument, inputVersion);
+            return projectEngine.ProcessDesignTime(documentSource, fileKind: document.FileKind, importSources.DrainToImmutable(), project.TagHelpers);
         }
 
         private static async Task<RazorSourceDocument> GetRazorSourceDocumentAsync(IDocumentSnapshot document, RazorProjectItem? projectItem)
@@ -494,7 +500,7 @@ internal class DocumentState
             return RazorSourceDocument.Create(sourceText, RazorSourceDocumentProperties.Create(document.FilePath, projectItem?.RelativePhysicalPath));
         }
 
-        private static async Task<ImmutableArray<ImportItem>> GetImportsAsync(IDocumentSnapshot document)
+        internal static async Task<ImmutableArray<ImportItem>> GetImportsAsync(IDocumentSnapshot document)
         {
             var imports = document.GetImports();
             using var result = new PooledArrayBuilder<ImportItem>(imports.Length);
@@ -508,7 +514,7 @@ internal class DocumentState
             return result.DrainToImmutable();
         }
 
-        private record struct ImportItem(string? FilePath, VersionStamp Version, IDocumentSnapshot Document)
+        internal record struct ImportItem(string? FilePath, VersionStamp Version, IDocumentSnapshot Document)
         {
             public readonly string? FileKind => Document.FileKind;
         }