Add support for GoToDefinition on source-generated files

This adds a new response element to gotodefinition responses: SourceGeneratedFileInfo. This is similar to MetadataSource, except that unlike MetadataSource it's not tracked on a type/project basis, but rather as a document/project basis. Retrieving info about a source generated file can be done through the SourceGeneratedFileService endpoints:

* SourceGeneratedFileInfo - Gets the file content of a source generated file.
* UpdateSourceGeneratedFileInfo - Gets the updated content of a source generated file, if it has changed since the last time information was returned.
* SourceGeneratedFileClosed - Sent to the server to inform it that the editor has closed the generated file and it can stop tracking Document version info for that file.

Currently, the only way to get the info needed to open a source-generated file is via the gotodefinition endpoint. We'll want to add info to find-usages as well, but that's a job for another day. Contributes to OmniSharp#1934.

src/OmniSharp.Abstractions/Models/v1/GotoDefinition/GotoDefinitionResponse.cs

@@ -1,16 +1,20 @@
+#nullable enable
+
 using Newtonsoft.Json;
 using OmniSharp.Models.Metadata;
+using OmniSharp.Models.v1.SourceGeneratedFile;
 
 namespace OmniSharp.Models.GotoDefinition
 {
     public class GotoDefinitionResponse : ICanBeEmptyResponse
     {
-        public string FileName { get; set; }
+        public string? FileName { get; set; }
         [JsonConverter(typeof(ZeroBasedIndexConverter))]
         public int Line { get; set; }
         [JsonConverter(typeof(ZeroBasedIndexConverter))]
         public int Column { get; set; }
-        public MetadataSource MetadataSource { get; set; }
+        public MetadataSource? MetadataSource { get; set; }
+        public SourceGeneratedFileInfo? SourceGeneratedInfo { get; set; }
         public bool IsEmpty => string.IsNullOrWhiteSpace(FileName) && MetadataSource == null;
     }
 }

src/OmniSharp.Abstractions/Models/v1/SourceGeneratedFile/SourceGeneratedFileClosedRequest.cs

@@ -0,0 +1,11 @@
+#nullable enable
+
+using OmniSharp.Mef;
+
+namespace OmniSharp.Models.v1.SourceGeneratedFile
+{
+    [OmniSharpEndpoint(OmniSharpEndpoints.SourceGeneratedFileClosed, typeof(SourceGeneratedFileClosedRequest), typeof(SourceGeneratedFileClosedResponse))]
+    public sealed record SourceGeneratedFileClosedRequest : SourceGeneratedFileInfo, IRequest
+    {
+    }
+}

src/OmniSharp.Abstractions/Models/v1/SourceGeneratedFile/SourceGeneratedFileClosedResponse.cs

@@ -0,0 +1,10 @@
+#nullable enable
+using System.Threading.Tasks;
+
+namespace OmniSharp.Models.v1.SourceGeneratedFile
+{
+    public sealed class SourceGeneratedFileClosedResponse
+    {
+        public static readonly Task<SourceGeneratedFileClosedResponse> Instance = Task.FromResult(new SourceGeneratedFileClosedResponse());
+    }
+}

src/OmniSharp.Abstractions/Models/v1/SourceGeneratedFile/SourceGeneratedFileInfo.cs

@@ -0,0 +1,12 @@
+#nullable enable
+
+using System;
+
+namespace OmniSharp.Models.v1.SourceGeneratedFile
+{
+    public record SourceGeneratedFileInfo
+    {
+        public Guid ProjectGuid { get; init; }
+        public Guid DocumentGuid { get; init; }
+    }
+}

src/OmniSharp.Abstractions/Models/v1/SourceGeneratedFile/SourceGeneratedFileRequest.cs

@@ -0,0 +1,12 @@
+#nullable enable
+
+using OmniSharp.Mef;
+using System;
+
+namespace OmniSharp.Models.v1.SourceGeneratedFile
+{
+    [OmniSharpEndpoint(OmniSharpEndpoints.SourceGeneratedFile, typeof(SourceGeneratedFileRequest), typeof(SourceGeneratedFileResponse))]
+    public sealed record SourceGeneratedFileRequest : SourceGeneratedFileInfo, IRequest
+    {
+    }
+}

src/OmniSharp.Abstractions/Models/v1/SourceGeneratedFile/SourceGeneratedFileResponse.cs

@@ -0,0 +1,10 @@
+#nullable enable
+
+namespace OmniSharp.Models.v1.SourceGeneratedFile
+{
+    public sealed record SourceGeneratedFileResponse
+    {
+        public string? SourceName { get; init; }
+        public string? Source { get; init; }
+    }
+}

src/OmniSharp.Abstractions/Models/v1/SourceGeneratedFile/UpdateSourceGeneratedFileRequest.cs

@@ -0,0 +1,10 @@
+#nullable enable
+using OmniSharp.Mef;
+
+namespace OmniSharp.Models.v1.SourceGeneratedFile
+{
+    [OmniSharpEndpoint(OmniSharpEndpoints.UpdateSourceGeneratedFile, typeof(UpdateSourceGeneratedFileRequest), typeof(UpdateSourceGeneratedFileResponse))]
+    public sealed record UpdateSourceGeneratedFileRequest : SourceGeneratedFileInfo, IRequest
+    {
+    }
+}

src/OmniSharp.Abstractions/Models/v1/SourceGeneratedFile/UpdateSourceGeneratedFileResponse.cs

@@ -0,0 +1,17 @@
+#nullable enable
+
+namespace OmniSharp.Models.v1.SourceGeneratedFile
+{
+    public record UpdateSourceGeneratedFileResponse
+    {
+        public UpdateType UpdateType { get; init; }
+        public string? Source { get; init; }
+    }
+
+    public enum UpdateType
+    {
+        Unchanged,
+        Deleted,
+        Modified
+    }
+}

src/OmniSharp.Abstractions/Models/v2/GotoDefinition/GotoDefinitionResponse.cs

@@ -1,6 +1,7 @@
 #nullable enable
 
 using OmniSharp.Models.Metadata;
+using OmniSharp.Models.v1.SourceGeneratedFile;
 using System.Collections.Generic;
 
 namespace OmniSharp.Models.V2.GotoDefinition
@@ -14,5 +15,6 @@ public record Definition
     {
         public Location Location { get; init; } = null!;
         public MetadataSource? MetadataSource { get; init; }
+        public SourceGeneratedFileInfo? SourceGeneratedFileInfo { get; init; }
     }
 }

src/OmniSharp.Abstractions/OmniSharpEndpoints.cs

@@ -51,6 +51,10 @@ public static class OmniSharpEndpoints
         public const string CompletionResolve = "/completion/resolve";
         public const string CompletionAfterInsert = "/completion/afterinsert";
 
+        public const string SourceGeneratedFile = "/sourcegeneratedfile";
+        public const string UpdateSourceGeneratedFile = "/updatesourcegeneratedfile";
+        public const string SourceGeneratedFileClosed = "/sourcegeneratedfileclosed";
+
         public static class V2
         {
             public const string GetCodeActions = "/v2/getcodeactions";

src/OmniSharp.Roslyn.CSharp/Services/Navigation/GoToDefinitionHelpers.cs

@@ -3,6 +3,8 @@
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.FindSymbols;
 using OmniSharp.Extensions;
+using OmniSharp.Models.v1.SourceGeneratedFile;
+using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -42,5 +44,21 @@ internal static class GoToDefinitionHelpers
 
             return null;
         }
+
+        internal static SourceGeneratedFileInfo? GetSourceGeneratedFileInfo(OmniSharpWorkspace workspace, Location location)
+        {
+            Debug.Assert(location.IsInSource);
+            var document = workspace.CurrentSolution.GetDocument(location.SourceTree);
+            if (document is not SourceGeneratedDocument)
+            {
+                return null;
+            }
+
+            return new SourceGeneratedFileInfo
+            {
+                ProjectGuid = document.Project.Id.Id,
+                DocumentGuid = document.Id.Id
+            };
+        }
     }
 }

src/OmniSharp.Roslyn.CSharp/Services/Navigation/GotoDefinitionService.cs

@@ -50,7 +50,8 @@ public async Task<GotoDefinitionResponse> Handle(GotoDefinitionRequest request)
                 {
                     FileName = lineSpan.Path,
                     Line = lineSpan.StartLinePosition.Line,
-                    Column = lineSpan.StartLinePosition.Character
+                    Column = lineSpan.StartLinePosition.Character,
+                    SourceGeneratedInfo = GoToDefinitionHelpers.GetSourceGeneratedFileInfo(_workspace, location)
                 };
             }
             else if (location.IsInMetadata && request.WantMetadata)

src/OmniSharp.Roslyn.CSharp/Services/Navigation/GotoDefinitionServiceV2.cs

@@ -49,7 +49,11 @@ public async Task<GotoDefinitionResponse> Handle(GotoDefinitionRequest request)
                 return new GotoDefinitionResponse()
                 {
                     Definitions = symbol.Locations
-                        .Select(location => new Definition { Location = location.GetMappedLineSpan().GetLocationFromFileLinePositionSpan() })
+                        .Select(location => new Definition
+                        {
+                            Location = location.GetMappedLineSpan().GetLocationFromFileLinePositionSpan(),
+                            SourceGeneratedFileInfo = GoToDefinitionHelpers.GetSourceGeneratedFileInfo(_workspace, location)
+                        })
                         .ToList()
                 };
             }

src/OmniSharp.Roslyn.CSharp/Services/Navigation/SourceGeneratedFileService.cs

@@ -0,0 +1,105 @@
+#nullable enable
+
+using Microsoft.CodeAnalysis;
+using Microsoft.Extensions.Logging;
+using OmniSharp.Mef;
+using OmniSharp.Models.v1.SourceGeneratedFile;
+using System.Collections.Generic;
+using System.Composition;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace OmniSharp.Roslyn.CSharp.Services.Navigation
+{
+    [Shared]
+    [OmniSharpHandler(OmniSharpEndpoints.SourceGeneratedFile, LanguageNames.CSharp)]
+    [OmniSharpHandler(OmniSharpEndpoints.UpdateSourceGeneratedFile, LanguageNames.CSharp)]
+    [OmniSharpHandler(OmniSharpEndpoints.SourceGeneratedFileClosed, LanguageNames.CSharp)]
+    public class SourceGeneratedFileService :
+        IRequestHandler<SourceGeneratedFileRequest, SourceGeneratedFileResponse>,
+        IRequestHandler<UpdateSourceGeneratedFileRequest, UpdateSourceGeneratedFileResponse>,
+        IRequestHandler<SourceGeneratedFileClosedRequest, SourceGeneratedFileClosedResponse>
+    {
+        private readonly OmniSharpWorkspace _workspace;
+        private readonly ILogger _logger;
+        private readonly Dictionary<DocumentId, VersionStamp> _lastSentVerisons = new();
+        private readonly object _lock = new();
+
+        [ImportingConstructor]
+        public SourceGeneratedFileService(OmniSharpWorkspace workspace, ILoggerFactory loggerFactory)
+        {
+            _workspace = workspace;
+            _logger = loggerFactory.CreateLogger<SourceGeneratedFileService>();
+        }
+
+        public async Task<SourceGeneratedFileResponse> Handle(SourceGeneratedFileRequest request)
+        {
+            var documentId = GetId(request);
+
+            var document = await _workspace.CurrentSolution.GetSourceGeneratedDocumentAsync(documentId, CancellationToken.None);
+
+            if (document is null)
+            {
+                _logger.LogError("Document with ID {0}:{1} was not found or not a source generated file", request.ProjectGuid, request.DocumentGuid);
+                return new SourceGeneratedFileResponse();
+            }
+
+            var text = await document.GetTextAsync();
+
+            var documentVerison = await document.GetTextVersionAsync();
+            lock (_lock)
+            {
+                _lastSentVerisons[documentId] = documentVerison;
+            }
+
+            return new SourceGeneratedFileResponse
+            {
+                Source = text.ToString(),
+                SourceName = document.FilePath
+            };
+        }
+
+        public async Task<UpdateSourceGeneratedFileResponse> Handle(UpdateSourceGeneratedFileRequest request)
+        {
+            var documentId = GetId(request);
+            var document = await _workspace.CurrentSolution.GetSourceGeneratedDocumentAsync(documentId, CancellationToken.None);
+            if (document == null)
+            {
+                lock (_lock)
+                {
+                    _ = _lastSentVerisons.Remove(documentId);
+                }
+                return new UpdateSourceGeneratedFileResponse() { UpdateType = UpdateType.Deleted };
+            }
+
+            var docVersion = await document.GetTextVersionAsync();
+            lock (_lock)
+            {
+                if (_lastSentVerisons.TryGetValue(documentId, out var lastVersion) && lastVersion == docVersion)
+                {
+                    return new UpdateSourceGeneratedFileResponse() { UpdateType = UpdateType.Unchanged };
+                }
+
+                _lastSentVerisons[documentId] = docVersion;
+            }
+
+            return new UpdateSourceGeneratedFileResponse()
+            {
+                UpdateType = UpdateType.Modified,
+                Source = (await document.GetTextAsync()).ToString()
+            };
+        }
+
+        public Task<SourceGeneratedFileClosedResponse> Handle(SourceGeneratedFileClosedRequest request)
+        {
+            lock (_lock)
+            {
+                _ = _lastSentVerisons.Remove(GetId(request));
+            }
+
+            return SourceGeneratedFileClosedResponse.Instance;
+        }
+
+        private DocumentId GetId(SourceGeneratedFileInfo info) => DocumentId.CreateFromSerialized(ProjectId.CreateFromSerialized(info.ProjectGuid), info.DocumentGuid);
+    }
+}

tests/OmniSharp.Roslyn.CSharp.Tests/AbstractGoToDefinitionFacts.cs

@@ -1,9 +1,12 @@
 using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
+using Microsoft.CodeAnalysis.Diagnostics;
 using OmniSharp.Mef;
 using OmniSharp.Models.Metadata;
+using OmniSharp.Models.v1.SourceGeneratedFile;
 using OmniSharp.Roslyn.CSharp.Services.Navigation;
 using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.Linq;
 using System.Threading.Tasks;
 using TestUtility;
@@ -506,6 +509,61 @@ public async Task ReturnsNoResultsButDoesNotThrowForNamespaces()
             Assert.Empty(GetInfo(response));
         }
 
+        [Fact]
+        public async Task ReturnsResultsForSourceGenerators()
+        {
+            const string Source = @"
+public class {|generatedClassName:Generated|}
+{
+    public int {|propertyName:Property|} { get; set; }
+}
+";
+            const string FileName = "real.cs";
+            TestFile generatedTestFile = new("GeneratedFile.cs", Source);
+            var testFile = new TestFile(FileName, @"
+class C
+{
+    public void M(Generated g)
+    {
+        _ = g.P$$roperty;
+}
+");
+
+            TestHelpers.AddProjectToWorkspace(SharedOmniSharpTestHost.Workspace,
+                "project.csproj",
+                new[] { "netcoreapp3.1" },
+                new[] { testFile },
+                analyzerRefs: ImmutableArray.Create<AnalyzerReference>(new TestGeneratorReference(
+                    context => context.AddSource("GeneratedFile", Source))));
+
+            var point = testFile.Content.GetPointFromPosition();
+
+            var gotoDefRequest = CreateRequest(FileName, point.Line, point.Offset, wantMetadata: true);
+            var gotoDefHandler = GetRequestHandler(SharedOmniSharpTestHost);
+            var response = await gotoDefHandler.Handle(gotoDefRequest);
+            var info = GetInfo(response).Single();
+
+            Assert.NotNull(info.SourceGeneratorInfo);
+
+            var expectedSpan = generatedTestFile.Content.GetSpans("propertyName").Single();
+            var expectedRange = generatedTestFile.Content.GetRangeFromSpan(expectedSpan);
+
+            Assert.Equal(expectedRange.Start.Line, info.Line);
+            Assert.Equal(expectedRange.Start.Offset, info.Column);
+
+            var sourceGeneratedFileHandler = SharedOmniSharpTestHost.GetRequestHandler<SourceGeneratedFileService>(OmniSharpEndpoints.SourceGeneratedFile);
+            var sourceGeneratedRequest = new SourceGeneratedFileRequest
+            {
+                DocumentGuid = info.SourceGeneratorInfo.DocumentGuid,
+                ProjectGuid = info.SourceGeneratorInfo.ProjectGuid
+            };
+
+            var sourceGeneratedFileResponse = await sourceGeneratedFileHandler.Handle(sourceGeneratedRequest);
+            Assert.NotNull(sourceGeneratedFileResponse);
+            Assert.Equal(generatedTestFile.Content.Code, sourceGeneratedFileResponse.Source);
+            Assert.Equal(@"OmniSharp.Roslyn.CSharp.Tests\OmniSharp.Roslyn.CSharp.Tests.TestSourceGenerator\GeneratedFile.cs", sourceGeneratedFileResponse.SourceName);
+        }
+
         protected async Task TestGoToSourceAsync(params TestFile[] testFiles)
         {
             var response = await GetResponseAsync(testFiles, wantMetadata: false);
@@ -581,6 +639,6 @@ protected async Task<TGotoDefinitionResponse> GetResponseAsync(TestFile[] testFi
 
         protected abstract TGotoDefinitionRequest CreateRequest(string fileName, int line, int column, bool wantMetadata, int timeout = 60000);
         protected abstract MetadataSource GetMetadataSource(TGotoDefinitionResponse response);
-        protected abstract IEnumerable<(int Line, int Column, string FileName)> GetInfo(TGotoDefinitionResponse response);
+        protected abstract IEnumerable<(int Line, int Column, string FileName, SourceGeneratedFileInfo SourceGeneratorInfo)> GetInfo(TGotoDefinitionResponse response);
     }
 }