Add option to preserve CancellationToken parameters

It can still be useful to preserve the CancellationToken in some scenarios. For example, for [cancelling a synchronous bulk copy operation][1].

```csharp
bulkCopy.SqlRowsCopied += (_, e) =>
{
    if (cancellationToken.IsCancellationRequested)
    {
        e.Abort = true;
    }
};
```

I plan to submit a pull request at [PhenX.EntityFrameworkCore.BulkInsert][2] to use *Sync Method Generator* in order to greatly simplify async + sync methods implementations. 😉

[1]: https://github.com/PhenX/PhenX.EntityFrameworkCore.BulkInsert/blob/137d2fc8fed17b5aa7e6f11fccc079b7f463aff0/src/PhenX.EntityFrameworkCore.BulkInsert.SqlServer/SqlServerBulkInsertProvider.cs#L68-L74
[2]: https://github.com/PhenX/PhenX.EntityFrameworkCore.BulkInsert

Author: 0xced

README.md

@@ -62,7 +62,7 @@ A list of changes applied to the new synchronized method:
   \*\* `Memory` and `ReadOnlyMemory` is preserved in sync methods if it is a type argument of a collection. This is due to a compiler limitation which states that a `ref struct` can't be the element type of an array.
 
 - Remove parameters
-  - [CancellationToken](https://learn.microsoft.com/en-us/dotnet/api/system.threading.cancellationtoken)
+  - [CancellationToken](https://learn.microsoft.com/en-us/dotnet/api/system.threading.cancellationtoken), unless the `PreserveCancellationToken` property is set to `true`.
   - [IProgress\<T>](https://learn.microsoft.com/en-us/dotnet/api/system.iprogress-1), unless the `PreserveProgress` property is set to `true`.
 - Invocation changes
   - Remove `ConfigureAwait` from [Tasks](https://learn.microsoft.com/en-us/dotnet/api/system.threading.tasks.task.configureawait) and [Asynchronous Enumerations](https://learn.microsoft.com/en-us/dotnet/api/system.threading.tasks.taskasyncenumerableextensions.configureawait)
@@ -106,6 +106,18 @@ public async Task MethodAsync(IProgress<double> progress)
 }
 ```
 
+#### PreserveCancellationToken
+
+By default, this source generator removes `CancellationToken` parameters from async methods. To preserve them, use the `PreserveCancellationToken` option.
+
+```cs
+[Zomp.SyncMethodGenerator.CreateSyncVersion(PreserveCancellationToken = true)]
+public async Task MethodAsync(CancellationToken cancellationToken = default)
+{
+    cancellationToken.ThrowIfCancellationRequested();
+}
+```
+
 ### CreateSyncVersionAttribute on a type
 
 You can also decorate your type (class, struct, record, or interface) to generate a sync version for every asynchronous method.

src/Zomp.SyncMethodGenerator/AsyncToSyncRewriter.cs

@@ -11,8 +11,9 @@ namespace Zomp.SyncMethodGenerator;
 /// <param name="semanticModel">The semantic model.</param>
 /// <param name="disableNullable">Instructs the source generator that nullable context should be disabled.</param>
 /// <param name="preserveProgress">Instructs the source generator to preserve <see cref="IProgress{T}"/> parameters.</param>
+/// <param name="preserveCancellationToken">Instructs the source generator to preserve <see cref="CancellationToken"/> parameters.</param>
 /// <param name="methodName">Method declaration syntax.</param>
-internal sealed class AsyncToSyncRewriter(SemanticModel semanticModel, bool disableNullable, bool preserveProgress, MethodDeclarationSyntax methodName) : CSharpSyntaxRewriter
+internal sealed class AsyncToSyncRewriter(SemanticModel semanticModel, bool disableNullable, bool preserveProgress, bool preserveCancellationToken, MethodDeclarationSyntax methodName) : CSharpSyntaxRewriter
 {
     public const string SyncOnly = "SYNC_ONLY";
 
@@ -455,8 +456,19 @@ var entries
             }
         }
 
+        // Remove [EnumeratorCancellation] attribute
+        var attributeLists = default(SyntaxList<AttributeListSyntax>);
+        foreach (var attributeList in node.AttributeLists)
+        {
+            var attributes = SeparatedList(attributeList.Attributes.Where(a => !IsEnumeratorCancellationAttribute(a)));
+            if (attributes.Count > 0)
+            {
+                attributeLists = attributeLists.Add(attributeList.WithAttributes(attributes));
+            }
+        }
+
         return node.Type is null || TypeAlreadyQualified(node.Type) ? @base
-            : @base.WithType(ProcessType(node.Type)).WithTriviaFrom(@base);
+            : @base.WithType(ProcessType(node.Type)).WithAttributeLists(attributeLists).WithTriviaFrom(@base);
     }
 
     /// <inheritdoc/>
@@ -1879,11 +1891,11 @@ private string GetNewName(IMethodSymbol methodSymbol)
     private string ReplaceWithSpan(ISymbol symbol)
         => Regex.Replace(symbol.Name, Memory, Span);
 
-    private bool ShouldRemoveType(ITypeSymbol symbol)
+    private bool ShouldRemoveType(ITypeSymbol symbol, bool isArgument = false)
     {
         if (symbol is IArrayTypeSymbol at)
         {
-            return ShouldRemoveType(at.ElementType);
+            return ShouldRemoveType(at.ElementType, isArgument);
         }
 
         if (symbol is not INamedTypeSymbol namedSymbol)
@@ -1903,7 +1915,7 @@ private bool ShouldRemoveType(ITypeSymbol symbol)
             }
         }
 
-        return (namedSymbol.IsIProgress && !preserveProgress) || namedSymbol.IsCancellationToken;
+        return (namedSymbol.IsIProgress && !preserveProgress) || (isArgument ? namedSymbol.IsCancellationToken : !preserveCancellationToken);
     }
 
     private bool ShouldRemoveArgument(ISymbol symbol, bool isNegated = false) => symbol switch
@@ -1912,10 +1924,10 @@ private bool ShouldRemoveType(ITypeSymbol symbol)
         IPropertySymbol { IsCancellationRequested: true } => !isNegated,
         IMethodSymbol ms =>
             IsSpecialMethod(ms) is SpecialMethod.None or SpecialMethod.Drop
-                && ((ShouldRemoveType(ms.ReturnType) && ms.MethodKind != MethodKind.LocalFunction)
-                    || (ms.ReceiverType is { } receiver && ShouldRemoveType(receiver)))
+                && ((ShouldRemoveType(ms.ReturnType, isArgument: true) && ms.MethodKind != MethodKind.LocalFunction)
+                    || (ms.ReceiverType is { } receiver && ShouldRemoveType(receiver, isArgument: true)))
                 && !HasSyncMethod(ms),
-        _ => ShouldRemoveType(GetReturnType(symbol)),
+        _ => ShouldRemoveType(GetReturnType(symbol), isArgument: true),
     };
 
     private bool PreProcess(
@@ -2155,6 +2167,12 @@ private SyncOnlyAttributeContext ProcessSyncOnlyAttributes(SyntaxTriviaList synt
 
     private ISymbol? GetSymbol(SyntaxNode node) => semanticModel.GetSymbolInfo(node).Symbol;
 
+    private bool IsEnumeratorCancellationAttribute(AttributeSyntax attributeSyntax)
+    {
+        var type = GetSymbol(attributeSyntax)?.ContainingType;
+        return type?.IsEnumeratorCancellationAttribute == true;
+    }
+
     private bool CanDropDeclaration(LocalDeclarationStatementSyntax local)
     {
         var symbol = GetSymbol(local.Declaration.Type);
@@ -2203,7 +2221,7 @@ private TypeSyntax ProcessSyntaxUsingSymbol(TypeSyntax typeSyntax)
     };
 
     private bool ChecksIfNegatedIsCancellationRequested(ExpressionSyntax condition)
-        => RemoveParentheses(condition) is PrefixUnaryExpressionSyntax { OperatorToken.RawKind: (int)SyntaxKind.ExclamationToken } pe
+        => !preserveCancellationToken && RemoveParentheses(condition) is PrefixUnaryExpressionSyntax { OperatorToken.RawKind: (int)SyntaxKind.ExclamationToken } pe
         && RemoveParentheses(pe.Operand) is MemberAccessExpressionSyntax { Name.Identifier.ValueText: IsCancellationRequested } mae
         && GetSymbol(mae) is { ContainingType.IsCancellationToken: true };
 

src/Zomp.SyncMethodGenerator/Extensions.cs

@@ -52,6 +52,12 @@ internal static class Extensions
             ContainingNamespace: { Name: Threading, ContainingNamespace: { Name: System, ContainingNamespace.IsGlobalNamespace: true } }
         };
 
+        public bool IsEnumeratorCancellationAttribute => symbol is
+        {
+            Name: EnumeratorCancellationAttribute, IsGenericType: false,
+            ContainingNamespace: { Name: CompilerServices, ContainingNamespace: { Name: Runtime, ContainingNamespace: { Name: System, ContainingNamespace.IsGlobalNamespace: true } } }
+        };
+
         public bool IsIProgress => symbol is
         {
             Name: IProgress, IsGenericType: true,
@@ -118,6 +124,7 @@ internal static class Extensions
     // Type names
     private const string Enumerator = nameof(Span<>.Enumerator);
     private const string CancellationToken = nameof(global::System.Threading.CancellationToken);
+    private const string EnumeratorCancellationAttribute = nameof(global::System.Runtime.CompilerServices.EnumeratorCancellationAttribute);
     private const string IAsyncEnumerable = nameof(IAsyncEnumerable<>);
     private const string IAsyncEnumerator = nameof(IAsyncEnumerator<>);
     private const string Task = nameof(Task<>);

src/Zomp.SyncMethodGenerator/SourceGenerationHelper.cs

@@ -36,6 +36,11 @@ internal class {{SyncMethodSourceGenerator.CreateSyncVersionAttribute}} : System
         /// Gets or sets a value indicating whether <see cref="System.IProgress{T}"/> parameters will be preserved in the generated code. False by default.
         /// </summary>
         public bool {{SyncMethodSourceGenerator.PreserveProgress}} { get; set; }
+
+        /// <summary>
+        /// Gets or sets a value indicating whether <see cref="System.Threading.CancellationToken"/> parameters will be preserved in the generated code. False by default.
+        /// </summary>
+        public bool {{SyncMethodSourceGenerator.PreserveCancellationToken}} { get; set; }
     }
 #endif
 }

src/Zomp.SyncMethodGenerator/SyncMethodSourceGenerator.cs

@@ -21,6 +21,7 @@ public class SyncMethodSourceGenerator : IIncrementalGenerator
 
     internal const string OmitNullableDirective = "OmitNullableDirective";
     internal const string PreserveProgress = "PreserveProgress";
+    internal const string PreserveCancellationToken = "PreserveCancellationToken";
 
     /// <inheritdoc/>
     public void Initialize(IncrementalGeneratorInitializationContext context)
@@ -237,9 +238,10 @@ static string BuildClassName(MethodParentDeclaration c)
         disableNullable |= explicitDisableNullable;
 
         var preserveProgress = syncMethodGeneratorAttributeData.NamedArguments.FirstOrDefault(c => c.Key == PreserveProgress) is { Value.Value: true };
+        var preserveCancellationToken = syncMethodGeneratorAttributeData.NamedArguments.FirstOrDefault(c => c.Key == PreserveCancellationToken) is { Value.Value: true };
 
         var toVisit = extensionParent ?? (SyntaxNode)methodDeclarationSyntax;
-        var rewriter = new AsyncToSyncRewriter(context.SemanticModel, disableNullable, preserveProgress, methodDeclarationSyntax);
+        var rewriter = new AsyncToSyncRewriter(context.SemanticModel, disableNullable, preserveProgress, preserveCancellationToken, methodDeclarationSyntax);
         var sn = rewriter.Visit(toVisit);
         var content = sn.ToFullString();
 

tests/Generator.Tests/PreserveCancellationTokenTests.cs

@@ -0,0 +1,50 @@
+namespace Generator.Tests;
+
+public class PreserveCancellationTokenTests
+{
+    [Fact]
+    public Task PreserveIsCancellationRequestedAndThrowIfCancellationRequested() => """
+[CreateSyncVersion(PreserveCancellationToken = true)]
+async Task MethodAsync(CancellationToken ct = default)
+{
+    if (!ct.IsCancellationRequested)
+    {
+        ct.ThrowIfCancellationRequested();
+    }
+}
+""".Verify();
+
+    [Fact]
+    public Task RemoveCancellationTokenForSyncMethods() => """
+[CreateSyncVersion(PreserveCancellationToken = true)]
+async Task MethodAsync(CancellationToken ct = default)
+{
+    Stream.Null.FlushAsync(ct);
+}
+""".Verify();
+
+    [Theory]
+    [InlineData("[EnumeratorCancellation]")]
+    [InlineData("[EnumeratorCancellationAttribute]")]
+    [InlineData("[System.Runtime.CompilerServices.EnumeratorCancellation]")]
+    [InlineData("[System.Runtime.CompilerServices.EnumeratorCancellationAttribute]")]
+    [InlineData("[global::System.Runtime.CompilerServices.EnumeratorCancellationAttribute]")]
+    public Task RemoveEnumeratorCancellationAttribute(string attribute) => $$"""
+[CreateSyncVersion(PreserveCancellationToken = true)]
+async IAsyncEnumerable<int> FibonacciAsync({{attribute}} CancellationToken ct = default)
+{
+    var f0 = 0;
+    var f1 = 1;
+    yield return f0;
+    yield return f1;
+    while (!ct.IsCancellationRequested)
+    {
+        var fn = f0 + f1;
+        yield return fn;
+        await Task.Yield();
+        f0 = f1;
+        f1 = fn;
+    }
+}
+""".Verify(disableUnique: true);
+}

tests/Generator.Tests/Snapshots/PreserveCancellationTokenTests.PreserveIsCancellationRequestedAndThrowIfCancellationRequested#g.verified.cs

@@ -0,0 +1,8 @@
+//HintName: Test.Class.MethodAsync.g.cs
+void Method(global::System.Threading.CancellationToken ct = default)
+{
+    if (!ct.IsCancellationRequested)
+    {
+        ct.ThrowIfCancellationRequested();
+    }
+}

tests/Generator.Tests/Snapshots/PreserveCancellationTokenTests.RemoveCancellationTokenForSyncMethods#g.verified.cs

@@ -0,0 +1,5 @@
+//HintName: Test.Class.MethodAsync.g.cs
+void Method(global::System.Threading.CancellationToken ct = default)
+{
+    global::System.IO.Stream.Null.Flush();
+}

tests/Generator.Tests/Snapshots/PreserveCancellationTokenTests.RemoveEnumeratorCancellationAttribute#FibonacciAsync.g.verified.cs

@@ -0,0 +1,15 @@
+//HintName: Test.Class.FibonacciAsync.g.cs
+global::System.Collections.Generic.IEnumerable<int> Fibonacci(global::System.Threading.CancellationToken ct = default)
+{
+    var f0 = 0;
+    var f1 = 1;
+    yield return f0;
+    yield return f1;
+    while (!ct.IsCancellationRequested)
+    {
+        var fn = f0 + f1;
+        yield return fn;
+        f0 = f1;
+        f1 = fn;
+    }
+}