Overwrite memoized return type after sum-type candidate is calculated (#4805)

* Overwrite memoized return type after sum-type candidate is calculated

* Fix mismatched types

* Fix code style

src/Psalm/Internal/Analyzer/Statements/Expression/Call/Method/AtomicMethodCallAnalysisResult.php

@@ -79,4 +79,9 @@ class AtomicMethodCallAnalysisResult
      * @var list<\Psalm\Internal\MethodIdentifier>
      */
     public $too_few_arguments_method_ids = [];
+
+    /**
+     * @var bool
+     */
+    public $can_memoize = false;
 }

src/Psalm/Internal/Analyzer/Statements/Expression/Call/Method/ExistingAtomicMethodCallAnalyzer.php

@@ -200,8 +200,6 @@ class ExistingAtomicMethodCallAnalyzer extends CallAnalyzer
 
         $declaring_method_id = $codebase->methods->getDeclaringMethodId($method_id);
 
-        $can_memoize = false;
-
         $return_type_candidate = MethodCallReturnTypeFetcher::fetch(
             $statements_analyzer,
             $codebase,
@@ -253,7 +251,7 @@ class ExistingAtomicMethodCallAnalyzer extends CallAnalyzer
 
         if ($method_storage) {
             if (!$context->collect_mutations && !$context->collect_initializations) {
-                $can_memoize = MethodCallPurityAnalyzer::analyze(
+                $result->can_memoize = MethodCallPurityAnalyzer::analyze(
                     $statements_analyzer,
                     $codebase,
                     $stmt,
@@ -349,23 +347,6 @@ class ExistingAtomicMethodCallAnalyzer extends CallAnalyzer
             }
         }
 
-        if (!$args && $lhs_var_id) {
-            if ($config->memoize_method_calls || $can_memoize) {
-                $method_var_id = $lhs_var_id . '->' . $method_name_lc . '()';
-
-                if (isset($context->vars_in_scope[$method_var_id])) {
-                    $return_type_candidate = clone $context->vars_in_scope[$method_var_id];
-
-                    if ($can_memoize) {
-                        /** @psalm-suppress UndefinedPropertyAssignment */
-                        $stmt->pure = true;
-                    }
-                } else {
-                    $context->vars_in_scope[$method_var_id] = $return_type_candidate;
-                }
-            }
-        }
-
         if ($codebase->methods_to_rename) {
             $declaring_method_id = $codebase->methods->getDeclaringMethodId($method_id);
 

src/Psalm/Internal/Analyzer/Statements/Expression/Call/MethodCallAnalyzer.php

@@ -28,6 +28,7 @@ use Psalm\Type\Atomic\TNamedObject;
 use function count;
 use function is_string;
 use function array_reduce;
+use function strtolower;
 
 /**
  * @internal
@@ -201,6 +202,20 @@ class MethodCallAnalyzer extends \Psalm\Internal\Analyzer\Statements\Expression\
                 $possible_new_class_types[] = $context->vars_in_scope[$lhs_var_id];
             }
         }
+        if (!$stmt->args && $lhs_var_id && $stmt->name instanceof PhpParser\Node\Identifier) {
+            if ($codebase->config->memoize_method_calls || $result->can_memoize) {
+                $method_var_id = $lhs_var_id . '->' . strtolower($stmt->name->name) . '()';
+                if (isset($context->vars_in_scope[$method_var_id])) {
+                    $result->return_type = clone $context->vars_in_scope[$method_var_id];
+                    if ($result->can_memoize) {
+                        /** @psalm-suppress UndefinedPropertyAssignment */
+                        $stmt->pure = true;
+                    }
+                } elseif ($result->return_type !== null) {
+                    $context->vars_in_scope[$method_var_id] = $result->return_type;
+                }
+            }
+        }
 
         if (count($possible_new_class_types) > 0) {
             $class_type = array_reduce(

tests/Template/ClassTemplateTest.php

@@ -1476,6 +1476,43 @@ class ClassTemplateTest extends TestCase
                     '$a_or_b' => 'A|B',
                 ],
             ],
+            'doNotCombineTypesWhenMemoized' => [
+                '<?php
+                    class A {}
+                    class B {}
+
+                    /**
+                     * @template T
+                     */
+                    class C {
+                        /**
+                         * @var T
+                         */
+                        private $t;
+
+                        /**
+                         * @param T $t
+                         */
+                        public function __construct($t) {
+                            $this->t = $t;
+                        }
+
+                        /**
+                         * @return T
+                         * @psalm-mutation-free
+                         */
+                        public function get() {
+                            return $this->t;
+                        }
+                    }
+
+                    /** @var C<A>|C<B> $random_collection **/
+                    $a_or_b = $random_collection->get();',
+                [
+                    '$random_collection' => 'C<A>|C<B>',
+                    '$a_or_b' => 'A|B',
+                ],
+            ],
             'inferClosureParamTypeFromContext' => [
                 '<?php
                     /**