Rewrite all extension-form calls of a document in one Fix-All pass

Funcky.Analyzers/Funcky.Analyzers.CodeFixes/CodeAnalysisExtensions/BooleanThen.cs

@@ -1,3 +1,4 @@
+using System.Collections.Immutable;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
@@ -11,10 +12,6 @@ namespace Funcky.Analyzers.CodeAnalysisExtensions;
 /// extension type nested in <c>BooleanExtensions</c> (or directly in it when declared as a classic extension method).</summary>
 internal static class BooleanThen
 {
-    public const string BooleanExtensionsNamespace = "Funcky.Extensions";
-
-    private static readonly SyntaxAnnotation ReplacementAnnotation = new();
-
     public static InvocationExpressionSyntax Invocation(ExpressionSyntax condition, SimpleNameSyntax thenName, ArgumentSyntax argument)
         => InvocationExpression(
             MemberAccessExpression(SyntaxKind.SimpleMemberAccessExpression, Receiver(condition), thenName),
@@ -46,31 +43,40 @@ public static bool IsThen(ISymbol? symbol, INamedTypeSymbol booleanExtensionsTyp
     /// <summary>Replaces <paramref name="target"/> with the <c>Then</c> call, adding <c>using Funcky.Extensions;</c> when the call does not
     /// bind without it. Returns <see langword="null"/> when the call still does not bind to Funcky's <c>Then</c> afterwards, because
     /// another <c>Then</c> extension in scope wins or makes the call ambiguous.</summary>
-    public static async Task<Document?> TryReplaceAsync(Document document, SyntaxNode root, SemanticModel semanticModel, ExpressionSyntax target, InvocationExpressionSyntax replacement, INamedTypeSymbol booleanExtensionsType, CancellationToken cancellationToken)
-    {
-        var rootWithReplacement = root.ReplaceNode(target, replacement.WithAdditionalAnnotations(ReplacementAnnotation));
+    public static Task<Document?> TryReplaceAsync(Document document, SyntaxNode root, SemanticModel semanticModel, ExpressionSyntax target, InvocationExpressionSyntax replacement, INamedTypeSymbol booleanExtensionsType, CancellationToken cancellationToken)
+        => TryReplaceAllAsync(document, root, semanticModel, ImmutableArray.Create(new ExtensionMemberRewrite(target, replacement)), booleanExtensionsType, cancellationToken);
 
-        if (IsThen(semanticModel.GetSpeculativeSymbolInfo(target.SpanStart, replacement, SpeculativeBindingOption.BindAsExpression).Symbol, booleanExtensionsType))
+    /// <summary>Applies all rewrites at once, adding <c>using Funcky.Extensions;</c> once when a <c>Then</c> call does not bind without it.
+    /// A call that still does not bind to Funcky's <c>Then</c> afterwards, because another <c>Then</c> extension in scope wins or makes
+    /// it ambiguous, is left as it is; the result is <see langword="null"/> when that holds for every call.</summary>
+    public static async Task<Document?> TryReplaceAllAsync(Document document, SyntaxNode root, SemanticModel semanticModel, ImmutableArray<ExtensionMemberRewrite> rewrites, INamedTypeSymbol booleanExtensionsType, CancellationToken cancellationToken)
+    {
+        if (rewrites.IsEmpty)
         {
-            return document.WithSyntaxRoot(rootWithReplacement);
+            return null;
         }
 
-        var documentWithUsing = document.WithSyntaxRoot(UsingDirectives.AddUsing(rootWithReplacement, rootWithReplacement.GetAnnotatedNodes(ReplacementAnnotation).Single(), BooleanExtensionsNamespace));
+        var annotatedRoot = FunckyExtensionsImport.ReplaceWithAnnotatedExtensionMemberCalls(root, semanticModel, rewrites, symbol => IsThen(symbol, booleanExtensionsType));
+        var unboundRewrites = await UnboundRewritesAsync(document.WithSyntaxRoot(annotatedRoot), rewrites.Length, cancellationToken).ConfigureAwait(false);
 
-        return await BindsToThenAsync(documentWithUsing, cancellationToken).ConfigureAwait(false)
-            ? documentWithUsing
-            : null;
+        return unboundRewrites.IsEmpty
+            ? document.WithSyntaxRoot(FunckyExtensionsImport.WithoutReplacementAnnotations(annotatedRoot))
+            : await TryReplaceAllAsync(document, root, semanticModel, rewrites.Where((_, index) => !unboundRewrites.Contains(index)).ToImmutableArray(), booleanExtensionsType, cancellationToken).ConfigureAwait(false);
     }
 
-    private static async Task<bool> BindsToThenAsync(Document document, CancellationToken cancellationToken)
+    /// <summary>The indices of the rewrites whose <c>Then</c> call does not bind to Funcky's extension in the changed document; all of
+    /// them when the document cannot be bound at all.</summary>
+    private static async Task<ImmutableArray<int>> UnboundRewritesAsync(Document document, int rewriteCount, CancellationToken cancellationToken)
     {
         var root = await document.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
         var semanticModel = await document.GetSemanticModelAsync(cancellationToken).ConfigureAwait(false);
 
-        return root is not null
-               && semanticModel is not null
-               && semanticModel.Compilation.GetBooleanExtensionsType() is { } booleanExtensionsType
-               && IsThen(semanticModel.GetSymbolInfo(root.GetAnnotatedNodes(ReplacementAnnotation).Single(), cancellationToken).Symbol, booleanExtensionsType);
+        return root is not null && semanticModel is not null && semanticModel.Compilation.GetBooleanExtensionsType() is { } booleanExtensionsType
+            ? FunckyExtensionsImport.Replacements(root)
+                .Where(replacement => !IsThen(semanticModel.GetSymbolInfo(replacement.Replacement, cancellationToken).Symbol, booleanExtensionsType))
+                .Select(replacement => replacement.Index)
+                .ToImmutableArray()
+            : ImmutableArray.CreateRange(Enumerable.Range(0, rewriteCount));
     }
 
     // A member access binds tighter than every operator, so anything but a primary expression needs parentheses.

Funcky.Analyzers/Funcky.Analyzers.CodeFixes/CodeAnalysisExtensions/ExtensionMemberRewrite.cs

@@ -0,0 +1,6 @@
+using Microsoft.CodeAnalysis.CSharp.Syntax;
+
+namespace Funcky.Analyzers.CodeAnalysisExtensions;
+
+/// <summary>An expression to replace with a call of a static extension member of a class in <c>Funcky.Extensions</c>.</summary>
+public sealed record ExtensionMemberRewrite(ExpressionSyntax Original, ExpressionSyntax Replacement);

Funcky.Analyzers/Funcky.Analyzers.CodeFixes/CodeAnalysisExtensions/FunckyExtensionsImport.cs

@@ -1,22 +1,74 @@
+using System.Collections.Immutable;
+using System.Globalization;
 using Funcky.Analyzers.ExtensionForms;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 
 namespace Funcky.Analyzers.CodeAnalysisExtensions;
 
-/// <summary>Replaces an expression with one that calls a static extension member of a class in <c>Funcky.Extensions</c>, and
-/// adds the <c>using</c> for that namespace when the replacement does not bind to such a member without it.</summary>
+/// <summary>Replaces expressions with ones that call a static extension member of a class in <c>Funcky.Extensions</c>, and
+/// adds the <c>using</c> for that namespace when a replacement does not bind to such a member without it.</summary>
 internal static class FunckyExtensionsImport
 {
-    private static readonly SyntaxAnnotation ReplacementAnnotation = new();
+    private const string ReplacementAnnotationKind = "Funcky.Analyzers.FunckyExtensionsImport.Replacement";
 
     public static Document ReplaceWithExtensionMemberCall(Document document, SyntaxNode root, SemanticModel semanticModel, ExpressionSyntax original, ExpressionSyntax replacement)
+        => ReplaceWithExtensionMemberCalls(document, root, semanticModel, ImmutableArray.Create(new ExtensionMemberRewrite(original, replacement)), ExtensionMemberSupport.IsDeclaredInFunckyExtensions);
+
+    /// <summary>Applies all rewrites to the document at once. The <c>using</c> is added once per group of usings (the enclosing
+    /// namespace declaration or the file) that has a replacement in need of it, so that any number of rewrites work in one pass.
+    /// A rewrite nested in another one is skipped, the outer replacement is built from the original and would discard it anyway;
+    /// its diagnostic stays and is fixed by a subsequent pass. <paramref name="bindsToExtensionMember"/> tells whether the symbol a
+    /// replacement binds to without the <c>using</c> is the intended extension member.</summary>
+    public static Document ReplaceWithExtensionMemberCalls(Document document, SyntaxNode root, SemanticModel semanticModel, ImmutableArray<ExtensionMemberRewrite> rewrites, Func<ISymbol?, bool> bindsToExtensionMember)
+        => document.WithSyntaxRoot(WithoutReplacementAnnotations(ReplaceWithAnnotatedExtensionMemberCalls(root, semanticModel, rewrites, bindsToExtensionMember)));
+
+    /// <summary>Like <see cref="ReplaceWithExtensionMemberCalls"/>, but the replacements in the returned root still carry the annotations
+    /// that <see cref="Replacements"/> finds, so that a caller can bind them in the changed document before it strips them with
+    /// <see cref="WithoutReplacementAnnotations"/>. A document handed to the host must not keep them: a later fix on the same document
+    /// would take them for its own.</summary>
+    public static SyntaxNode ReplaceWithAnnotatedExtensionMemberCalls(SyntaxNode root, SemanticModel semanticModel, ImmutableArray<ExtensionMemberRewrite> rewrites, Func<ISymbol?, bool> bindsToExtensionMember)
     {
-        var annotatedReplacement = replacement.WithAdditionalAnnotations(ReplacementAnnotation);
-        var rootWithReplacement = root.ReplaceNode(original, annotatedReplacement);
+        var originals = rewrites.Select(rewrite => rewrite.Original).ToImmutableHashSet();
+        var outermostRewrites = rewrites
+            .Select((rewrite, index) => (Rewrite: rewrite, Index: index))
+            .Where(pair => !pair.Rewrite.Original.Ancestors().Any(originals.Contains))
+            .GroupBy(pair => pair.Rewrite.Original)
+            .ToDictionary(group => group.Key, group => group.First());
+        var rootWithReplacements = root.ReplaceNodes(outermostRewrites.Keys, (original, _) => outermostRewrites[original].Rewrite.Replacement.WithAdditionalAnnotations(Annotation(outermostRewrites[original].Index)));
+        var replacementsInNeedOfImport = outermostRewrites.Values
+            .Where(pair => !bindsToExtensionMember(semanticModel.GetSpeculativeSymbolInfo(pair.Rewrite.Original.SpanStart, pair.Rewrite.Replacement, SpeculativeBindingOption.BindAsExpression).Symbol))
+            .Select(pair => pair.Index)
+            .ToImmutableArray();
 
-        return ExtensionMemberSupport.IsDeclaredInFunckyExtensions(semanticModel.GetSpeculativeSymbolInfo(original.SpanStart, replacement, SpeculativeBindingOption.BindAsExpression).Symbol)
-            ? document.WithSyntaxRoot(rootWithReplacement)
-            : document.WithSyntaxRoot(UsingDirectives.AddUsing(rootWithReplacement, rootWithReplacement.GetAnnotatedNodes(ReplacementAnnotation).Single(), ExtensionMemberSupport.ExtensionsNamespace));
+        return AddUsings(rootWithReplacements, replacementsInNeedOfImport);
     }
+
+    /// <summary>The replacements in a root produced by <see cref="ReplaceWithAnnotatedExtensionMemberCalls"/>, each with the index of the rewrite it stems from.</summary>
+    public static IEnumerable<(int Index, SyntaxNode Replacement)> Replacements(SyntaxNode root)
+        => root.GetAnnotatedNodes(ReplacementAnnotationKind).Select(node => (Index(node), node));
+
+    public static SyntaxNode WithoutReplacementAnnotations(SyntaxNode root)
+        => root.ReplaceNodes(root.GetAnnotatedNodes(ReplacementAnnotationKind), (_, replacement) => replacement.WithoutAnnotations(ReplacementAnnotationKind));
+
+    private static SyntaxNode AddUsings(SyntaxNode root, ImmutableArray<int> replacementsInNeedOfImport)
+    {
+        var anchors = Replacements(root)
+            .Where(replacement => replacementsInNeedOfImport.Contains(replacement.Index))
+            .GroupBy(replacement => UsingContainer(replacement.Replacement)?.Span ?? default)
+            .Select(group => group.First().Index)
+            .ToImmutableArray();
+
+        return anchors.Aggregate(root, (current, anchor) => UsingDirectives.AddUsing(current, Replacements(current).Single(replacement => replacement.Index == anchor).Replacement, ExtensionMemberSupport.ExtensionsNamespace));
+    }
+
+    // Where UsingDirectives.AddUsing puts the directive: the enclosing namespace declaration with usings, or the file.
+    private static BaseNamespaceDeclarationSyntax? UsingContainer(SyntaxNode node)
+        => node.FirstAncestorOrSelf<BaseNamespaceDeclarationSyntax>(namespaceDeclaration => namespaceDeclaration.Usings.Count > 0);
+
+    private static SyntaxAnnotation Annotation(int index)
+        => new(ReplacementAnnotationKind, index.ToString(CultureInfo.InvariantCulture));
+
+    private static int Index(SyntaxNode replacement)
+        => int.Parse(replacement.GetAnnotations(ReplacementAnnotationKind).Single().Data!, CultureInfo.InvariantCulture);
 }

Funcky.Analyzers/Funcky.Analyzers.CodeFixes/ExtensionForms/PreferEnumerableSequenceConstructorsCodeFix.cs

@@ -3,7 +3,6 @@
 using Funcky.Analyzers.CodeAnalysisExtensions;
 using Funcky.Analyzers.ExtensionForms;
 using Microsoft.CodeAnalysis;
-using Microsoft.CodeAnalysis.CodeActions;
 using Microsoft.CodeAnalysis.CodeFixes;
 using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
@@ -14,28 +13,20 @@ namespace Funcky.Analyzers;
 
 [Shared]
 [ExportCodeFixProvider(LanguageNames.CSharp, Name = nameof(PreferEnumerableSequenceConstructorsCodeFix))]
-public sealed class PreferEnumerableSequenceConstructorsCodeFix : CodeFixProvider
+public sealed class PreferEnumerableSequenceConstructorsCodeFix : ExtensionMemberCodeFix
 {
     private const string EnumerableTypeName = "global::System.Linq.Enumerable";
 
     public override ImmutableArray<string> FixableDiagnosticIds => ImmutableArray.Create(PreferEnumerableSequenceConstructorsAnalyzer.PreferEnumerable.Id);
 
-    public override FixAllProvider GetFixAllProvider() => WellKnownFixAllProviders.BatchFixer;
+    protected override string GetTitle(ExtensionMemberRewrite rewrite) => "Replace with the Enumerable form";
 
-    public override async Task RegisterCodeFixesAsync(CodeFixContext context)
-    {
-        var root = await context.Document.GetSyntaxRootAsync(context.CancellationToken).ConfigureAwait(false);
-
-        foreach (var diagnostic in context.Diagnostics)
-        {
-            if (root?.FindInvocationExpression(diagnostic.Location.SourceSpan) is { } invocation
-                && invocation.Span == diagnostic.Location.SourceSpan
-                && WithEnumerableReceiver(invocation.Expression) is { } replacement)
-            {
-                context.RegisterCodeFix(new UseEnumerableCodeFixAction(context.Document, invocation, replacement), diagnostic);
-            }
-        }
-    }
+    protected override ExtensionMemberRewrite? TryCreateRewrite(SyntaxNode root, SemanticModel semanticModel, Diagnostic diagnostic, CancellationToken cancellationToken)
+        => root.FindInvocationExpression(diagnostic.Location.SourceSpan) is { } invocation
+           && invocation.Span == diagnostic.Location.SourceSpan
+           && WithEnumerableReceiver(invocation.Expression) is { } replacement
+            ? new ExtensionMemberRewrite(invocation, replacement)
+            : null;
 
     // `Sequence.Return(x)` and `Funcky.Sequence.Return(x)` keep their name (and type arguments); `Return(x)` through a
     // `using static` gets the Enumerable receiver in front.
@@ -49,21 +40,4 @@ public override async Task RegisterCodeFixesAsync(CodeFixContext context)
 
     private static ExpressionSyntax Enumerable()
         => ParseExpression(EnumerableTypeName).WithAdditionalAnnotations(Simplifier.Annotation);
-
-    private sealed class UseEnumerableCodeFixAction(Document document, InvocationExpressionSyntax invocation, InvocationExpressionSyntax replacement) : CodeAction
-    {
-        public override string Title => "Replace with the Enumerable form";
-
-        public override string EquivalenceKey => nameof(PreferEnumerableSequenceConstructorsCodeFix);
-
-        protected override async Task<Document> GetChangedDocumentAsync(CancellationToken cancellationToken)
-        {
-            var root = await document.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
-            var semanticModel = await document.GetSemanticModelAsync(cancellationToken).ConfigureAwait(false);
-
-            return root is null || semanticModel is null
-                ? document
-                : FunckyExtensionsImport.ReplaceWithExtensionMemberCall(document, root, semanticModel, invocation, replacement);
-        }
-    }
 }