Tighten up rules arouund when mutation-free methods get memoised

src/Psalm/Context.php

@@ -647,6 +647,8 @@ class Context
             return;
         }
 
+        $existing_type->allow_mutations = true;
+
         $this->removeVarFromConflictingClauses(
             $remove_var_id,
             $existing_type->hasMixed()
@@ -671,12 +673,14 @@ class Context
         }
     }
 
-    public function removeAllObjectVars(): void
+    public function removeMutableObjectVars(): void
     {
         $vars_to_remove = [];
 
-        foreach ($this->vars_in_scope as $var_id => $_) {
+        foreach ($this->vars_in_scope as $var_id => $type) {
-            if (strpos($var_id, '->') !== false || strpos($var_id, '::') !== false) {
+            if ($type->has_mutations
+                && (strpos($var_id, '->') !== false || strpos($var_id, '::') !== false)
+            ) {
                 $vars_to_remove[] = $var_id;
             }
         }

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

@@ -968,7 +968,7 @@ class FunctionCallAnalyzer extends CallAnalyzer
                 }
 
                 if (!$config->remember_property_assignments_after_call) {
-                    $context->removeAllObjectVars();
+                    $context->removeMutableObjectVars();
                 }
             } elseif ($function_call_info->function_id
                 && (($function_call_info->function_storage

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

@@ -84,4 +84,9 @@ class AtomicMethodCallAnalysisResult
      * @var bool
      */
     public $can_memoize = false;
+
+    /**
+     * @var bool
+     */
+    public $immutable_call = false;
 }

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

@@ -251,7 +251,7 @@ class ExistingAtomicMethodCallAnalyzer extends CallAnalyzer
 
         if ($method_storage) {
             if (!$context->collect_mutations && !$context->collect_initializations) {
-                $result->can_memoize = MethodCallPurityAnalyzer::analyze(
+                MethodCallPurityAnalyzer::analyze(
                     $statements_analyzer,
                     $codebase,
                     $stmt,
@@ -261,7 +261,8 @@ class ExistingAtomicMethodCallAnalyzer extends CallAnalyzer
                     $method_storage,
                     $class_storage,
                     $context,
-                    $config
+                    $config,
+                    $result
                 );
             }
 

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

@@ -23,10 +23,9 @@ class MethodCallPurityAnalyzer
         \Psalm\Storage\MethodStorage $method_storage,
         \Psalm\Storage\ClassLikeStorage $class_storage,
         Context $context,
-        \Psalm\Config $config
+        \Psalm\Config $config,
-    ) : bool {
+        AtomicMethodCallAnalysisResult $result
-        $can_memoize = false;
+    ) : void {
-
         $method_pure_compatible = $method_storage->external_mutation_free
             && $statements_analyzer->node_data->isPureCompatible($stmt->var);
 
@@ -81,16 +80,23 @@ class MethodCallPurityAnalyzer
             if ($method_storage->mutation_free
                 && (!$method_storage->mutation_free_inferred
                     || $method_storage->final)
+                && ($method_storage->immutable || $config->remember_property_assignments_after_call)
             ) {
                 if ($context->inside_conditional
                     && !$method_storage->assertions
                     && !$method_storage->if_true_assertions
                 ) {
                     /** @psalm-suppress UndefinedPropertyAssignment */
-                    $stmt->pure = true;
+                    $stmt->memoizable = true;
+
+                    if ($method_storage->immutable) {
+                        /** @psalm-suppress UndefinedPropertyAssignment */
+                        $stmt->pure = true;
+                    }
                 }
 
-                $can_memoize = true;
+                $result->can_memoize = true;
+                $result->immutable_call = $method_storage->immutable;
             }
 
             if ($codebase->find_unused_variables
@@ -128,7 +134,7 @@ class MethodCallPurityAnalyzer
             && !$method_storage->mutation_free
             && !$method_pure_compatible
         ) {
-            $context->removeAllObjectVars();
+            $context->removeMutableObjectVars();
         } elseif ($method_storage->this_property_mutations) {
             foreach ($method_storage->this_property_mutations as $name => $_) {
                 $mutation_var_id = $lhs_var_id . '->' . $name;
@@ -144,7 +150,5 @@ class MethodCallPurityAnalyzer
                 }
             }
         }
-
-        return $can_memoize;
     }
 }

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

@@ -205,14 +205,17 @@ class MethodCallAnalyzer extends \Psalm\Internal\Analyzer\Statements\Expression\
         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;
+                    $context->vars_in_scope[$method_var_id]->has_mutations = false;
+                }
+
+                if ($result->can_memoize) {
+                    /** @psalm-suppress UndefinedPropertyAssignment */
+                    $stmt->memoizable = true;
                 }
             }
         }

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

@@ -222,7 +222,7 @@ class NewAnalyzer extends \Psalm\Internal\Analyzer\Statements\Expression\CallAna
         }
 
         if (!$config->remember_property_assignments_after_call && !$context->collect_initializations) {
-            $context->removeAllObjectVars();
+            $context->removeMutableObjectVars();
         }
 
         return true;

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

@@ -225,7 +225,7 @@ class StaticCallAnalyzer extends CallAnalyzer
         }
 
         if (!$config->remember_property_assignments_after_call && !$context->collect_initializations) {
-            $context->removeAllObjectVars();
+            $context->removeMutableObjectVars();
         }
 
         if (!$statements_analyzer->node_data->getType($stmt)) {

src/Psalm/Internal/Analyzer/Statements/Expression/CallAnalyzer.php

@@ -874,10 +874,10 @@ class CallAnalyzer
                         $readonly_template_result = new TemplateResult($template_type_map, $template_type_map);
 
                          \Psalm\Internal\Type\TemplateInferredTypeReplacer::replace(
-                            $op_vars_in_scope[$var_id],
+                             $op_vars_in_scope[$var_id],
-                            $readonly_template_result,
+                             $readonly_template_result,
-                            $codebase
+                             $codebase
-                        );
+                         );
                     }
 
                     $op_vars_in_scope[$var_id]->from_docblock = true;

src/Psalm/Internal/Analyzer/Statements/Expression/ExpressionIdentifier.php

@@ -199,7 +199,7 @@ class ExpressionIdentifier
         ) {
             $config = \Psalm\Config::getInstance();
 
-            if ($config->memoize_method_calls || isset($stmt->pure)) {
+            if ($config->memoize_method_calls || isset($stmt->memoizable)) {
                 $lhs_var_name = self::getArrayVarId(
                     $stmt->var,
                     $this_class_name,

src/Psalm/Internal/Analyzer/Statements/Expression/Fetch/AtomicPropertyFetchAnalyzer.php

@@ -425,6 +425,10 @@ class AtomicPropertyFetchAnalyzer
             $in_assignment
         );
 
+        if ($class_storage->mutation_free) {
+            $class_property_type->has_mutations = false;
+        }
+
         if ($stmt_type = $statements_analyzer->node_data->getType($stmt)) {
             $statements_analyzer->node_data->setType(
                 $stmt,
@@ -1087,6 +1091,7 @@ class AtomicPropertyFetchAnalyzer
                 );
             }
         }
+
         return $class_property_type;
     }
 }

src/Psalm/Internal/Codebase/Populator.php

@@ -219,6 +219,7 @@ class Populator
                 if (!$method->is_static && !$method->external_mutation_free) {
                     $method->mutation_free = $storage->mutation_free;
                     $method->external_mutation_free = $storage->external_mutation_free;
+                    $method->immutable = $storage->mutation_free;
                 }
             }
 

src/Psalm/Internal/Type/SimpleAssertionReconciler.php

@@ -1669,7 +1669,7 @@ class SimpleAssertionReconciler extends \Psalm\Type\Reconciler
                 $did_remove_type = true;
             } elseif ($type instanceof TTemplateParam) {
                 if ($type->as->hasArray() || $type->as->hasMixed()) {
-                   $type = clone $type;
+                    $type = clone $type;
 
                     $type->as = self::reconcileArray(
                         $type->as,

src/Psalm/Storage/MethodStorage.php

@@ -65,6 +65,11 @@ class MethodStorage extends FunctionLikeStorage
      */
     public $external_mutation_free = false;
 
+    /**
+     * @var bool
+     */
+    public $immutable = false;
+
     /**
      * @var bool
      */

src/Psalm/Type/Union.php

@@ -174,6 +174,11 @@ class Union implements TypeNode
      */
     public $allow_mutations = true;
 
+    /**
+     * @var bool
+     */
+    public $has_mutations = true;
+
     /** @var null|string */
     private $id;
 

tests/PropertyTypeTest.php

@@ -31,8 +31,7 @@ class PropertyTypeTest extends TestCase
                 }
 
                 class X {
-                    /** @var ?int **/
+                    public ?int $x = null;
-                    public $x;
 
                     public function getX(): int {
                         $this->x = 5;
@@ -64,8 +63,7 @@ class PropertyTypeTest extends TestCase
                 }
 
                 class X {
-                    /** @var ?int **/
+                    public ?int $x = null;
-                    public $x;
 
                     public function getX(): int {
                         $this->x = 5;
@@ -80,7 +78,7 @@ class PropertyTypeTest extends TestCase
         $this->analyzeFile('somefile.php', new Context());
     }
 
-    public function testForgetPropertyAssignmentsInBranchWithThrow(): void
+    public function testForgetPropertyAssignmentsInBranch(): void
     {
         Config::getInstance()->remember_property_assignments_after_call = false;
 
@@ -99,19 +97,130 @@ class PropertyTypeTest extends TestCase
                 }
 
                 class X {
-                    /** @var ?int **/
+                    public ?int $x = null;
-                    public $x;
+                }
 
-                    public function getX(bool $b): int {
+                function testX(X $x): void {
-                        $this->x = 5;
+                    $x->x = 5;
 
-                        if ($b) {
+                    if (rand(0, 1)) {
-                            XCollector::modify();
+                        XCollector::modify();
-                            throw new \Exception("bad");
+                    }
+
+                    if ($x->x === null) {}
+                }'
+        );
+
+        $this->analyzeFile('somefile.php', new Context());
+    }
+
+    public function testForgetFinalMethodCalls(): void
+    {
+        Config::getInstance()->remember_property_assignments_after_call = false;
+
+        $this->addFile(
+            'somefile.php',
+            '<?php
+                class XCollector {
+                    /** @var X[] */
+                    private static array $xs = [];
+
+                    public static function modify() : void {
+                        foreach (self::$xs as $x) {
+                            $x->x = null;
                         }
+                    }
+                }
 
+                class X {
+                    public ?int $x = null;
+
+                    public function __construct(?int $x) {
+                        $this->x = $x;
+                    }
+
+                    public final function getX() : ?int {
                         return $this->x;
                     }
+                }
+
+                function testX(X $x): void {
+                    if ($x->getX()) {
+                        XCollector::modify();
+                        if ($x->getX() === null) {}
+                    }
+                }'
+        );
+
+        $this->analyzeFile('somefile.php', new Context());
+    }
+
+    public function testRememberImmutableMethodCalls(): void
+    {
+        Config::getInstance()->remember_property_assignments_after_call = false;
+
+        $this->expectExceptionMessage('TypeDoesNotContainNull - somefile.php:22:29');
+        $this->expectException(\Psalm\Exception\CodeException::class);
+
+        $this->addFile(
+            'somefile.php',
+            '<?php
+                class XCollector {
+                    public static function modify() : void {}
+                }
+
+                /** @psalm-immutable */
+                class X {
+                    public ?int $x = null;
+
+                    public function __construct(?int $x) {
+                        $this->x = $x;
+                    }
+
+                    public function getX() : ?int {
+                        return $this->x;
+                    }
+                }
+
+                function testX(X $x): void {
+                    if ($x->getX()) {
+                        XCollector::modify();
+                        if ($x->getX() === null) {}
+                    }
+                }'
+        );
+
+        $this->analyzeFile('somefile.php', new Context());
+    }
+
+    public function testRememberImmutableProperties(): void
+    {
+        Config::getInstance()->remember_property_assignments_after_call = false;
+
+        $this->expectExceptionMessage('TypeDoesNotContainNull - somefile.php:18:29');
+        $this->expectException(\Psalm\Exception\CodeException::class);
+
+        $this->addFile(
+            'somefile.php',
+            '<?php
+                class XCollector {
+                    public static function modify() : void {}
+                }
+
+                /** @psalm-immutable */
+                class X {
+                    public ?int $x = null;
+
+                    public function __construct(?int $x) {
+                        $this->x = $x;
+                    }
+                }
+
+                function testX(X $x): void {
+                    if ($x->x) {
+                        XCollector::modify();
+                        if ($x->x === null) {}
+                    }
                 }'
         );
 

tests/PureAnnotationTest.php

@@ -333,7 +333,8 @@ class PureAnnotationTest extends TestCase
                         public function foo() : void {}
 
                         public function doSomething(): void {
-                            if ($this->checkNotNullNested() && $this->other->foo()) {}
+                            $this->checkNotNullNested();
+                            $this->other->foo();
                         }
                     }'
             ],