Remove ImportDocumentSnapshot and (finally) fix nullability annotations on IDocumentSnapshot

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/CodeDocumentReferenceHolder.cs

@@ -30,7 +30,7 @@ public void DocumentProcessed(RazorCodeDocument codeDocument, IDocumentSnapshot
         // multiple parses/regenerations across LSP requests that are all for the same document version.
         lock (_codeDocumentCache)
         {
-            var key = new DocumentKey(documentSnapshot.Project.Key, documentSnapshot.FilePath.AssumeNotNull());
+            var key = new DocumentKey(documentSnapshot.Project.Key, documentSnapshot.FilePath);
             _codeDocumentCache[key] = codeDocument;
         }
     }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Diagnostics/RazorDiagnosticsPublisher.Comparer.cs

@@ -26,9 +26,6 @@ public bool Equals(IDocumentSnapshot? x, IDocumentSnapshot? y)
         }
 
         public int GetHashCode(IDocumentSnapshot obj)
-        {
-            var filePath = obj.FilePath.AssumeNotNull();
-            return FilePathComparer.Instance.GetHashCode(filePath);
-        }
+            => FilePathComparer.Instance.GetHashCode(obj.FilePath);
     }
 }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Diagnostics/RazorDiagnosticsPublisher.cs

@@ -121,7 +121,7 @@ private async Task PublishDiagnosticsAsync(IDocumentSnapshot document, Cancellat
 
         lock (_publishedDiagnostics)
         {
-            var filePath = document.FilePath.AssumeNotNull();
+            var filePath = document.FilePath;
 
             // See if these are the same diagnostics as last time. If so, we don't need to publish.
             if (_publishedDiagnostics.TryGetValue(filePath, out var previousDiagnostics))

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/GeneratedDocumentSynchronizer.cs

@@ -19,18 +19,18 @@ internal class GeneratedDocumentSynchronizer(
     public void DocumentProcessed(RazorCodeDocument codeDocument, IDocumentSnapshot document)
     {
         var hostDocumentVersion = document.Version;
-        var filePath = document.FilePath.AssumeNotNull();
+        var filePath = document.FilePath;
 
         // If the document isn't open, and we're not updating buffers for closed documents, then we don't need to do anything.
-        if (!_projectManager.IsDocumentOpen(document.FilePath) &&
+        if (!_projectManager.IsDocumentOpen(filePath) &&
             !_languageServerFeatureOptions.UpdateBuffersForClosedDocuments)
         {
             return;
         }
 
         // If the document has been removed from the project, then don't do anything, or version numbers will be thrown off
         if (!_projectManager.TryGetLoadedProject(document.Project.Key, out var project) ||
-            !project.ContainsDocument(document.FilePath))
+            !project.ContainsDocument(filePath))
         {
             return;
         }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/OpenDocumentGenerator.Comparer.cs

@@ -37,7 +37,7 @@ public int GetHashCode(IDocumentSnapshot obj)
         {
             var hash = HashCodeCombiner.Start();
             hash.Add(obj.Project.Key.Id, FilePathComparer.Instance);
-            hash.Add(obj.FileKind.AssumeNotNull(), FilePathComparer.Instance);
+            hash.Add(obj.FileKind, FilePathComparer.Instance);
             return hash.CombinedHash;
         }
     }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/OpenDocumentGenerator.cs

@@ -148,7 +148,7 @@ private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
                     {
                         foreach (var relatedDocument in oldProject.GetRelatedDocuments(document))
                         {
-                            var relatedDocumentFilePath = relatedDocument.FilePath.AssumeNotNull();
+                            var relatedDocumentFilePath = relatedDocument.FilePath;
 
                             if (newProject.TryGetDocument(relatedDocumentFilePath, out var newRelatedDocument))
                             {
@@ -169,7 +169,7 @@ private void ProjectManager_Changed(object? sender, ProjectChangeEventArgs args)
 
         void EnqueueIfNecessary(IDocumentSnapshot document)
         {
-            if (!_projectManager.IsDocumentOpen(document.FilePath.AssumeNotNull()) &&
+            if (!_projectManager.IsDocumentOpen(document.FilePath) &&
                 !_options.UpdateBuffersForClosedDocuments)
             {
                 return;

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/GoToDefinition/AbstractRazorComponentDefinitionService.cs

@@ -3,7 +3,6 @@
 
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.AspNetCore.Razor;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.Logging;
@@ -62,7 +61,7 @@ internal abstract class AbstractRazorComponentDefinitionService(
             return null;
         }
 
-        var componentFilePath = componentDocument.FilePath.AssumeNotNull();
+        var componentFilePath = componentDocument.FilePath;
 
         _logger.LogInformation($"Definition found at file path: {componentFilePath}");
 

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

@@ -23,8 +23,8 @@ internal class DocumentContext(Uri uri, IDocumentSnapshot snapshot, VSProjectCon
 
     public Uri Uri { get; } = uri;
     public IDocumentSnapshot Snapshot { get; } = snapshot;
-    public string FilePath => Snapshot.FilePath.AssumeNotNull();
-    public string FileKind => Snapshot.FileKind.AssumeNotNull();
+    public string FilePath => Snapshot.FilePath;
+    public string FileKind => Snapshot.FileKind;
     public IProjectSnapshot Project => Snapshot.Project;
 
     public TextDocumentIdentifier GetTextDocumentIdentifier()

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

@@ -82,15 +82,8 @@ public async ValueTask<RazorCodeDocument> GetGeneratedOutputAsync(bool forceDesi
         return output;
     }
 
-    private async Task<RazorCodeDocument> GetDesignTimeGeneratedOutputAsync(CancellationToken cancellationToken)
-    {
-        var tagHelpers = await Project.GetTagHelpersAsync(cancellationToken).ConfigureAwait(false);
-        var projectEngine = Project.GetProjectEngine();
-        var imports = await DocumentState.GetImportsAsync(this, projectEngine, cancellationToken).ConfigureAwait(false);
-        return await DocumentState
-            .GenerateCodeDocumentAsync(this, projectEngine, imports, tagHelpers, forceRuntimeCodeGeneration: false, cancellationToken)
-            .ConfigureAwait(false);
-    }
+    private Task<RazorCodeDocument> GetDesignTimeGeneratedOutputAsync(CancellationToken cancellationToken)
+        => DocumentState.GenerateCodeDocumentAsync(this, Project.GetProjectEngine(), forceRuntimeCodeGeneration: false, cancellationToken);
 
     /// <summary>
     ///  Retrieves a cached Roslyn <see cref="SyntaxTree"/> from the generated C# document.

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

@@ -186,7 +186,7 @@ static void PropagateToTaskCompletionSource(
             var configurationVersion = project.ConfigurationVersion;
             var projectWorkspaceStateVersion = project.ProjectWorkspaceStateVersion;
             var documentCollectionVersion = project.DocumentCollectionVersion;
-            var imports = await GetImportsAsync(document, project.GetProjectEngine(), cancellationToken).ConfigureAwait(false);
+            var importItems = await GetImportItemsAsync(document, project.GetProjectEngine(), cancellationToken).ConfigureAwait(false);
             var documentVersion = await document.GetTextVersionAsync(cancellationToken).ConfigureAwait(false);
 
             // OK now that have the previous output and all of the versions, we can see if anything
@@ -207,7 +207,7 @@ static void PropagateToTaskCompletionSource(
                 inputVersion = documentCollectionVersion;
             }
 
-            foreach (var import in imports)
+            foreach (var import in importItems)
             {
                 var importVersion = import.Version;
                 if (inputVersion.GetNewerVersion(importVersion) == importVersion)
@@ -232,9 +232,8 @@ static void PropagateToTaskCompletionSource(
                 }
             }
 
-            var tagHelpers = await project.GetTagHelpersAsync(cancellationToken).ConfigureAwait(false);
             var forceRuntimeCodeGeneration = project.LanguageServerFeatureOptions.ForceRuntimeCodeGeneration;
-            var codeDocument = await GenerateCodeDocumentAsync(document, project.GetProjectEngine(), imports, tagHelpers, forceRuntimeCodeGeneration, cancellationToken).ConfigureAwait(false);
+            var codeDocument = await GenerateCodeDocumentAsync(document, project.GetProjectEngine(), importItems, forceRuntimeCodeGeneration, cancellationToken).ConfigureAwait(false);
             return (codeDocument, inputVersion);
         }
 

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

@@ -1,12 +1,20 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
+using Microsoft.CodeAnalysis.Text;
+
 namespace Microsoft.CodeAnalysis.Razor.ProjectSystem;
 
 internal partial class DocumentState
 {
-    internal record struct ImportItem(string? FilePath, VersionStamp Version, IDocumentSnapshot Document)
+    private readonly record struct ImportItem(
+        string? FilePath,
+        string? FileKind,
+        SourceText Text,
+        VersionStamp Version)
     {
-        public readonly string? FileKind => Document.FileKind;
+        // Note: The default import does not have file path, file kind, or version stamp.
+        public static ImportItem CreateDefault(SourceText text)
+            => new(FilePath: null, FileKind: null, text, Version: default);
     }
 }

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

@@ -193,98 +193,113 @@ public virtual DocumentState WithTextLoader(TextLoader textLoader)
         return new DocumentState(HostDocument, Version + 1, textAndVersion: null, textLoader);
     }
 
-    // Internal, because we are temporarily sharing code with CohostDocumentSnapshot
-    internal static ImmutableArray<IDocumentSnapshot> GetImportsCore(IProjectSnapshot project, RazorProjectEngine projectEngine, string filePath, string fileKind)
+    internal static async Task<RazorCodeDocument> GenerateCodeDocumentAsync(
+        IDocumentSnapshot document,
+        RazorProjectEngine projectEngine,
+        bool forceRuntimeCodeGeneration,
+        CancellationToken cancellationToken)
     {
-        var projectItem = projectEngine.FileSystem.GetItem(filePath, fileKind);
+        var importItems = await GetImportItemsAsync(document, projectEngine, cancellationToken).ConfigureAwait(false);
 
-        using var importItems = new PooledArrayBuilder<RazorProjectItem>();
+        return await GenerateCodeDocumentAsync(
+            document, projectEngine, importItems, forceRuntimeCodeGeneration, cancellationToken).ConfigureAwait(false);
+    }
+
+    private static async Task<RazorCodeDocument> GenerateCodeDocumentAsync(
+        IDocumentSnapshot document,
+        RazorProjectEngine projectEngine,
+        ImmutableArray<ImportItem> imports,
+        bool forceRuntimeCodeGeneration,
+        CancellationToken cancellationToken)
+    {
+        var importSources = GetImportSources(imports, projectEngine);
+        var tagHelpers = await document.Project.GetTagHelpersAsync(cancellationToken).ConfigureAwait(false);
+        var source = await GetSourceAsync(document, projectEngine, cancellationToken).ConfigureAwait(false);
+
+        return forceRuntimeCodeGeneration
+            ? projectEngine.Process(source, document.FileKind, importSources, tagHelpers)
+            : projectEngine.ProcessDesignTime(source, document.FileKind, importSources, tagHelpers);
+    }
+
+    private static async Task<ImmutableArray<ImportItem>> GetImportItemsAsync(IDocumentSnapshot document, RazorProjectEngine projectEngine, CancellationToken cancellationToken)
+    {
+        var projectItem = projectEngine.FileSystem.GetItem(document.FilePath, document.FileKind);
+
+        using var importProjectItems = new PooledArrayBuilder<RazorProjectItem>();
 
         foreach (var feature in projectEngine.ProjectFeatures.OfType<IImportProjectFeature>())
         {
             if (feature.GetImports(projectItem) is { } featureImports)
             {
-                importItems.AddRange(featureImports);
+                importProjectItems.AddRange(featureImports);
             }
         }
 
-        if (importItems.Count == 0)
+        if (importProjectItems.Count == 0)
         {
             return [];
         }
 
-        using var imports = new PooledArrayBuilder<IDocumentSnapshot>(capacity: importItems.Count);
+        var project = document.Project;
+
+        using var importItems = new PooledArrayBuilder<ImportItem>(capacity: importProjectItems.Count);
 
-        foreach (var item in importItems)
+        foreach (var importProjectItem in importProjectItems)
         {
-            if (item is NotFoundProjectItem)
+            if (importProjectItem is NotFoundProjectItem)
             {
                 continue;
             }
 
-            if (item.PhysicalPath is null)
+            if (importProjectItem.PhysicalPath is null)
             {
                 // This is a default import.
-                var defaultImport = new ImportDocumentSnapshot(project, item);
-                imports.Add(defaultImport);
+                using var stream = importProjectItem.Read();
+                var text = SourceText.From(stream);
+                var defaultImport = ImportItem.CreateDefault(text);
+
+                importItems.Add(defaultImport);
             }
-            else if (project.TryGetDocument(item.PhysicalPath, out var import))
+            else if (project.TryGetDocument(importProjectItem.PhysicalPath, out var importDocument))
             {
-                imports.Add(import);
+                var text = await importDocument.GetTextAsync(cancellationToken).ConfigureAwait(false);
+                var versionStamp = await importDocument.GetTextVersionAsync(cancellationToken).ConfigureAwait(false);
+                var importItem = new ImportItem(importDocument.FilePath, importDocument.FileKind, text, versionStamp);
+
+                importItems.Add(importItem);
             }
         }
 
-        return imports.DrainToImmutable();
+        return importItems.DrainToImmutable();
     }
 
-    internal static async Task<RazorCodeDocument> GenerateCodeDocumentAsync(
-        IDocumentSnapshot document,
-        RazorProjectEngine projectEngine,
-        ImmutableArray<ImportItem> imports,
-        ImmutableArray<TagHelperDescriptor> tagHelpers,
-        bool forceRuntimeCodeGeneration,
-        CancellationToken cancellationToken)
+    private static ImmutableArray<RazorSourceDocument> GetImportSources(ImmutableArray<ImportItem> importItems, RazorProjectEngine projectEngine)
     {
-        // OK we have to generate the code.
-        using var importSources = new PooledArrayBuilder<RazorSourceDocument>(imports.Length);
-        foreach (var item in imports)
-        {
-            var importProjectItem = item.FilePath is null ? null : projectEngine.FileSystem.GetItem(item.FilePath, item.FileKind);
-            var sourceDocument = await GetRazorSourceDocumentAsync(item.Document, importProjectItem, cancellationToken).ConfigureAwait(false);
-            importSources.Add(sourceDocument);
-        }
+        using var importSources = new PooledArrayBuilder<RazorSourceDocument>(importItems.Length);
 
-        var projectItem = document.FilePath is null ? null : projectEngine.FileSystem.GetItem(document.FilePath, document.FileKind);
-        var documentSource = await GetRazorSourceDocumentAsync(document, projectItem, cancellationToken).ConfigureAwait(false);
-
-        if (forceRuntimeCodeGeneration)
+        foreach (var importItem in importItems)
         {
-            return projectEngine.Process(documentSource, fileKind: document.FileKind, importSources.DrainToImmutable(), tagHelpers);
-        }
+            var importProjectItem = importItem is { FilePath: string filePath, FileKind: var fileKind }
+                ? projectEngine.FileSystem.GetItem(filePath, fileKind)
+                : null;
 
-        return projectEngine.ProcessDesignTime(documentSource, fileKind: document.FileKind, importSources.DrainToImmutable(), tagHelpers);
-    }
-
-    internal static async Task<ImmutableArray<ImportItem>> GetImportsAsync(IDocumentSnapshot document, RazorProjectEngine projectEngine, CancellationToken cancellationToken)
-    {
-        var imports = GetImportsCore(document.Project, projectEngine, document.FilePath.AssumeNotNull(), document.FileKind.AssumeNotNull());
-        using var result = new PooledArrayBuilder<ImportItem>(imports.Length);
+            var properties = RazorSourceDocumentProperties.Create(importItem.FilePath, importProjectItem?.RelativePhysicalPath);
+            var importSource = RazorSourceDocument.Create(importItem.Text, properties);
 
-        foreach (var snapshot in imports)
-        {
-            var versionStamp = await snapshot.GetTextVersionAsync(cancellationToken).ConfigureAwait(false);
-            result.Add(new ImportItem(snapshot.FilePath, versionStamp, snapshot));
+            importSources.Add(importSource);
         }
 
-        return result.DrainToImmutable();
+        return importSources.DrainToImmutable();
     }
 
-    private static async Task<RazorSourceDocument> GetRazorSourceDocumentAsync(
-        IDocumentSnapshot document,
-        RazorProjectItem? projectItem,
-        CancellationToken cancellationToken)
+    private static async Task<RazorSourceDocument> GetSourceAsync(IDocumentSnapshot document, RazorProjectEngine projectEngine, CancellationToken cancellationToken)
     {
-        var sourceText = await document.GetTextAsync(cancellationToken).ConfigureAwait(false);
-        return RazorSourceDocument.Create(sourceText, RazorSourceDocumentProperties.Create(document.FilePath, projectItem?.RelativePhysicalPath));
+        var projectItem = document is { FilePath: string filePath, FileKind: var fileKind }
+            ? projectEngine.FileSystem.GetItem(filePath, fileKind)
+            : null;
+
+        var text = await document.GetTextAsync(cancellationToken).ConfigureAwait(false);
+        var properties = RazorSourceDocumentProperties.Create(document.FilePath, projectItem?.RelativePhysicalPath);
+        return RazorSourceDocument.Create(text, properties);
     }
 }

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

@@ -6,7 +6,6 @@
 #endif
 
 using System.Diagnostics;
-using Microsoft.AspNetCore.Razor;
 using Microsoft.AspNetCore.Razor.ProjectSystem;
 using Microsoft.AspNetCore.Razor.Serialization;
 using Microsoft.AspNetCore.Razor.Utilities;
@@ -16,7 +15,7 @@ namespace Microsoft.CodeAnalysis.Razor.ProjectSystem;
 internal static class Extensions
 {
     public static DocumentSnapshotHandle ToHandle(this IDocumentSnapshot snapshot)
-        => new(snapshot.FilePath.AssumeNotNull(), snapshot.TargetPath.AssumeNotNull(), snapshot.FileKind.AssumeNotNull());
+        => new(snapshot.FilePath, snapshot.TargetPath, snapshot.FileKind);
 
     public static ProjectKey ToProjectKey(this Project project)
     {

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

@@ -11,9 +11,9 @@ namespace Microsoft.CodeAnalysis.Razor.ProjectSystem;
 
 internal interface IDocumentSnapshot
 {
-    string? FileKind { get; }
-    string? FilePath { get; }
-    string? TargetPath { get; }
+    string FileKind { get; }
+    string FilePath { get; }
+    string TargetPath { get; }
     IProjectSnapshot Project { get; }
 
     int Version { get; }