From f63a40f69b182e0b92a8d37ad0e2a3f02902c613 Mon Sep 17 00:00:00 2001
From: Matt Brown <github@muglug.com>
Date: Fri, 2 Oct 2020 00:47:42 -0400
Subject: [PATCH] Fix #4266 - prevent OOM when analysing closure unioned with
 invokable class

---
 .../Expression/Call/FunctionCallAnalyzer.php  |  8 +++--
 .../Expression/Call/MethodCallAnalyzer.php    | 10 +++++-
 .../ReturnTypeManipulationTest.php            | 32 +++++++++++++++++++
 3 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/src/Psalm/Internal/Analyzer/Statements/Expression/Call/FunctionCallAnalyzer.php b/src/Psalm/Internal/Analyzer/Statements/Expression/Call/FunctionCallAnalyzer.php
index 832705001..90c2243b1 100644
--- a/src/Psalm/Internal/Analyzer/Statements/Expression/Call/FunctionCallAnalyzer.php
+++ b/src/Psalm/Internal/Analyzer/Statements/Expression/Call/FunctionCallAnalyzer.php
@@ -695,7 +695,8 @@ class FunctionCallAnalyzer extends CallAnalyzer
                         $stmt,
                         $real_stmt,
                         $function_name,
-                        $context
+                        $context,
+                        $var_type_part
                     );
                 }
             }
@@ -746,7 +747,8 @@ class FunctionCallAnalyzer extends CallAnalyzer
         PhpParser\Node\Expr\FuncCall $stmt,
         PhpParser\Node\Expr\FuncCall $real_stmt,
         PhpParser\Node\Expr $function_name,
-        Context $context
+        Context $context,
+        Type\Atomic $atomic_type
     ) : void {
         $old_data_provider = $statements_analyzer->node_data;
 
@@ -772,6 +774,8 @@ class FunctionCallAnalyzer extends CallAnalyzer
             $statements_analyzer->addSuppressedIssues(['PossiblyInvalidMethodCall']);
         }
 
+        $statements_analyzer->node_data->setType($function_name, new Type\Union([$atomic_type]));
+
         \Psalm\Internal\Analyzer\Statements\Expression\Call\MethodCallAnalyzer::analyze(
             $statements_analyzer,
             $fake_method_call,
diff --git a/src/Psalm/Internal/Analyzer/Statements/Expression/Call/MethodCallAnalyzer.php b/src/Psalm/Internal/Analyzer/Statements/Expression/Call/MethodCallAnalyzer.php
index 9e1796419..767f9e472 100644
--- a/src/Psalm/Internal/Analyzer/Statements/Expression/Call/MethodCallAnalyzer.php
+++ b/src/Psalm/Internal/Analyzer/Statements/Expression/Call/MethodCallAnalyzer.php
@@ -47,7 +47,15 @@ class MethodCallAnalyzer extends \Psalm\Internal\Analyzer\Statements\Expression\
         $was_inside_use = $context->inside_use;
         $context->inside_use = true;
 
-        if (ExpressionAnalyzer::analyze($statements_analyzer, $stmt->var, $context) === false) {
+        $existing_stmt_var_type = null;
+
+        if (!$real_method_call) {
+            $existing_stmt_var_type = $statements_analyzer->node_data->getType($stmt->var);
+        }
+
+        if ($existing_stmt_var_type) {
+            $statements_analyzer->node_data->setType($stmt->var, $existing_stmt_var_type);
+        } elseif (ExpressionAnalyzer::analyze($statements_analyzer, $stmt->var, $context) === false) {
             return false;
         }
 
diff --git a/tests/FileManipulation/ReturnTypeManipulationTest.php b/tests/FileManipulation/ReturnTypeManipulationTest.php
index d9fcb7b0b..488dfc02d 100644
--- a/tests/FileManipulation/ReturnTypeManipulationTest.php
+++ b/tests/FileManipulation/ReturnTypeManipulationTest.php
@@ -575,6 +575,38 @@ class ReturnTypeManipulationTest extends FileManipulationTest
                 false,
                 true,
             ],
+            'dontOOM' => [
+                '<?php
+                    class FC {
+                        public function __invoke() : void {}
+                    }
+
+                    function foo(): string {
+                        if (rand(0, 1)) {
+                            $cb = new FC();
+                        } else {
+                            $cb = function() {};
+                        }
+                        $cb();
+                    }',
+                '<?php
+                    class FC {
+                        public function __invoke() : void {}
+                    }
+
+                    function foo(): string {
+                        if (rand(0, 1)) {
+                            $cb = new FC();
+                        } else {
+                            $cb = function() {};
+                        }
+                        $cb();
+                    }',
+                '7.3',
+                ['InvalidReturnType'],
+                false,
+                true,
+            ],
         ];
     }
 }