Delay loading of CodeRefactoringProvider's till absolutely needed (#78495)

Delay loading of CodeRefactoringProvider's till absolutely needed. For
example, we do not want to load refactoring provider for
AdditonalDocument XAML files unless user is editing such files.

Author: jasonmalinowski

src/EditorFeatures/Test/CodeRefactorings/CodeRefactoringServiceTest.cs

@@ -8,6 +8,7 @@
 using System.Collections.Immutable;
 using System.Composition;
 using System.Linq;
+using System.Reflection;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.CodeActions;
@@ -229,6 +230,63 @@ public async Task TestAnalyzerConfigDocumentRefactoringAsync()
         Assert.Equal(refactoring2.Title, globalConfigRefactoringTitle);
     }
 
+    [Fact]
+    public void TestDelayedLoading()
+    {
+        var composition = FeaturesTestCompositions.Features.AddParts(
+            typeof(NonSourceFileRefactoringWithDocumentKindsAndExtensions),
+            typeof(NonSourceFileRefactoringWithDocumentKinds),
+            typeof(NonSourceFileRefactoringWithDocumentExtensions),
+            typeof(NonSourceFileRefactoringWithoutDocumentKindsAndExtensions));
+
+        using var workspace = TestWorkspace.CreateCSharp("", composition: composition);
+        var refactoringService = (CodeRefactorings.CodeRefactoringService)workspace.GetService<ICodeRefactoringService>();
+
+        var project = workspace.CurrentSolution.Projects.Single()
+            .AddAdditionalDocument("test.TXT", "", filePath: "test.TXT").Project
+            .AddAdditionalDocument("test", "", filePath: "test").Project
+            .AddAdditionalDocument("test.log", "", filePath: "test.log").Project
+            .AddDocument("test.cs", "", filePath: "test.cs").Project
+            .AddDocument("test.editorconfig", "", filePath: "test.editorconfig").Project;
+
+        VerifyProviders(refactoringService,
+            project.AdditionalDocuments.Single(t => t.Name == "test.TXT"),
+            typeof(NonSourceFileRefactoringWithDocumentKindsAndExtensions),
+            typeof(NonSourceFileRefactoringWithDocumentKinds));
+
+        VerifyProviders(refactoringService,
+            project.AdditionalDocuments.Single(t => t.Name == "test"),
+            typeof(NonSourceFileRefactoringWithDocumentKinds));
+
+        VerifyProviders(refactoringService,
+            project.AdditionalDocuments.Single(t => t.Name == "test.log"),
+            typeof(NonSourceFileRefactoringWithDocumentKinds));
+
+        VerifyProviders(refactoringService,
+            project.Documents.Single(t => t.Name == "test.editorconfig"),
+            typeof(NonSourceFileRefactoringWithDocumentExtensions),
+            typeof(NonSourceFileRefactoringWithoutDocumentKindsAndExtensions));
+
+        VerifyProviders(refactoringService,
+            project.Documents.Single(t => t.Name == "test.cs"),
+            typeof(NonSourceFileRefactoringWithoutDocumentKindsAndExtensions));
+    }
+
+    private static void VerifyProviders(CodeRefactorings.CodeRefactoringService service, TextDocument document, params Type[] expectedProviderTypes)
+    {
+        // Exclude providers which have not been setup by test
+        var assembly = Assembly.GetExecutingAssembly();
+        var actualProviders = service.GetProviders(document)
+            .Where(p => p.GetType().Assembly == assembly)
+            .ToArray();
+
+        foreach (var type in expectedProviderTypes)
+        {
+            Assert.Contains(actualProviders, p => p.GetType() == type);
+        }
+        Assert.Equal(expectedProviderTypes.Length, actualProviders.Length);
+    }
+
     internal abstract class AbstractNonSourceFileRefactoring : CodeRefactoringProvider
     {
         public string Title { get; }

src/Features/Core/Portable/CodeRefactorings/CodeRefactoringService.cs

@@ -30,43 +30,73 @@ namespace Microsoft.CodeAnalysis.CodeRefactorings;
 internal sealed class CodeRefactoringService(
     [ImportMany] IEnumerable<Lazy<CodeRefactoringProvider, CodeChangeProviderMetadata>> providers) : ICodeRefactoringService
 {
-    private readonly Lazy<ImmutableDictionary<string, Lazy<ImmutableArray<CodeRefactoringProvider>>>> _lazyLanguageToProvidersMap = new Lazy<ImmutableDictionary<string, Lazy<ImmutableArray<CodeRefactoringProvider>>>>(
-            () =>
-                ImmutableDictionary.CreateRange(
-                    DistributeLanguages(providers)
-                        .GroupBy(lz => lz.Metadata.Language)
-                        .Select(grp => KeyValuePairUtil.Create(
-                            grp.Key,
-                            new Lazy<ImmutableArray<CodeRefactoringProvider>>(() => [.. ExtensionOrderer.Order(grp).Select(lz => lz.Value)])))));
+    private readonly Lazy<ImmutableDictionary<ProviderKey, Lazy<ImmutableArray<CodeRefactoringProvider>>>> _lazyLanguageDocumentToProvidersMap =
+        new Lazy<ImmutableDictionary<ProviderKey, Lazy<ImmutableArray<CodeRefactoringProvider>>>>(() =>
+            ImmutableDictionary.CreateRange(
+                DistributeLanguagesAndDocuments(providers)
+                    .GroupBy(lz => new ProviderKey(lz.Metadata.Language, lz.Metadata.DocumentKind, lz.Metadata.DocumentExtension))
+                    .Select(grp => KeyValuePairUtil.Create(grp.Key,
+                        new Lazy<ImmutableArray<CodeRefactoringProvider>>(() => [.. ExtensionOrderer.Order(grp).Select(lz => lz.Value)])))));
+
     private readonly Lazy<ImmutableDictionary<CodeRefactoringProvider, CodeChangeProviderMetadata>> _lazyRefactoringToMetadataMap = new(() => providers.Where(provider => provider.IsValueCreated).ToImmutableDictionary(provider => provider.Value, provider => provider.Metadata));
 
     private ImmutableDictionary<CodeRefactoringProvider, FixAllProviderInfo?> _fixAllProviderMap = ImmutableDictionary<CodeRefactoringProvider, FixAllProviderInfo?>.Empty;
 
-    private static IEnumerable<Lazy<CodeRefactoringProvider, OrderableLanguageMetadata>> DistributeLanguages(IEnumerable<Lazy<CodeRefactoringProvider, CodeChangeProviderMetadata>> providers)
+    private static IEnumerable<Lazy<CodeRefactoringProvider, OrderableLanguageDocumentMetadata>> DistributeLanguagesAndDocuments(IEnumerable<Lazy<CodeRefactoringProvider, CodeChangeProviderMetadata>> providers)
     {
         foreach (var provider in providers)
         {
             foreach (var language in provider.Metadata.Languages)
             {
-                var orderable = new OrderableLanguageMetadata(
-                    provider.Metadata.Name, language, provider.Metadata.AfterTyped, provider.Metadata.BeforeTyped);
-                yield return new Lazy<CodeRefactoringProvider, OrderableLanguageMetadata>(() => provider.Value, orderable);
+                foreach (var documentKind in provider.Metadata.DocumentKinds)
+                {
+                    // Document kinds come from ExportCodeRefactoringProviderAttribute which throws
+                    // if values do not match enum TextDocumentKind. Here we'll throw too.
+                    var kind = (TextDocumentKind)Enum.Parse(typeof(TextDocumentKind), documentKind, ignoreCase: true);
+                    var documentExtensions = provider.Metadata.DocumentExtensions.Where(e => !string.IsNullOrEmpty(e));
+                    if (!documentExtensions.Any())
+                    {
+                        // Metadata without file extension applies to all files.
+                        documentExtensions = [""];
+                    }
+
+                    foreach (var documentExtension in documentExtensions)
+                    {
+                        var orderable = new OrderableLanguageDocumentMetadata(
+                            provider.Metadata.Name ?? "", language, kind, documentExtension, provider.Metadata.AfterTyped, provider.Metadata.BeforeTyped);
+                        yield return new Lazy<CodeRefactoringProvider, OrderableLanguageDocumentMetadata>(() => provider.Value, orderable);
+                    }
+                }
             }
         }
     }
 
-    private ImmutableDictionary<string, Lazy<ImmutableArray<CodeRefactoringProvider>>> LanguageToProvidersMap
-        => _lazyLanguageToProvidersMap.Value;
+    private ImmutableDictionary<ProviderKey, Lazy<ImmutableArray<CodeRefactoringProvider>>> LanguageDocumentToProvidersMap
+        => _lazyLanguageDocumentToProvidersMap.Value;
 
     private ImmutableDictionary<CodeRefactoringProvider, CodeChangeProviderMetadata> RefactoringToMetadataMap
         => _lazyRefactoringToMetadataMap.Value;
 
-    private ConcatImmutableArray<CodeRefactoringProvider> GetProviders(TextDocument document)
+    internal ConcatImmutableArray<CodeRefactoringProvider> GetProviders(TextDocument document)
     {
         var allRefactorings = ImmutableArray<CodeRefactoringProvider>.Empty;
-        if (LanguageToProvidersMap.TryGetValue(document.Project.Language, out var lazyProviders))
+
+        // Include providers which apply to all extensions
+        var key = new ProviderKey(document.Project.Language, document.Kind, "");
+        if (LanguageDocumentToProvidersMap.TryGetValue(key, out var lazyProviders))
+        {
+            allRefactorings = lazyProviders.Value;
+        }
+
+        // Get providers for specific combination of language, doc kind and extension,
+        // e.g. (C#, AdditionalDocument, .xaml)
+        if (FileNameUtilities.GetExtension(document.FilePath) is string documentExtension && documentExtension.Length > 0)
         {
-            allRefactorings = ProjectCodeRefactoringProvider.FilterExtensions(document, lazyProviders.Value, GetExtensionInfo);
+            key = new ProviderKey(document.Project.Language, document.Kind, documentExtension);
+            if (LanguageDocumentToProvidersMap.TryGetValue(key, out lazyProviders))
+            {
+                allRefactorings = allRefactorings.Concat(lazyProviders.Value);
+            }
         }
 
         return allRefactorings.ConcatFast(GetProjectRefactorings(document));
@@ -271,4 +301,24 @@ protected override bool TryGetExtensionsFromReference(AnalyzerReference referenc
             return false;
         }
     }
+
+    private record struct ProviderKey(string Language, TextDocumentKind DocumentKind, string DocumentExtension) : IEquatable<ProviderKey>
+    {
+        public bool Equals(ProviderKey other)
+        {
+            // We create keys from two sources:
+            //   * MEF's ExportCodeRefactoringProviderAttribute when building the map for available providers.
+            //   * TextDocument when looking up the map for available providers.
+            // Text documents can point to files with different extensions, e.g. MyPage.xaml and MyControl.XAML.
+            // Thus we need case insensitive comparison for DocumentExtension.
+            return Language == other.Language &&
+                   DocumentKind == other.DocumentKind &&
+                   StringComparer.OrdinalIgnoreCase.Equals(DocumentExtension, other.DocumentExtension);
+        }
+
+        public override int GetHashCode()
+        {
+            return (Language, DocumentKind, StringComparer.OrdinalIgnoreCase.GetHashCode(DocumentExtension)).GetHashCode();
+        }
+    }
 }

src/Workspaces/SharedUtilitiesAndExtensions/Workspace/Core/Workspace/Mef/CodeChangeProviderMetadata.cs

@@ -12,16 +12,22 @@ namespace Microsoft.CodeAnalysis.Host.Mef;
 internal sealed class CodeChangeProviderMetadata : OrderableMetadata, ILanguagesMetadata
 {
     public IEnumerable<string> Languages { get; }
+    public IEnumerable<string> DocumentKinds { get; }
+    public IEnumerable<string> DocumentExtensions { get; }
 
     public CodeChangeProviderMetadata(IDictionary<string, object> data)
         : base(data)
     {
         this.Languages = ((IReadOnlyDictionary<string, object>)data).GetEnumerableMetadata<string>("Languages");
+        this.DocumentKinds = ((IReadOnlyDictionary<string, object>)data).GetEnumerableMetadata<string>("DocumentKinds");
+        this.DocumentExtensions = ((IReadOnlyDictionary<string, object>)data).GetEnumerableMetadata<string>("DocumentExtensions");
     }
 
     public CodeChangeProviderMetadata(string name, IEnumerable<string> after = null, IEnumerable<string> before = null, params string[] languages)
         : base(name, after, before)
     {
         this.Languages = languages;
+        this.DocumentKinds = [];
+        this.DocumentExtensions = [];
     }
 }

src/Workspaces/SharedUtilitiesAndExtensions/Workspace/Core/Workspace/Mef/OrderableLanguageDocumentMetadata.cs

@@ -0,0 +1,22 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+#nullable disable
+
+using System.Collections.Generic;
+
+namespace Microsoft.CodeAnalysis.Host.Mef;
+
+internal sealed class OrderableLanguageDocumentMetadata : OrderableLanguageMetadata
+{
+    public TextDocumentKind DocumentKind { get; }
+    public string DocumentExtension { get; }
+
+    public OrderableLanguageDocumentMetadata(string name, string language, TextDocumentKind documentKind, string documentExtension, IEnumerable<string> after, IEnumerable<string> before)
+        : base(name, language, after, before)
+    {
+        DocumentKind = documentKind;
+        DocumentExtension = documentExtension;
+    }
+}

src/Workspaces/SharedUtilitiesAndExtensions/Workspace/Core/WorkspaceExtensions.projitems

@@ -169,6 +169,7 @@
     <Compile Include="$(MSBuildThisFileDirectory)Workspace\Mef\MefLanguageServices.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Workspace\Mef\MefUtilities.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Workspace\Mef\MefWorkspaceServices.cs" />
+    <Compile Include="$(MSBuildThisFileDirectory)Workspace\Mef\OrderableLanguageDocumentMetadata.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Workspace\Mef\OrderableLanguageMetadata.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Workspace\Mef\OrderableMetadata.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Workspace\Mef\WorkspaceServiceMetadata.cs" />