From c0ca383020943f0554eb80b440ce144a74fe3f51 Mon Sep 17 00:00:00 2001
From: adrew <klimichkartorgnusov@gmail.com>
Date: Mon, 27 Dec 2021 19:35:37 +0300
Subject: [PATCH 1/2] Contextual inference for closure param types

---
 .../Expression/Call/ArgumentsAnalyzer.php     | 17 ++++
 .../Expression/Call/FunctionCallAnalyzer.php  | 10 ++-
 tests/CallableTest.php                        | 90 +++++++++++++++++++
 3 files changed, 115 insertions(+), 2 deletions(-)

diff --git a/src/Psalm/Internal/Analyzer/Statements/Expression/Call/ArgumentsAnalyzer.php b/src/Psalm/Internal/Analyzer/Statements/Expression/Call/ArgumentsAnalyzer.php
index b7741e84f5e..1404cb97a45 100644
--- a/src/Psalm/Internal/Analyzer/Statements/Expression/Call/ArgumentsAnalyzer.php
+++ b/src/Psalm/Internal/Analyzer/Statements/Expression/Call/ArgumentsAnalyzer.php
@@ -231,6 +231,23 @@ public static function analyze(
                 );
             }
 
+            $inferred_arg_type = $statements_analyzer->node_data->getType($arg->value);
+
+            if (null !== $inferred_arg_type && null !== $template_result && null !== $param && null !== $param->type) {
+                $codebase = $statements_analyzer->getCodebase();
+
+                TemplateStandinTypeReplacer::replace(
+                    clone $param->type,
+                    $template_result,
+                    $codebase,
+                    $statements_analyzer,
+                    $inferred_arg_type,
+                    $argument_offset,
+                    $context->self,
+                    $context->calling_method_id ?: $context->calling_function_id
+                );
+            }
+
             if ($toggled_class_exists) {
                 $context->inside_class_exists = false;
             }
diff --git a/src/Psalm/Internal/Analyzer/Statements/Expression/Call/FunctionCallAnalyzer.php b/src/Psalm/Internal/Analyzer/Statements/Expression/Call/FunctionCallAnalyzer.php
index d59208dfd36..225aa21b27b 100644
--- a/src/Psalm/Internal/Analyzer/Statements/Expression/Call/FunctionCallAnalyzer.php
+++ b/src/Psalm/Internal/Analyzer/Statements/Expression/Call/FunctionCallAnalyzer.php
@@ -10,7 +10,6 @@
 use Psalm\Internal\Algebra\FormulaGenerator;
 use Psalm\Internal\Analyzer\AlgebraAnalyzer;
 use Psalm\Internal\Analyzer\FunctionLikeAnalyzer;
-use Psalm\Internal\Analyzer\Statements\Expression\Call\MethodCallAnalyzer;
 use Psalm\Internal\Analyzer\Statements\Expression\CallAnalyzer;
 use Psalm\Internal\Analyzer\Statements\ExpressionAnalyzer;
 use Psalm\Internal\Analyzer\StatementsAnalyzer;
@@ -165,13 +164,20 @@ public static function analyze(
         }
 
         if (!$is_first_class_callable) {
+            $template_result = null;
+
+            if (isset($function_call_info->function_storage->template_types)) {
+                $template_result = new TemplateResult($function_call_info->function_storage->template_types ?: [], []);
+            }
+
             ArgumentsAnalyzer::analyze(
                 $statements_analyzer,
                 $stmt->getArgs(),
                 $function_call_info->function_params,
                 $function_call_info->function_id,
                 $function_call_info->allow_named_args,
-                $context
+                $context,
+                $template_result
             );
         }
 
diff --git a/tests/CallableTest.php b/tests/CallableTest.php
index de0770ac950..9b0a2fb1250 100644
--- a/tests/CallableTest.php
+++ b/tests/CallableTest.php
@@ -147,6 +147,96 @@ function asTupled(ArrayList $list): ArrayList
                     '$b' => 'ArrayList<array{int}>',
                 ],
             ],
+            'inferArgByPreviousFunctionArg' => [
+                '<?php
+                    /**
+                     * @template A
+                     * @template B
+                     *
+                     * @param iterable<array-key, A> $_collection
+                     * @param callable(A): B $_ab
+                     * @return list<B>
+                     */
+                    function map(iterable $_collection, callable $_ab) { return []; }
+
+                    /** @template T */
+                    final class Foo
+                    {
+                        /** @return Foo<int> */
+                        public function toInt() { throw new RuntimeException("???"); }
+                    }
+
+                    /** @var list<Foo<string>> */
+                    $items = [];
+
+                    $inferred = map($items, function ($i) {
+                        return $i->toInt();
+                    });',
+                'assertions' => [
+                    '$inferred' => 'list<Foo<int>>',
+                ],
+            ],
+            'inferTemplateForExplicitlyTypedArgByPreviousFunctionArg' => [
+                '<?php
+                    /**
+                     * @template A
+                     * @template B
+                     *
+                     * @param iterable<array-key, A> $_collection
+                     * @param callable(A): B $_ab
+                     * @return list<B>
+                     */
+                    function map(iterable $_collection, callable $_ab) { return []; }
+
+                    /** @template T */
+                    final class Foo
+                    {
+                        /** @return Foo<int> */
+                        public function toInt() { throw new RuntimeException("???"); }
+                    }
+
+                    /** @var list<Foo<string>> */
+                    $items = [];
+
+                    $inferred = map($items, function (Foo $i) {
+                        return $i->toInt();
+                    });',
+                'assertions' => [
+                    '$inferred' => 'list<Foo<int>>',
+                ],
+            ],
+            'doNotInferTemplateForExplicitlyTypedWithPhpdocArgByPreviousFunctionArg' => [
+                '<?php
+                    /**
+                     * @template A
+                     * @template B
+                     *
+                     * @param iterable<array-key, A> $_collection
+                     * @param callable(A): B $_ab
+                     * @return list<B>
+                     */
+                    function map(iterable $_collection, callable $_ab) { return []; }
+
+                    /** @template T */
+                    final class Foo
+                    {
+                        /** @return Foo<int> */
+                        public function toInt() { throw new RuntimeException("???"); }
+                    }
+
+                    /** @var list<Foo<string>> */
+                    $items = [];
+
+                    $inferred = map($items,
+                        /** @param Foo $i */
+                        function ($i) {
+                            return $i;
+                        }
+                    );',
+                'assertions' => [
+                    '$inferred' => 'list<Foo>',
+                ],
+            ],
             'varReturnType' => [
                 '<?php
                     $add_one = function(int $a) : int {

From 87deb19a87b3f616774c2c1c917f58c1d8bb7139 Mon Sep 17 00:00:00 2001
From: adrew <klimichkartorgnusov@gmail.com>
Date: Mon, 27 Dec 2021 19:43:06 +0300
Subject: [PATCH 2/2] Remove dead code from test

---
 tests/CallableTest.php | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/tests/CallableTest.php b/tests/CallableTest.php
index 9b0a2fb1250..3864ec7a00c 100644
--- a/tests/CallableTest.php
+++ b/tests/CallableTest.php
@@ -218,11 +218,7 @@ public function toInt() { throw new RuntimeException("???"); }
                     function map(iterable $_collection, callable $_ab) { return []; }
 
                     /** @template T */
-                    final class Foo
-                    {
-                        /** @return Foo<int> */
-                        public function toInt() { throw new RuntimeException("???"); }
-                    }
+                    final class Foo { }
 
                     /** @var list<Foo<string>> */
                     $items = [];