Remove UseProvidedBody string matching from DelegateBodySyntaxExtractor

EasySourceGenerators.Generators/DataBuilding/DataMethodBodyBuilders.cs

@@ -38,6 +38,9 @@ public IMethodBodyBuilderStage4<TParam1, T> WithParameter<TParam1>() =>
 
 public record DataMethodBodyBuilderStage4<TParam1, TReturnType>(BodyGenerationData Data) : IMethodBodyBuilderStage4<TParam1, TReturnType>
 {
+    public IMethodBodyBuilderStage5WithConstants<TParam1, TReturnType, TConstants> WithCompileTimeConstants<TConstants>(Func<TConstants> compileTimeConstantsFactory) =>
+        new DataMethodBodyBuilderStage5WithConstants<TParam1, TReturnType, TConstants>(Data with { CompileTimeConstants = compileTimeConstantsFactory() });
+
     public IMethodBodyGenerator UseProvidedBody(Func<TParam1, TReturnType> body) => new DataMethodBodyGenerator(Data with { RuntimeDelegateBody = body });
 
     public IMethodBodyGenerator BodyReturningConstant(Func<TReturnType> constantValueFactory) =>
@@ -46,6 +49,9 @@ public IMethodBodyGenerator BodyReturningConstant(Func<TReturnType> constantValu
 
 public record DataMethodBodyBuilderStage4NoArg<TReturnType>(BodyGenerationData Data) : IMethodBodyBuilderStage4NoArg<TReturnType>
 {
+    public IMethodBodyBuilderStage5NoArgWithConstants<TReturnType, TConstants> WithCompileTimeConstants<TConstants>(Func<TConstants> compileTimeConstantsFactory) =>
+        new DataMethodBodyBuilderStage5NoArgWithConstants<TReturnType, TConstants>(Data with { CompileTimeConstants = compileTimeConstantsFactory() });
+
     public IMethodBodyGenerator UseProvidedBody(Func<TReturnType> body) => new DataMethodBodyGenerator(Data with { RuntimeDelegateBody = body });
 
     public IMethodBodyGenerator BodyReturningConstant(Func<TReturnType> constantValueFactory) =>
@@ -54,10 +60,42 @@ public IMethodBodyGenerator BodyReturningConstant(Func<TReturnType> constantValu
 
 public record DataMethodBodyBuilderStage4ReturnVoid<TParam1>(BodyGenerationData BodyGenerationData) : IMethodBodyBuilderStage4ReturnVoid<TParam1>
 {
+    public IMethodBodyBuilderStage5ReturnVoidWithConstants<TParam1, TConstants> WithCompileTimeConstants<TConstants>(Func<TConstants> compileTimeConstantsFactory) =>
+        new DataMethodBodyBuilderStage5ReturnVoidWithConstants<TParam1, TConstants>(BodyGenerationData with { CompileTimeConstants = compileTimeConstantsFactory() });
+
     public IMethodBodyGenerator UseProvidedBody(Action<TParam1> body) => new DataMethodBodyGenerator(BodyGenerationData with { RuntimeDelegateBody = body });
 }
 
 public record DataMethodBodyBuilderStage4ReturnVoidNoArg(BodyGenerationData BodyGenerationData) : IMethodBodyBuilderStage4ReturnVoidNoArg
 {
+    public IMethodBodyBuilderStage5ReturnVoidNoArgWithConstants<TConstants> WithCompileTimeConstants<TConstants>(Func<TConstants> compileTimeConstantsFactory) =>
+        new DataMethodBodyBuilderStage5ReturnVoidNoArgWithConstants<TConstants>(BodyGenerationData with { CompileTimeConstants = compileTimeConstantsFactory() });
+
     public IMethodBodyGenerator UseProvidedBody(Action body) => new DataMethodBodyGenerator(BodyGenerationData with { RuntimeDelegateBody = body });
+}
+
+public record DataMethodBodyBuilderStage5WithConstants<TParam1, TReturnType, TConstants>(BodyGenerationData Data) : IMethodBodyBuilderStage5WithConstants<TParam1, TReturnType, TConstants>
+{
+    public IMethodBodyGenerator UseProvidedBody(Func<TConstants, TParam1, TReturnType> body) => new DataMethodBodyGenerator(Data with { RuntimeDelegateBody = body });
+
+    public IMethodBodyGenerator BodyReturningConstant(Func<TConstants, TReturnType> constantValueFactory) =>
+        new DataMethodBodyGenerator(Data with { ReturnConstantValueFactory = constantValueFactory });
+}
+
+public record DataMethodBodyBuilderStage5NoArgWithConstants<TReturnType, TConstants>(BodyGenerationData Data) : IMethodBodyBuilderStage5NoArgWithConstants<TReturnType, TConstants>
+{
+    public IMethodBodyGenerator UseProvidedBody(Func<TConstants, TReturnType> body) => new DataMethodBodyGenerator(Data with { RuntimeDelegateBody = body });
+
+    public IMethodBodyGenerator BodyReturningConstant(Func<TConstants, TReturnType> constantValueFactory) =>
+        new DataMethodBodyGenerator(Data with { ReturnConstantValueFactory = constantValueFactory });
+}
+
+public record DataMethodBodyBuilderStage5ReturnVoidWithConstants<TParam1, TConstants>(BodyGenerationData Data) : IMethodBodyBuilderStage5ReturnVoidWithConstants<TParam1, TConstants>
+{
+    public IMethodBodyGenerator UseProvidedBody(Action<TConstants, TParam1> body) => new DataMethodBodyGenerator(Data with { RuntimeDelegateBody = body });
+}
+
+public record DataMethodBodyBuilderStage5ReturnVoidNoArgWithConstants<TConstants>(BodyGenerationData Data) : IMethodBodyBuilderStage5ReturnVoidNoArgWithConstants<TConstants>
+{
+    public IMethodBodyGenerator UseProvidedBody(Action<TConstants> body) => new DataMethodBodyGenerator(Data with { RuntimeDelegateBody = body });
 }

EasySourceGenerators.Generators/IncrementalGenerators/BodyGenerationDataExtractor.cs

@@ -15,6 +15,8 @@ internal static class BodyGenerationDataExtractor
     /// Checks for <c>ReturnConstantValueFactory</c> first, then <c>RuntimeDelegateBody</c>.
     /// Returns a <see cref="FluentBodyResult"/> with the extracted value, or <c>null</c> return value
     /// if neither factory nor body are present.
+    /// Sets <see cref="FluentBodyResult.HasDelegateBody"/> when <c>RuntimeDelegateBody</c> is present,
+    /// indicating that the delegate body source code should be extracted from the syntax tree.
     /// </summary>
     internal static FluentBodyResult Extract(object methodResult, bool isVoidReturnType)
     {
@@ -26,23 +28,35 @@ internal static FluentBodyResult Extract(object methodResult, bool isVoidReturnT
         {
             // The method returned something that isn't a DataMethodBodyGenerator.
             // This may happen when the fluent chain is incomplete (e.g., user returned an intermediate builder).
-            return new FluentBodyResult(null, isVoidReturnType);
+            return new FluentBodyResult(null, isVoidReturnType, HasDelegateBody: false);
         }
 
         object? bodyGenerationData = dataProperty.GetValue(methodResult);
         if (bodyGenerationData == null)
         {
-            return new FluentBodyResult(null, isVoidReturnType);
+            return new FluentBodyResult(null, isVoidReturnType, HasDelegateBody: false);
         }
 
         Type dataType = bodyGenerationData.GetType();
         PropertyInfo? returnTypeProperty = dataType.GetProperty("ReturnType");
         Type? dataReturnType = returnTypeProperty?.GetValue(bodyGenerationData) as Type;
         bool isVoid = dataReturnType == typeof(void);
 
+        bool hasDelegateBody = HasRuntimeDelegateBody(dataType, bodyGenerationData);
+
         return TryExtractFromConstantFactory(dataType, bodyGenerationData, isVoid)
-               ?? TryExtractFromRuntimeBody(dataType, bodyGenerationData, isVoid)
-               ?? new FluentBodyResult(null, isVoid);
+               ?? TryExtractFromRuntimeBody(dataType, bodyGenerationData, isVoid, hasDelegateBody)
+               ?? new FluentBodyResult(null, isVoid, hasDelegateBody);
+    }
+
+    /// <summary>
+    /// Checks whether <c>RuntimeDelegateBody</c> is set (non-null) in the body generation data.
+    /// </summary>
+    private static bool HasRuntimeDelegateBody(Type dataType, object bodyGenerationData)
+    {
+        PropertyInfo? runtimeBodyProperty = dataType.GetProperty("RuntimeDelegateBody");
+        Delegate? runtimeBody = runtimeBodyProperty?.GetValue(bodyGenerationData) as Delegate;
+        return runtimeBody != null;
     }
 
     /// <summary>
@@ -61,7 +75,7 @@ internal static FluentBodyResult Extract(object methodResult, bool isVoidReturnT
         }
 
         object? constantValue = constantFactory.DynamicInvoke();
-        return new FluentBodyResult(constantValue?.ToString(), isVoid);
+        return new FluentBodyResult(constantValue?.ToString(), isVoid, HasDelegateBody: false);
     }
 
     /// <summary>
@@ -72,7 +86,8 @@ internal static FluentBodyResult Extract(object methodResult, bool isVoidReturnT
     private static FluentBodyResult? TryExtractFromRuntimeBody(
         Type dataType,
         object bodyGenerationData,
-        bool isVoid)
+        bool isVoid,
+        bool hasDelegateBody)
     {
         PropertyInfo? runtimeBodyProperty = dataType.GetProperty("RuntimeDelegateBody");
         Delegate? runtimeBody = runtimeBodyProperty?.GetValue(bodyGenerationData) as Delegate;
@@ -85,10 +100,10 @@ internal static FluentBodyResult Extract(object methodResult, bool isVoidReturnT
         if (bodyParams.Length == 0)
         {
             object? bodyResult = runtimeBody.DynamicInvoke();
-            return new FluentBodyResult(bodyResult?.ToString(), isVoid);
+            return new FluentBodyResult(bodyResult?.ToString(), isVoid, hasDelegateBody);
         }
 
         // For delegates with parameters, we can't invoke at compile time without values
-        return new FluentBodyResult(null, isVoid);
+        return new FluentBodyResult(null, isVoid, hasDelegateBody);
     }
 }

EasySourceGenerators.Generators/IncrementalGenerators/DelegateBodySyntaxExtractor.cs

@@ -6,30 +6,25 @@
 namespace EasySourceGenerators.Generators.IncrementalGenerators;
 
 /// <summary>
-/// Extracts the delegate body source code from a <c>UseProvidedBody(...)</c> invocation
-/// within a generator method's syntax tree. The extracted body is re-indented to match
+/// Extracts the delegate body source code from the outermost invocation's lambda argument
+/// in a generator method's return expression. The extracted body is re-indented to match
 /// the target method body indentation (8 spaces).
 /// </summary>
 internal static class DelegateBodySyntaxExtractor
 {
     private const string MethodBodyIndent = "        ";
 
     /// <summary>
-    /// Attempts to find a <c>UseProvidedBody(...)</c> call in the given generator method syntax
-    /// and extract the lambda body. Returns <c>null</c> if no such call is found.
-    /// For expression lambdas, returns a single <c>return {expr};</c> line.
+    /// Attempts to find the lambda argument of the outermost invocation in the generator
+    /// method's return expression and extract the lambda body. Returns <c>null</c> if no
+    /// such lambda is found.
+    /// For expression lambdas, returns the expression text.
     /// For block lambdas, returns the block body re-indented to the method body level.
     /// </summary>
     internal static string? TryExtractDelegateBody(MethodDeclarationSyntax generatorMethodSyntax)
     {
-        InvocationExpressionSyntax? invocation = generatorMethodSyntax
-            .DescendantNodes()
-            .OfType<InvocationExpressionSyntax>()
-            .FirstOrDefault(inv =>
-                inv.Expression is MemberAccessExpressionSyntax memberAccess &&
-                memberAccess.Name.Identifier.Text == "UseProvidedBody");
-
-        if (invocation == null)
+        ExpressionSyntax? returnExpression = GetReturnExpression(generatorMethodSyntax);
+        if (returnExpression is not InvocationExpressionSyntax invocation)
         {
             return null;
         }
@@ -54,6 +49,29 @@ inv.Expression is MemberAccessExpressionSyntax memberAccess &&
         return null;
     }
 
+    /// <summary>
+    /// Gets the return expression from a generator method. Handles both expression-body
+    /// methods (<c>=&gt; expr</c>) and block-body methods (<c>{ return expr; }</c>).
+    /// Assumes the generator method has a simple structure with at most one return statement.
+    /// </summary>
+    private static ExpressionSyntax? GetReturnExpression(MethodDeclarationSyntax method)
+    {
+        if (method.ExpressionBody != null)
+        {
+            return method.ExpressionBody.Expression;
+        }
+
+        if (method.Body != null)
+        {
+            ReturnStatementSyntax? returnStatement = method.Body.Statements
+                .OfType<ReturnStatementSyntax>()
+                .FirstOrDefault();
+            return returnStatement?.Expression;
+        }
+
+        return null;
+    }
+
     /// <summary>
     /// Extracts the content of a block body (between <c>{</c> and <c>}</c>),
     /// determines the base indentation, and re-indents all lines to the method body level.

EasySourceGenerators.Generators/IncrementalGenerators/GeneratesMethodExecutionRuntime.cs

@@ -14,10 +14,13 @@ internal sealed record SwitchBodyData(
 
 /// <summary>
 /// Result extracted from <see cref="DataBuilding.BodyGenerationData"/> after executing a fluent body generator method.
+/// <see cref="HasDelegateBody"/> indicates that the generator used <c>UseProvidedBody</c>,
+/// signaling that the delegate body source code should be extracted from the syntax tree.
 /// </summary>
 internal sealed record FluentBodyResult(
     string? ReturnValue,
-    bool IsVoid);
+    bool IsVoid,
+    bool HasDelegateBody);
 
 /// <summary>
 /// Orchestrates the execution of generator methods at compile time.

EasySourceGenerators.Generators/IncrementalGenerators/GeneratesMethodGenerationPipeline.cs

@@ -96,9 +96,10 @@ private static string GenerateSourceForGroup(
     }
 
     /// <summary>
-    /// Generates source code from a fluent body pattern. First attempts to extract the delegate
-    /// body from a <c>UseProvidedBody</c> call in the syntax tree. If no such call is found,
-    /// falls back to executing the generator method and extracting the return value.
+    /// Generates source code from a fluent body pattern. Executes the generator method first
+    /// to obtain <see cref="FluentBodyResult"/>. If the result indicates a delegate body was
+    /// provided (via <see cref="FluentBodyResult.HasDelegateBody"/>), attempts to extract the
+    /// lambda body from the syntax tree. Otherwise, uses the runtime-evaluated return value.
     /// </summary>
     private static string GenerateFromFluentBodyPattern(
         SourceProductionContext context,
@@ -107,18 +108,6 @@ private static string GenerateFromFluentBodyPattern(
         INamedTypeSymbol containingType,
         Compilation compilation)
     {
-        string? delegateBody = DelegateBodySyntaxExtractor.TryExtractDelegateBody(methodInfo.Syntax);
-        if (delegateBody != null)
-        {
-            bool isVoidReturn = partialMethod.ReturnType.SpecialType == SpecialType.System_Void;
-            string bodyLines = FormatDelegateBodyForEmit(delegateBody, isVoidReturn);
-
-            return GeneratesMethodPatternSourceBuilder.GeneratePartialMethodWithBody(
-                containingType,
-                partialMethod,
-                bodyLines);
-        }
-
         (FluentBodyResult? result, string? error) = GeneratesMethodExecutionRuntime.ExecuteFluentBodyGeneratorMethod(
             methodInfo.Symbol,
             partialMethod,
@@ -134,10 +123,25 @@ private static string GenerateFromFluentBodyPattern(
             return string.Empty;
         }
 
+        if (result!.HasDelegateBody)
+        {
+            string? delegateBody = DelegateBodySyntaxExtractor.TryExtractDelegateBody(methodInfo.Syntax);
+            if (delegateBody != null)
+            {
+                bool isVoidReturn = partialMethod.ReturnType.SpecialType == SpecialType.System_Void;
+                string bodyLines = FormatDelegateBodyForEmit(delegateBody, isVoidReturn);
+
+                return GeneratesMethodPatternSourceBuilder.GeneratePartialMethodWithBody(
+                    containingType,
+                    partialMethod,
+                    bodyLines);
+            }
+        }
+
         return GeneratesMethodPatternSourceBuilder.GenerateSimplePartialMethod(
             containingType,
             partialMethod,
-            result!.ReturnValue);
+            result.ReturnValue);
     }
 
     /// <summary>