[Defluent] Do not remove return from parameter and property on ReturnNewFluentChainMethodCallToNonFluentRector (#589)

* [Defluent] Do not remove return on ReturnNewFluentChainMethodCallToNonFluentRector

* [Defluent] Do not remove return on ReturnNewFluentChainMethodCallToNonFluentRector

* eol

* eol

* from parameter

* need to check on property as well

* Fixed 🎉

* final touch: clean up

rules-tests/Defluent/Rector/Return_/ReturnNewFluentChainMethodCallToNonFluentRector/Fixture/do_not_remove_return_from_parameter.php.inc

@@ -0,0 +1,32 @@
+<?php
+
+namespace Rector\Tests\Defluent\Rector\Return_\ReturnNewFluentChainMethodCallToNonFluentRector\Fixture;
+
+use Rector\Tests\Defluent\Rector\Return_\ReturnNewFluentChainMethodCallToNonFluentRector\Source\FluentInterfaceClass;
+
+class DoNotRemoveReturnFromParameter
+{
+    public function someFunction(FluentInterfaceClass $fluentInterfaceClass)
+    {
+        return $fluentInterfaceClass->someFunction()->otherFunction();
+    }
+}
+
+?>
+-----
+<?php
+
+namespace Rector\Tests\Defluent\Rector\Return_\ReturnNewFluentChainMethodCallToNonFluentRector\Fixture;
+
+use Rector\Tests\Defluent\Rector\Return_\ReturnNewFluentChainMethodCallToNonFluentRector\Source\FluentInterfaceClass;
+
+class DoNotRemoveReturnFromParameter
+{
+    public function someFunction(FluentInterfaceClass $fluentInterfaceClass)
+    {
+        $fluentInterfaceClass->someFunction();
+        return $fluentInterfaceClass->otherFunction();
+    }
+}
+
+?>

rules-tests/Defluent/Rector/Return_/ReturnNewFluentChainMethodCallToNonFluentRector/Fixture/do_not_remove_return_from_property.php.inc

@@ -0,0 +1,46 @@
+<?php
+
+namespace Rector\Tests\Defluent\Rector\Return_\ReturnNewFluentChainMethodCallToNonFluentRector\Fixture;
+
+use Rector\Tests\Defluent\Rector\Return_\ReturnNewFluentChainMethodCallToNonFluentRector\Source\FluentInterfaceClass;
+
+class DoNotRemoveReturnFromProperty
+{
+    private FluentInterfaceClass $property;
+
+    public function __construct(FluentInterfaceClass $fluentInterfaceClass)
+    {
+        $this->property = $fluentInterfaceClass;
+    }
+
+    public function someFunction()
+    {
+        return $this->property->someFunction()->otherFunction();
+    }
+}
+
+?>
+-----
+<?php
+
+namespace Rector\Tests\Defluent\Rector\Return_\ReturnNewFluentChainMethodCallToNonFluentRector\Fixture;
+
+use Rector\Tests\Defluent\Rector\Return_\ReturnNewFluentChainMethodCallToNonFluentRector\Source\FluentInterfaceClass;
+
+class DoNotRemoveReturnFromProperty
+{
+    private FluentInterfaceClass $property;
+
+    public function __construct(FluentInterfaceClass $fluentInterfaceClass)
+    {
+        $this->property = $fluentInterfaceClass;
+    }
+
+    public function someFunction()
+    {
+        $this->property->someFunction();
+        return $this->property->otherFunction();
+    }
+}
+
+?>

rules/Defluent/Rector/Return_/ReturnNewFluentChainMethodCallToNonFluentRector.php

@@ -85,10 +85,21 @@ CODE_SAMPLE
             return null;
         }
 
-        $this->fluentNodeRemover->removeCurrentNode($node);
-        $this->addNodesAfterNode($assignAndRootExprAndNodesToAdd->getNodesToAdd(), $node);
+        $nodesToAdd = $assignAndRootExprAndNodesToAdd->getNodesToAdd();
+        $lastNodeToAdd = end($nodesToAdd);
 
-        return null;
+        if (! $lastNodeToAdd) {
+            return null;
+        }
+
+        if (! $lastNodeToAdd instanceof Return_) {
+            $nodesToAdd[array_key_last($nodesToAdd)] = new Return_($lastNodeToAdd);
+        }
+
+        $this->fluentNodeRemover->removeCurrentNode($node);
+        $this->addNodesAfterNode($nodesToAdd, $node);
+
+        return $node;
     }
 
     private function matchReturnMethodCall(Return_ $return): ?Expr