From 0d90f9f748210959695c027b842c2790740e12db Mon Sep 17 00:00:00 2001
From: Tomas Votruba <tomas.vot@gmail.com>
Date: Thu, 4 Sep 2025 17:38:35 +0200
Subject: [PATCH 1/2] add fixture to TypeWillReturnCallableArrowFunctionRector

---
 .../Fixture/with_callback.php.inc             | 37 +++++++++++++++++++
 1 file changed, 37 insertions(+)
 create mode 100644 rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Fixture/with_callback.php.inc

diff --git a/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Fixture/with_callback.php.inc b/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Fixture/with_callback.php.inc
new file mode 100644
index 00000000..370555a7
--- /dev/null
+++ b/rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Fixture/with_callback.php.inc
@@ -0,0 +1,37 @@
+<?php
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Fixture;
+
+use PHPUnit\Framework\TestCase;
+use Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Source\SomeMockedClass;
+
+final class FillKnownParamType extends TestCase
+{
+    public function test($value): void
+    {
+        $this->createMock(SomeMockedClass::class)
+            ->method('someMethod')
+            ->with($this->callback(fn ($name) => $value));
+    }
+}
+
+?>
+-----
+<?php
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Fixture;
+
+use PHPUnit\Framework\TestCase;
+use Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Source\SomeMockedClass;
+
+final class FillKnownParamType extends TestCase
+{
+    public function test($value): void
+    {
+        $this->createMock(SomeMockedClass::class)
+            ->method('someMethod')
+            ->with($this->callback(fn (string $name): int => $value));
+    }
+}
+
+?>

From ca6f86b0dfe4cf6b8758ac536e7f9dfd2f58d3e5 Mon Sep 17 00:00:00 2001
From: Tomas Votruba <tomas.vot@gmail.com>
Date: Thu, 4 Sep 2025 17:44:20 +0200
Subject: [PATCH 2/2] add with() + callback() support to
 TypeWillReturnCallableArrowFunctionRector

---
 ...eWillReturnCallableArrowFunctionRector.php | 42 ++++++++++++++-----
 1 file changed, 32 insertions(+), 10 deletions(-)

diff --git a/rules/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector.php b/rules/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector.php
index fd99bb5e..e415d9cc 100644
--- a/rules/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector.php
+++ b/rules/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector.php
@@ -146,13 +146,8 @@ public function refactor(Node $node): ?Class_
                 return null;
             }
 
-            if (! $this->isName($node->name, self::WILL_RETURN_CALLBACK)) {
-                return null;
-            }
-
-            $innerArg = $node->getArgs()[0]
-                ->value;
-            if (! $innerArg instanceof ArrowFunction && ! $innerArg instanceof Closure) {
+            $innerClosure = $this->matchInnerClosure($node);
+            if (! $innerClosure instanceof Node) {
                 return null;
             }
 
@@ -205,7 +200,7 @@ public function refactor(Node $node): ?Class_
                 return null;
             }
 
-            foreach ($innerArg->params as $key => $param) {
+            foreach ($innerClosure->params as $key => $param) {
                 // avoid typing variadic parameters
                 if ($param->variadic) {
                     continue;
@@ -239,14 +234,14 @@ public function refactor(Node $node): ?Class_
                 $hasChanged = true;
             }
 
-            if (! $innerArg->returnType instanceof Node) {
+            if (! $innerClosure->returnType instanceof Node) {
                 $returnTypeNode = $this->staticTypeMapper->mapPHPStanTypeToPhpParserNode(
                     $parameterTypesAndReturnType->getReturnType(),
                     TypeKind::RETURN
                 );
 
                 if ($returnTypeNode instanceof Node) {
-                    $innerArg->returnType = $returnTypeNode;
+                    $innerClosure->returnType = $returnTypeNode;
                     $hasChanged = true;
                 }
             }
@@ -259,6 +254,33 @@ public function refactor(Node $node): ?Class_
         return $node;
     }
 
+    public function matchInnerClosure(MethodCall $methodCall): null|ArrowFunction|Closure
+    {
+        if ($this->isName($methodCall->name, 'with')) {
+            // special case for nested callback
+            $withFirstArg = $methodCall->getArgs()[0];
+
+            if ($withFirstArg->value instanceof MethodCall) {
+                $nestedMethodCall = $withFirstArg->value;
+                if ($this->isName($nestedMethodCall->name, 'callback')) {
+                    $nestedArg = $nestedMethodCall->getArgs()[0];
+                    if ($nestedArg->value instanceof ArrowFunction || $nestedArg->value instanceof Closure) {
+                        return $nestedArg->value;
+                    }
+                }
+            }
+        }
+
+        if ($this->isName($methodCall->name, self::WILL_RETURN_CALLBACK)) {
+            $innerArg = $methodCall->getArgs()[0];
+            if ($innerArg->value instanceof ArrowFunction || $innerArg->value instanceof Closure) {
+                return $innerArg->value;
+            }
+        }
+
+        return null;
+    }
+
     /**
      * @param array<string, string> $propertyNameToMockedTypes
      */