From 393b2dbf422825db6c1ee0dd24c4ddc42d31f05c Mon Sep 17 00:00:00 2001
From: martincostello <martin@martincostello.com>
Date: Fri, 27 Oct 2023 13:52:09 +0100
Subject: [PATCH] Potentially resolve AOT warning

Attempt to resolve the AOT warning in `DelegatingComponent`.
---
 src/Polly.Core/Polly.Core.csproj              |  1 +
 .../Utils/Pipeline/CompositeComponent.cs      | 19 ++++++++++++++++---
 .../Pipeline/PipelineComponentFactory.cs      | 14 ++++++++++++++
 test/Polly.AotTest/Program.cs                 | 10 ++++++++++
 4 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/src/Polly.Core/Polly.Core.csproj b/src/Polly.Core/Polly.Core.csproj
index a7d1d4b160c..96c4e532503 100644
--- a/src/Polly.Core/Polly.Core.csproj
+++ b/src/Polly.Core/Polly.Core.csproj
@@ -21,6 +21,7 @@
 
   <ItemGroup>
     <Using Include="Polly.Utils" />
+    <InternalsVisibleToTest Include="Polly.AotTest" />
     <InternalsVisibleToTest Include="Polly.Core.Tests" />
     <InternalsVisibleToTest Include="Polly.Testing" />
     <InternalsVisibleToTest Include="Polly.TestUtils" />
diff --git a/src/Polly.Core/Utils/Pipeline/CompositeComponent.cs b/src/Polly.Core/Utils/Pipeline/CompositeComponent.cs
index abaa8692f4f..e94dda2f9e0 100644
--- a/src/Polly.Core/Utils/Pipeline/CompositeComponent.cs
+++ b/src/Polly.Core/Utils/Pipeline/CompositeComponent.cs
@@ -133,20 +133,33 @@ internal override ValueTask<Outcome<TResult>> ExecuteCore<TResult, TState>(
             ResilienceContext context,
             TState state)
         {
+            // Custom state object is used to cast the callback and state to prevent infinite
+            // generic type recursion warning IL3054 when referenced in a native AoT application.
+            // See https://github.com/App-vNext/Polly/issues/1732 for further context.
             return _component.ExecuteCore(
-                static (context, state) =>
+                static (context, wrapper) =>
                 {
+                    var callback = (Func<ResilienceContext, TState, ValueTask<Outcome<TResult>>>)wrapper.Callback;
+                    var state = (TState)wrapper.State;
+
                     if (context.CancellationToken.IsCancellationRequested)
                     {
                         return Outcome.FromExceptionAsValueTask<TResult>(new OperationCanceledException(context.CancellationToken).TrySetStackTrace());
                     }
 
-                    return state.Next!.ExecuteCore(state.callback, context, state.state);
+                    return wrapper.Next.ExecuteCore(callback, context, state);
                 },
                 context,
-                (Next, callback, state));
+                new StateWrapper(Next!, callback, state!));
         }
 
         public override ValueTask DisposeAsync() => default;
+
+        private struct StateWrapper(PipelineComponent next, object callback, object state)
+        {
+            public PipelineComponent Next = next;
+            public object Callback = callback;
+            public object State = state;
+        }
     }
 }
diff --git a/src/Polly.Core/Utils/Pipeline/PipelineComponentFactory.cs b/src/Polly.Core/Utils/Pipeline/PipelineComponentFactory.cs
index a671dc7e0f1..9dede50d372 100644
--- a/src/Polly.Core/Utils/Pipeline/PipelineComponentFactory.cs
+++ b/src/Polly.Core/Utils/Pipeline/PipelineComponentFactory.cs
@@ -14,10 +14,24 @@ internal static class PipelineComponentFactory
 
     public static PipelineComponent WithDisposableCallbacks(PipelineComponent component, IEnumerable<Action> callbacks)
     {
+#if NET6_0_OR_GREATER
+        if (callbacks.TryGetNonEnumeratedCount(out var count))
+        {
+            if (count == 0)
+            {
+                return component;
+            }
+        }
+        else if (!callbacks.Any())
+        {
+            return component;
+        }
+#else
         if (!callbacks.Any())
         {
             return component;
         }
+#endif
 
         return new ComponentWithDisposeCallbacks(component, callbacks.ToList());
     }
diff --git a/test/Polly.AotTest/Program.cs b/test/Polly.AotTest/Program.cs
index 58191d4533e..b4b09982d6b 100644
--- a/test/Polly.AotTest/Program.cs
+++ b/test/Polly.AotTest/Program.cs
@@ -40,6 +40,16 @@
 p.ExecuteCore<int, int>(state => default, default);
 p.ExecuteCore<int, int>(state => default, default);
 
+var p1 = Polly.Utils.Pipeline.PipelineComponent.Empty;
+var p2 = Polly.Utils.Pipeline.PipelineComponent.Empty;
+var p3 = Polly.Utils.Pipeline.CompositeComponent.Create([p1, p2], null!, null!);
+var p4 = new Polly.Utils.Pipeline.BridgeComponent(null!);
+
+await p1.ExecuteCore<int, int>((state, context) => default, default!, default);
+await p2.ExecuteCore<int, int>((state, context) => default, default!, default);
+await p3.ExecuteCore<int, int>((state, context) => default, default!, default);
+await p4.ExecuteCore<int, int>((state, context) => default, default!, default);
+
 var app = builder.Build();
 
 var sampleTodos = new Todo[] {