Enable support for an LSP client to open source generated files

src/EditorFeatures/TestUtilities/LanguageServer/AbstractLanguageServerProtocolTests.cs

@@ -427,6 +427,20 @@ protected static void AddMappedDocument(Workspace workspace, string markup)
             workspace.TryApplyChanges(newSolution);
         }
 
+        protected static async Task AddGeneratorAsync(ISourceGenerator generator, EditorTestWorkspace workspace)
+        {
+            var analyzerReference = new TestGeneratorReference(generator);
+
+            var solution = workspace.CurrentSolution
+                .Projects.Single()
+                .AddAnalyzerReference(analyzerReference)
+                .Solution;
+
+            await workspace.ChangeSolutionAsync(solution);
+            await WaitForWorkspaceOperationsAsync(workspace);
+
+        }
+
         internal static async Task<Dictionary<string, IList<LSP.Location>>> GetAnnotatedLocationsAsync(EditorTestWorkspace workspace, Solution solution)
         {
             var locations = new Dictionary<string, IList<LSP.Location>>();
@@ -620,6 +634,19 @@ private static RoslynLanguageServer CreateLanguageServer(Stream inputStream, Str
                 return languageServer;
             }
 
+            public async Task<Document> GetDocumentAsync(Uri uri)
+            {
+                var document = await GetCurrentSolution().GetDocumentAsync(new LSP.TextDocumentIdentifier { Uri = uri }, CancellationToken.None).ConfigureAwait(false);
+                Contract.ThrowIfNull(document, $"Unable to find document with {uri} in solution");
+                return document;
+            }
+
+            public async Task<SourceText> GetDocumentTextAsync(Uri uri)
+            {
+                var document = await GetDocumentAsync(uri).ConfigureAwait(false);
+                return await document.GetTextAsync(CancellationToken.None).ConfigureAwait(false);
+            }
+
             public async Task<ResponseType?> ExecuteRequestAsync<RequestType, ResponseType>(string methodName, RequestType request, CancellationToken cancellationToken) where RequestType : class
             {
                 // If creating the LanguageServer threw we might timeout without this.
@@ -655,7 +682,7 @@ public async Task OpenDocumentAsync(Uri documentUri, string? text = null, string
                 {
                     // LSP open files don't care about the project context, just the file contents with the URI.
                     // So pick any of the linked documents to get the text from.
-                    var sourceText = await TestWorkspace.CurrentSolution.GetDocuments(documentUri).First().GetTextAsync(CancellationToken.None).ConfigureAwait(false);
+                    var sourceText = await GetDocumentTextAsync(documentUri).ConfigureAwait(false);
                     text = sourceText.ToString();
                 }
 

src/Features/Core/Portable/Navigation/INavigableItem.cs

@@ -48,7 +48,7 @@ internal interface INavigableItem
 
     ImmutableArray<INavigableItem> ChildItems { get; }
 
-    public record NavigableDocument(NavigableProject Project, string Name, string? FilePath, IReadOnlyList<string> Folders, DocumentId Id, bool IsSourceGeneratedDocument, Workspace? Workspace)
+    public record NavigableDocument(NavigableProject Project, string Name, string? FilePath, IReadOnlyList<string> Folders, DocumentId Id, SourceGeneratedDocumentIdentity? SourceGeneratedDocumentIdentity, Workspace? Workspace)
     {
         public static NavigableDocument FromDocument(Document document)
             => new(
@@ -57,7 +57,7 @@ public static NavigableDocument FromDocument(Document document)
                 document.FilePath,
                 document.Folders,
                 document.Id,
-                IsSourceGeneratedDocument: document is SourceGeneratedDocument,
+                SourceGeneratedDocumentIdentity: (document as SourceGeneratedDocument)?.Identity,
                 document.Project.Solution.TryGetWorkspace());
 
         /// <summary>
@@ -66,7 +66,7 @@ public static NavigableDocument FromDocument(Document document)
         /// navigable item was constructed during a Find Symbols operation on the same solution instance.
         /// </summary>
         internal ValueTask<Document> GetRequiredDocumentAsync(Solution solution, CancellationToken cancellationToken)
-            => solution.GetRequiredDocumentAsync(Id, includeSourceGenerated: IsSourceGeneratedDocument, cancellationToken);
+            => solution.GetRequiredDocumentAsync(Id, includeSourceGenerated: SourceGeneratedDocumentIdentity is not null, cancellationToken);
 
         /// <summary>
         /// Get the <see cref="SourceText"/> of the <see cref="CodeAnalysis.Document"/> within

src/Features/Lsif/Generator/Generator.cs

@@ -210,9 +210,9 @@ public async Task GenerateForProjectAsync(
             // use this document can benefit from that single shared model.
             var semanticModel = await document.GetRequiredSemanticModelAsync(cancellationToken);
 
-            var (uri, contentBase64Encoded) = await GetUriAndContentAsync(document, cancellationToken);
+            var contentBase64Encoded = await GetBase64EncodedContentAsync(document, cancellationToken);
 
-            var documentVertex = new Graph.LsifDocument(uri, GetLanguageKind(semanticModel.Language), contentBase64Encoded, idFactory);
+            var documentVertex = new Graph.LsifDocument(document.GetURI(), GetLanguageKind(semanticModel.Language), contentBase64Encoded, idFactory);
             lsifJsonWriter.Write(documentVertex);
             lsifJsonWriter.Write(new Event(Event.EventKind.Begin, documentVertex.GetId(), idFactory));
 
@@ -443,32 +443,21 @@ private static (DefinitionRangeTag tag, TextSpan fullRange)? CreateRangeTagAndCo
             return (new DefinitionRangeTag(syntaxToken.Text, symbolKind, fullRange), fullRangeSpan);
         }
 
-        private static async Task<(Uri uri, string? contentBase64Encoded)> GetUriAndContentAsync(
+        private static async Task<string?> GetBase64EncodedContentAsync(
             Document document, CancellationToken cancellationToken)
         {
-            Contract.ThrowIfNull(document.FilePath);
-
-            string? contentBase64Encoded = null;
-            Uri uri;
-
             if (document is SourceGeneratedDocument)
             {
                 var text = await document.GetValueTextAsync(cancellationToken);
 
                 // We always use UTF-8 encoding when writing out file contents, as that's expected by LSIF implementations.
                 // TODO: when we move to .NET Core, is there a way to reduce allocations here?
-                contentBase64Encoded = Convert.ToBase64String(Encoding.UTF8.GetBytes(text.ToString()));
-
-                // There is a triple slash here, so the "host" portion of the URI is empty, similar to
-                // how file URIs work.
-                uri = ProtocolConversions.CreateUriFromSourceGeneratedFilePath(document.FilePath);
+                return Convert.ToBase64String(Encoding.UTF8.GetBytes(text.ToString()));
             }
             else
             {
-                uri = ProtocolConversions.CreateAbsoluteUri(document.FilePath);
+                return null;
             }
-
-            return (uri, contentBase64Encoded);
         }
 
         private static async Task GenerateSemanticTokensAsync(

src/LanguageServer/BannedSymbols.txt

@@ -11,13 +11,13 @@ T:System.ComponentModel.Composition.PartNotDiscoverableAttribute; Use types from
 T:System.ComponentModel.Composition.SharedAttribute; Use types from System.Composition instead
 T:System.ComponentModel.Composition.SharingBoundaryAttribute; Use types from System.Composition instead
 T:System.ComponentModel.Composition.Convention.AttributedModelProvider; Use types from System.Composition instead
-M:System.Uri.#ctor(System.String); Use ProtocolConversions.CreateAbsoluteUri or ProtocolConversions.CreateUriFromSourceGeneratedFilePath
-M:System.Uri.#ctor(System.String,System.Boolean); Use ProtocolConversions.CreateAbsoluteUri or ProtocolConversions.CreateUriFromSourceGeneratedFilePath
-M:System.Uri.#ctor(System.String,System.UriCreationOptions); Use ProtocolConversions.CreateAbsoluteUri or ProtocolConversions.CreateUriFromSourceGeneratedFilePath
-M:System.Uri.#ctor(System.String,System.UriKind); Use ProtocolConversions.CreateAbsoluteUri or ProtocolConversions.CreateUriFromSourceGeneratedFilePath
-M:System.Uri.#ctor(System.Uri,System.String); Use ProtocolConversions.CreateAbsoluteUri or ProtocolConversions.CreateUriFromSourceGeneratedFilePath
-M:System.Uri.#ctor(System.Uri,System.Uri); Use ProtocolConversions.CreateAbsoluteUri or ProtocolConversions.CreateUriFromSourceGeneratedFilePath
-M:System.Uri.#ctor(System.Uri,System.String,System.Boolean); Use ProtocolConversions.CreateAbsoluteUri or ProtocolConversions.CreateUriFromSourceGeneratedFilePath
-M:System.Uri.TryCreate(System.String,System.UriKind,System.Uri@); Use ProtocolConversions.CreateAbsoluteUri or ProtocolConversions.CreateUriFromSourceGeneratedFilePath
-M:System.Uri.TryCreate(System.Uri,System.String,System.Uri@); Use ProtocolConversions.CreateAbsoluteUri or ProtocolConversions.CreateUriFromSourceGeneratedFilePath
-M:System.Uri.TryCreate(System.Uri,System.Uri,System.Uri@); Use ProtocolConversions.CreateAbsoluteUri or ProtocolConversions.CreateUriFromSourceGeneratedFilePath
+M:System.Uri.#ctor(System.String); Use Extensions.GetURI() if you have a document, or ProtocolConversions.CreateAbsoluteUri if you are sure the path is a real file path and not a generated file path
+M:System.Uri.#ctor(System.String,System.Boolean); Use Extensions.GetURI() if you have a document, or ProtocolConversions.CreateAbsoluteUri if you are sure the path is a real file path and not a generated file path
+M:System.Uri.#ctor(System.String,System.UriCreationOptions); Use Extensions.GetURI() if you have a document, or ProtocolConversions.CreateAbsoluteUri if you are sure the path is a real file path and not a generated file path
+M:System.Uri.#ctor(System.String,System.UriKind); Use Extensions.GetURI() if you have a document, or ProtocolConversions.CreateAbsoluteUri if you are sure the path is a real file path and not a generated file path
+M:System.Uri.#ctor(System.Uri,System.String); Use Extensions.GetURI() if you have a document, or ProtocolConversions.CreateAbsoluteUri if you are sure the path is a real file path and not a generated file path
+M:System.Uri.#ctor(System.Uri,System.Uri); Use Extensions.GetURI() if you have a document, or ProtocolConversions.CreateAbsoluteUri if you are sure the path is a real file path and not a generated file path
+M:System.Uri.#ctor(System.Uri,System.String,System.Boolean); Use Extensions.GetURI() if you have a document, or ProtocolConversions.CreateAbsoluteUri if you are sure the path is a real file path and not a generated file path
+M:System.Uri.TryCreate(System.String,System.UriKind,System.Uri@); Use Extensions.GetURI() if you have a document, or ProtocolConversions.CreateAbsoluteUri if you are sure the path is a real file path and not a generated file path
+M:System.Uri.TryCreate(System.Uri,System.String,System.Uri@); Use Extensions.GetURI() if you have a document, or ProtocolConversions.CreateAbsoluteUri if you are sure the path is a real file path and not a generated file path
+M:System.Uri.TryCreate(System.Uri,System.Uri,System.Uri@); Use Extensions.GetURI() if you have a document, or ProtocolConversions.CreateAbsoluteUri if you are sure the path is a real file path and not a generated file path

src/LanguageServer/Protocol/Extensions/Extensions.cs

@@ -24,8 +24,8 @@ internal static class Extensions
         public static Uri GetURI(this TextDocument document)
         {
             Contract.ThrowIfNull(document.FilePath);
-            return document is SourceGeneratedDocument
-                ? ProtocolConversions.CreateUriFromSourceGeneratedFilePath(document.FilePath)
+            return document is SourceGeneratedDocument sourceGeneratedDocument
+                ? SourceGeneratedDocumentUri.Create(sourceGeneratedDocument.Identity)
                 : ProtocolConversions.CreateAbsoluteUri(document.FilePath);
         }
 
@@ -56,42 +56,63 @@ public static Uri CreateUriForDocumentWithoutFilePath(this TextDocument document
             return ProtocolConversions.CreateAbsoluteUri(path);
         }
 
-        public static ImmutableArray<Document> GetDocuments(this Solution solution, Uri documentUri)
-            => GetDocuments(solution, ProtocolConversions.GetDocumentFilePathFromUri(documentUri));
-
-        public static ImmutableArray<Document> GetDocuments(this Solution solution, string documentPath)
-        {
-            var documentIds = solution.GetDocumentIdsWithFilePath(documentPath);
-
-            // We don't call GetRequiredDocument here as the id could be referring to an additional document.
-            var documents = documentIds.Select(solution.GetDocument).WhereNotNull().ToImmutableArray();
-            return documents;
-        }
-
         /// <summary>
         /// Get all regular and additional <see cref="TextDocument"/>s for the given <paramref name="documentUri"/>.
+        /// This will not return source generated documents.
         /// </summary>
         public static ImmutableArray<TextDocument> GetTextDocuments(this Solution solution, Uri documentUri)
         {
             var documentIds = GetDocumentIds(solution, documentUri);
 
             var documents = documentIds
-                .Select(solution.GetDocument)
-                .Concat(documentIds.Select(solution.GetAdditionalDocument))
+                .Select(solution.GetTextDocument)
                 .WhereNotNull()
                 .ToImmutableArray();
             return documents;
         }
 
         public static ImmutableArray<DocumentId> GetDocumentIds(this Solution solution, Uri documentUri)
-            => solution.GetDocumentIdsWithFilePath(ProtocolConversions.GetDocumentFilePathFromUri(documentUri));
+        {
+            // If this is not our special scheme for generated documents, then we can just look for documents with that file path.
+            if (documentUri.Scheme != SourceGeneratedDocumentUri.Scheme)
+                return solution.GetDocumentIdsWithFilePath(ProtocolConversions.GetDocumentFilePathFromUri(documentUri));
+
+            // We can get a null documentId if we were unable to find the project associated with the
+            // generated document - this can happen if say a project is unloaded.  There may be LSP requests
+            // already in-flight which may ask for a generated document from that project.  So we return null
+            var documentId = SourceGeneratedDocumentUri.DeserializeIdentity(solution, documentUri)?.DocumentId;
+
+            return documentId is not null ? [documentId] : [];
+        }
 
-        public static Document? GetDocument(this Solution solution, TextDocumentIdentifier documentIdentifier)
+        /// <summary>
+        /// Finds the document for a TextDocumentIdentifier, potentially returning a source-generated file.
+        /// </summary>
+        public static async ValueTask<Document?> GetDocumentAsync(this Solution solution, TextDocumentIdentifier documentIdentifier, CancellationToken cancellationToken)
+        {
+            var textDocument = await solution.GetTextDocumentAsync(documentIdentifier, cancellationToken).ConfigureAwait(false);
+            Contract.ThrowIfTrue(textDocument is not null && textDocument is not Document, $"{textDocument!.Id} is not a Document");
+            return textDocument as Document;
+        }
+
+        /// <summary>
+        /// Finds the TextDocument for a TextDocumentIdentifier, potentially returning a source-generated file.
+        /// </summary>
+        public static async ValueTask<TextDocument?> GetTextDocumentAsync(this Solution solution, TextDocumentIdentifier documentIdentifier, CancellationToken cancellationToken)
         {
-            var documents = solution.GetDocuments(documentIdentifier.Uri);
+            // If it's the URI scheme for source generated files, delegate to our other helper, otherwise we can handle anything else here.
+            if (documentIdentifier.Uri.Scheme == SourceGeneratedDocumentUri.Scheme)
+            {
+                // In the case of a URI scheme for source generated files, we generate a different URI for each project, thus this URI cannot be linked into multiple projects;
+                // this means we can safely call .SingleOrDefault() and not worry about calling FindDocumentInProjectContext.
+                var documentId = solution.GetDocumentIds(documentIdentifier.Uri).SingleOrDefault();
+                return await solution.GetDocumentAsync(documentId, includeSourceGenerated: true, cancellationToken).ConfigureAwait(false);
+            }
+
+            var documents = solution.GetTextDocuments(documentIdentifier.Uri);
             return documents.Length == 0
                 ? null
-                : documents.FindDocumentInProjectContext(documentIdentifier, (sln, id) => sln.GetRequiredDocument(id));
+                : documents.FindDocumentInProjectContext(documentIdentifier, (sln, id) => sln.GetRequiredTextDocument(id));
         }
 
         private static T FindItemInProjectContext<T>(

src/LanguageServer/Protocol/Extensions/ProtocolConversions.cs

@@ -35,14 +35,9 @@ internal static partial class ProtocolConversions
     {
         private const string CSharpMarkdownLanguageName = "csharp";
         private const string VisualBasicMarkdownLanguageName = "vb";
-        private const string SourceGeneratedDocumentBaseUri = "source-generated:///";
         private const string BlockCodeFence = "```";
         private const string InlineCodeFence = "`";
 
-#pragma warning disable RS0030 // Do not use banned APIs
-        private static readonly Uri s_sourceGeneratedDocumentBaseUri = new(SourceGeneratedDocumentBaseUri, UriKind.Absolute);
-#pragma warning restore
-
         private static readonly char[] s_dirSeparators = [PathUtilities.DirectorySeparatorChar, PathUtilities.AltDirectorySeparatorChar];
 
         private static readonly Regex s_markdownEscapeRegex = new(@"([\\`\*_\{\}\[\]\(\)#+\-\.!])", RegexOptions.Compiled);
@@ -188,7 +183,9 @@ static async Task<char> GetInsertionCharacterAsync(Document document, int positi
         }
 
         public static string GetDocumentFilePathFromUri(Uri uri)
-            => uri.IsFile ? uri.LocalPath : uri.AbsoluteUri;
+        {
+            return uri.IsFile ? uri.LocalPath : uri.AbsoluteUri;
+        }
 
         /// <summary>
         /// Converts an absolute local file path or an absolute URL string to <see cref="Uri"/>.
@@ -262,28 +259,6 @@ static string EscapeUriPart(string stringToEscape)
 #pragma warning restore
         }
 
-        public static Uri CreateUriFromSourceGeneratedFilePath(string filePath)
-        {
-            Debug.Assert(!PathUtilities.IsAbsolute(filePath));
-
-            // Fast path for common cases:
-            if (IsAscii(filePath))
-            {
-#pragma warning disable RS0030 // Do not use banned APIs
-                return new Uri(s_sourceGeneratedDocumentBaseUri, filePath);
-#pragma warning restore
-            }
-
-            // Workaround for https://github.com/dotnet/runtime/issues/89538:
-
-            var parts = filePath.Split(s_dirSeparators);
-            var url = SourceGeneratedDocumentBaseUri + string.Join("/", parts.Select(Uri.EscapeDataString));
-
-#pragma warning disable RS0030 // Do not use banned APIs
-            return new Uri(url, UriKind.Absolute);
-#pragma warning restore
-        }
-
         private static bool IsAscii(char c)
             => (uint)c <= '\x007f';
 
@@ -512,13 +487,12 @@ public static LSP.Range TextSpanToRange(TextSpan textSpan, SourceText text)
                 Range = MappedSpanResultToRange(mappedSpan)
             };
 
-            static async Task<LSP.Location?> ConvertTextSpanToLocationAsync(
+            static async Task<LSP.Location> ConvertTextSpanToLocationAsync(
                 TextDocument document,
                 TextSpan span,
                 bool isStale,
                 CancellationToken cancellationToken)
             {
-                Debug.Assert(document.FilePath != null);
                 var uri = document.GetURI();
 
                 var text = await document.GetValueTextAsync(cancellationToken).ConfigureAwait(false);