From ed1b2ca043908520dc1bdbb65135e831ce0252b0 Mon Sep 17 00:00:00 2001
From: Jan Nedbal <janedbal@gmail.com>
Date: Thu, 7 Sep 2023 13:20:37 +0200
Subject: [PATCH] Support function and closure

---
 .../ForbidReturnValueInYieldingMethodRule.php | 33 ++++++++++++++++---
 .../normal.php                                | 20 +++++++++++
 2 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/src/Rule/ForbidReturnValueInYieldingMethodRule.php b/src/Rule/ForbidReturnValueInYieldingMethodRule.php
index a73108c..356f57d 100644
--- a/src/Rule/ForbidReturnValueInYieldingMethodRule.php
+++ b/src/Rule/ForbidReturnValueInYieldingMethodRule.php
@@ -5,15 +5,19 @@
 use Generator;
 use PhpParser\Node;
 use PHPStan\Analyser\Scope;
+use PHPStan\Node\ClosureReturnStatementsNode;
 use PHPStan\Node\MethodReturnStatementsNode;
+use PHPStan\Node\ReturnStatementsNode;
 use PHPStan\Reflection\ParametersAcceptorSelector;
 use PHPStan\Rules\Rule;
 use PHPStan\Rules\RuleError;
 use PHPStan\Rules\RuleErrorBuilder;
+use PHPStan\Type\MixedType;
+use PHPStan\Type\Type;
 use function in_array;
 
 /**
- * @implements Rule<MethodReturnStatementsNode>
+ * @implements Rule<ReturnStatementsNode>
  */
 class ForbidReturnValueInYieldingMethodRule implements Rule
 {
@@ -27,11 +31,11 @@ public function __construct(bool $reportRegardlessOfReturnType)
 
     public function getNodeType(): string
     {
-        return MethodReturnStatementsNode::class;
+        return ReturnStatementsNode::class;
     }
 
     /**
-     * @param MethodReturnStatementsNode $node
+     * @param ReturnStatementsNode $node
      * @return list<RuleError>
      */
     public function processNode(
@@ -44,7 +48,7 @@ public function processNode(
         }
 
         if (!$this->reportRegardlessOfReturnType) {
-            $methodReturnType = ParametersAcceptorSelector::selectSingle($node->getMethodReflection()->getVariants())->getReturnType();
+            $methodReturnType = $this->getReturnType($node, $scope);
 
             if (in_array(Generator::class, $methodReturnType->getObjectClassNames(), true)) {
                 return [];
@@ -61,7 +65,11 @@ public function processNode(
                     ? 'this approach is denied'
                     : 'but this method is not marked to return Generator';
 
-                $errors[] = RuleErrorBuilder::message("Returned value from yielding method can be accessed only via Generator::getReturn, $suffix.")
+                $callType = $node instanceof MethodReturnStatementsNode // @phpstan-ignore-line ignore bc promise
+                    ? 'method'
+                    : 'function';
+
+                $errors[] = RuleErrorBuilder::message("Returned value from yielding $callType can be accessed only via Generator::getReturn, $suffix.")
                     ->line($returnNode->getLine())
                     ->build();
             }
@@ -70,4 +78,19 @@ public function processNode(
         return $errors;
     }
 
+    private function getReturnType(ReturnStatementsNode $node, Scope $scope): Type
+    {
+        $methodReflection = $scope->getFunction();
+
+        if ($node instanceof ClosureReturnStatementsNode) { // @phpstan-ignore-line ignore bc promise
+            return $scope->getFunctionType($node->getClosureExpr()->getReturnType(), false, false);
+        }
+
+        if ($methodReflection !== null) {
+            return ParametersAcceptorSelector::selectSingle($methodReflection->getVariants())->getReturnType();
+        }
+
+        return new MixedType();
+    }
+
 }
diff --git a/tests/Rule/data/ForbidReturnValueInYieldingMethodRule/normal.php b/tests/Rule/data/ForbidReturnValueInYieldingMethodRule/normal.php
index 3450d8d..b047080 100644
--- a/tests/Rule/data/ForbidReturnValueInYieldingMethodRule/normal.php
+++ b/tests/Rule/data/ForbidReturnValueInYieldingMethodRule/normal.php
@@ -45,3 +45,23 @@ public function returnGenerator(): \Generator {
         return;
     }
 }
+
+function returnIterable(): iterable {
+    yield 1;
+    return 2; // error: Returned value from yielding function can be accessed only via Generator::getReturn, but this method is not marked to return Generator.
+}
+
+function returnGenerator(): \Generator {
+    yield 1;
+    return 2;
+}
+
+function (): iterable {
+    yield 1;
+    return 2; // error: Returned value from yielding function can be accessed only via Generator::getReturn, but this method is not marked to return Generator.
+};
+
+function (): \Generator {
+    yield 1;
+    return 2;
+};