Initial inferred Closure Returns (#2945)

- Added detection for when we can infer
  closure types from parent fn return types
- Implemented functionality for infering types of
  returned closures
- Added ability to infer type of return from typed
  closures
- Added a new getFunctionLikeStorage method in
  Codebase to support easily getting a function
  despite being a method, closure, or function
- Added some utilities to the MethodIdentifier
  to facilitate creating MethodIdentifier's from
  string method references

Closes #2896

Signed-off-by: RJ Garcia <rj@bighead.net>

src/Psalm/Codebase.php

@@ -1,6 +1,7 @@
 <?php
 namespace Psalm;
 
+use Psalm\Internal\Analyzer\StatementsAnalyzer;
 use function array_combine;
 use function array_merge;
 use function count;
@@ -821,6 +822,24 @@ class Codebase
         return $this->classlikes->traitHasCorrectCase($fq_trait_name);
     }
 
+    /**
+     * Given a function id, return the function like storage for
+     * a method, closure, or function.
+     */
+    public function getFunctionLikeStorage(
+        StatementsAnalyzer $statements_analyzer,
+        string $function_id
+    ): FunctionLikeStorage {
+        $doesMethodExist =
+            \Psalm\Internal\MethodIdentifier::isValidMethodIdReference($function_id)
+            && $this->methodExists($function_id);
+        if ($doesMethodExist) {
+            return $this->methods->getStorage(\Psalm\Internal\MethodIdentifier::wrap($function_id));
+        }
+
+        return $this->functions->getStorage($statements_analyzer, $function_id);
+    }
+
     /**
      * Whether or not a given method exists
      *
@@ -836,15 +855,8 @@ class Codebase
         $calling_function_id = null,
         string $file_path = null
     ) {
-        if (is_string($method_id)) {
-            // remove trailing backslash if it exists
-            $method_id = preg_replace('/^\\\\/', '', $method_id);
-            $method_id_parts = explode('::', $method_id);
-            $method_id = new \Psalm\Internal\MethodIdentifier($method_id_parts[0], strtolower($method_id_parts[1]));
-        }
-
         return $this->methods->methodExists(
-            $method_id,
+            \Psalm\Internal\MethodIdentifier::wrap($method_id),
             $calling_function_id,
             $code_location,
             null,

src/Psalm/Internal/Analyzer/Statements/ReturnAnalyzer.php

@@ -2,7 +2,9 @@
 namespace Psalm\Internal\Analyzer\Statements;
 
 use PhpParser;
+use Psalm\Codebase;
 use Psalm\Internal\Analyzer\ClassLikeAnalyzer;
+use Psalm\Internal\Analyzer\ClosureAnalyzer;
 use Psalm\Internal\Analyzer\CommentAnalyzer;
 use Psalm\Internal\Analyzer\FunctionLikeAnalyzer;
 use Psalm\Internal\Analyzer\Statements\Expression\Call\ClassTemplateParamCollector;
@@ -12,6 +14,7 @@ use Psalm\Internal\Analyzer\TypeAnalyzer;
 use Psalm\CodeLocation;
 use Psalm\Context;
 use Psalm\Exception\DocblockParseException;
+use Psalm\Internal\Analyzer\TypeComparisonResult;
 use Psalm\Internal\Taint\Sink;
 use Psalm\Internal\Taint\Source;
 use Psalm\Issue\FalsableReturnStatement;
@@ -23,6 +26,8 @@ use Psalm\Issue\MixedReturnTypeCoercion;
 use Psalm\Issue\NoValue;
 use Psalm\Issue\NullableReturnStatement;
 use Psalm\IssueBuffer;
+use Psalm\Storage\FunctionLikeParameter;
+use Psalm\Storage\FunctionLikeStorage;
 use Psalm\Type;
 use function explode;
 use function strtolower;
@@ -122,6 +127,7 @@ class ReturnAnalyzer
 
         if ($stmt->expr) {
             $context->inside_call = true;
+            self::potentiallyInferTypesOnClosureFromParentReturnType($statements_analyzer, $stmt, $context);
 
             if (ExpressionAnalyzer::analyze($statements_analyzer, $stmt->expr, $context) === false) {
                 return false;
@@ -540,4 +546,98 @@ class ReturnAnalyzer
             );
         }
     }
+
+    /**
+     * If a function returns a closure, we try to infer the param/return types of
+     * the inner closure.
+     * @see \Psalm\Tests\ReturnTypeTest:756
+     */
+    private static function potentiallyInferTypesOnClosureFromParentReturnType(
+        StatementsAnalyzer $statements_analyzer,
+        PhpParser\Node\Stmt\Return_ $stmt,
+        Context $context
+    ): void {
+        // if not returning from inside of a function, return
+        if (!$context->calling_function_id) {
+            return;
+        }
+        // only handle if we returning a closure specifically
+        if (!$stmt->expr instanceof PhpParser\Node\Expr\Closure) {
+            return;
+        }
+
+        $closure_id = (new ClosureAnalyzer($stmt->expr, $statements_analyzer))->getId();
+        $closure_storage = $statements_analyzer
+            ->getCodebase()
+            ->getFunctionLikeStorage($statements_analyzer, $closure_id);
+
+        $parent_fn_storage = $statements_analyzer
+            ->getCodebase()
+            ->getFunctionLikeStorage($statements_analyzer, $context->calling_function_id);
+
+        // if no return type on parent, infer from return
+        if ($parent_fn_storage->return_type === null) {
+            $parent_fn_storage->return_type = self::functionLikeStorageToUnionType($closure_storage);
+            return;
+        }
+        // can't infer returned closure if the parent doesn't have a callable return type
+        if (!$parent_fn_storage->return_type->hasCallableType()) {
+            return;
+        }
+        // cannot infer if we have union/intersection types
+        if (\count($parent_fn_storage->return_type->getAtomicTypes()) !== 1) {
+            return;
+        }
+
+        /** @var Type\Atomic\TFn|Type\Atomic\TCallable $parent_callable_return_type */
+        $parent_callable_return_type = \current($parent_fn_storage->return_type->getAtomicTypes());
+        // if no params or return type was designed for parent callable return, then allow inferring of parent
+        // return type from child closure definition
+        if ($parent_callable_return_type->params === null && $parent_callable_return_type->return_type === null) {
+            $parent_fn_storage->return_type = self::functionLikeStorageToUnionType($closure_storage);
+            return;
+        }
+
+        foreach ($closure_storage->params as $key => $param) {
+            $parent_param = $parent_callable_return_type->params[$key] ?? null;
+            $param->type = self::inferInnerClosureTypeFromParent(
+                $statements_analyzer->getCodebase(),
+                $param->type,
+                $parent_param ? $parent_param->type : null
+            );
+        }
+        $closure_storage->return_type = self::inferInnerClosureTypeFromParent(
+            $statements_analyzer->getCodebase(),
+            $closure_storage->return_type,
+            $parent_callable_return_type->return_type
+        );
+    }
+
+    private static function functionLikeStorageToUnionType(
+        FunctionLikeStorage $closure
+    ): Type\Union {
+        return new Type\Union([
+            new Type\Atomic\TCallable('callable', $closure->params, $closure->return_type, $closure->pure)
+        ]);
+    }
+
+    /**
+     * - If non parent type, do nothing
+     * - If no return type, infer from parent
+     * - If parent return type is more specific, infer from parent
+     * - else, do nothing
+     */
+    private static function inferInnerClosureTypeFromParent(
+        Codebase $codebase,
+        ?Type\Union $return_type,
+        ?Type\Union $parent_return_type
+    ): ?Type\Union {
+        if (!$parent_return_type) {
+            return $return_type;
+        }
+        if (!$return_type || TypeAnalyzer::isContainedBy($codebase, $parent_return_type, $return_type)) {
+            return $parent_return_type;
+        }
+        return $return_type;
+    }
 }

src/Psalm/Internal/MethodIdentifier.php

@@ -17,6 +17,32 @@ class MethodIdentifier
         $this->method_name = $method_name;
     }
 
+    /**
+     * Takes any valid reference to a method id and converts
+     * it into a MethodIdentifier
+     * @param string|MethodIdentifier $method_id
+     */
+    public static function wrap($method_id): self
+    {
+        return \is_string($method_id) ? static::fromMethodIdReference($method_id) : $method_id;
+    }
+
+    public static function isValidMethodIdReference(string $method_id): bool
+    {
+        return \strpos($method_id, '::') !== false;
+    }
+
+    public static function fromMethodIdReference(string $method_id): self
+    {
+        if (!static::isValidMethodIdReference($method_id)) {
+            throw new \InvalidArgumentException('Invalid method id reference provided: ' . $method_id);
+        }
+        // remove trailing backslash if it exists
+        $method_id = \preg_replace('/^\\\\/', '', $method_id);
+        $method_id_parts = \explode('::', $method_id);
+        return new static($method_id_parts[0], \strtolower($method_id_parts[1]));
+    }
+
     /** @return string */
     public function __toString()
     {

tests/ReturnTypeTest.php

@@ -699,6 +699,146 @@ class ReturnTypeTest extends TestCase
                         }
                     }',
             ],
+            'infersClosureReturnTypes' => [
+                '<?php
+                    /**
+                     * @template T
+                     * @template U
+                     * @param callable(T): U $predicate
+                     * @return Closure(iterable<T>): iterable<U>
+                     */
+                    function map(callable $predicate): callable {
+                        return function($iter) use ($predicate) {
+                            foreach ($iter as $key => $value) {
+                                yield $key => $predicate($value);
+                            }
+                        };
+                    }
+
+                    $res = map(function(int $i): string { return (string) $i; })([1,2,3]);
+                ',
+                'assertions' => [
+                    '$res' => 'iterable<mixed, string>',
+                ],
+            ],
+            'infersClosureReturnTypesWithPartialTypehinting' => [
+                '<?php
+                    /**
+                     * @template T
+                     * @template U
+                     * @param callable(T): U $predicate
+                     * @return Closure(iterable<T>): iterable<U>
+                     */
+                    function map(callable $predicate): callable {
+                        return function(iterable $iter) use ($predicate): iterable {
+                            foreach ($iter as $key => $value) {
+                                yield $key => $predicate($value);
+                            }
+                        };
+                    }
+
+                    $res = map(function(int $i): string { return (string) $i; })([1,2,3]);
+                ',
+                'assertions' => [
+                    '$res' => 'iterable<mixed, string>',
+                ],
+            ],
+            'infersCallableReturnTypes' => [
+                '<?php
+                    /**
+                     * @template T
+                     * @template U
+                     * @param callable(T): U $predicate
+                     * @return callable(iterable<T>): iterable<U>
+                     */
+                    function map(callable $predicate): callable {
+                        return function($iter) use ($predicate) {
+                            foreach ($iter as $key => $value) {
+                                yield $key => $predicate($value);
+                            }
+                        };
+                    }
+
+                    $res = map(function(int $i): string { return (string) $i; })([1,2,3]);
+                ',
+                'assertions' => [
+                    '$res' => 'iterable<mixed, string>',
+                ],
+            ],
+            'infersCallableReturnTypesWithPartialTypehinting' => [
+                '<?php
+                    /**
+                     * @template T
+                     * @template U
+                     * @param callable(T): U $predicate
+                     * @return callable(iterable<T>): iterable<U>
+                     */
+                    function map(callable $predicate): callable {
+                        return function(iterable $iter) use ($predicate): iterable {
+                            foreach ($iter as $key => $value) {
+                                yield $key => $predicate($value);
+                            }
+                        };
+                    }
+
+                    $res = map(function(int $i): string { return (string) $i; })([1,2,3]);
+                ',
+                'assertions' => [
+                    '$res' => 'iterable<mixed, string>',
+                ],
+            ],
+            'infersReturnTypeFromReturnedCallable' => [
+                '<?php
+                    /**
+                     * @template T
+                     * @template U
+                     * @param callable(T): U $predicate
+                     */
+                    function map(callable $predicate) {
+                        return
+                        /**
+                         * @param iterable<T> $iter
+                         * @return iterable<U>
+                         */
+                        function($iter) use ($predicate): iterable {
+                            foreach ($iter as $key => $value) {
+                                yield $key => $predicate($value);
+                            }
+                        };
+                    }
+
+                    $res = map(function(int $i): string { return (string) $i; })([1,2,3]);
+                ',
+                'assertions' => [
+                    '$res' => 'iterable<mixed, string>',
+                ],
+            ],
+            'infersReturnTypeFromReturnedCallableWithPartialReturnStatement' => [
+                '<?php
+                    /**
+                     * @template T
+                     * @template U
+                     * @param callable(T): U $predicate
+                     */
+                    function map(callable $predicate): callable {
+                        return
+                        /**
+                         * @param iterable<T> $iter
+                         * @return iterable<U>
+                         */
+                        function($iter) use ($predicate): iterable {
+                            foreach ($iter as $key => $value) {
+                                yield $key => $predicate($value);
+                            }
+                        };
+                    }
+
+                    $res = map(function(int $i): string { return (string) $i; })([1,2,3]);
+                ',
+                'assertions' => [
+                    '$res' => 'iterable<mixed, string>',
+                ],
+            ],
         ];
     }
 
@@ -1022,6 +1162,82 @@ class ReturnTypeTest extends TestCase
                     ): string {}',
                 'error_message' => 'InvalidReturnType - src/somefile.php:4:24',
             ],
+            'cannotInferReturnClosureWithoutReturn' => [
+                '<?php
+                /**
+                 * @template T
+                 * @template U
+                 * @param callable(T): U $predicate
+                 * @return callable(iterable<T>): iterable<U>
+                 */
+                function map(callable $predicate): callable {
+                    $a = function($iter) use ($predicate) {
+                        foreach ($iter as $key => $value) {
+                            yield $key => $predicate($value);
+                        }
+                    };
+                    return $a;
+                }
+
+                $res = map(function(int $i): string { return (string) $i; })([1,2,3]);
+                ',
+                'error_message' => 'MixedAssignment - src/somefile.php:10:51 - Cannot assign $value to a mixed type',
+            ],
+            'cannotInferReturnClosureWithMoreSpecificTypes' => [
+                '<?php
+                /**
+                 * @template T
+                 * @template U
+                 * @param callable(T): U $predicate
+                 * @return callable(iterable<T>): iterable<U>
+                 */
+                function map(callable $predicate): callable {
+                    return
+                    /** @param iterable<int> $iter */
+                    function($iter) use ($predicate) {
+                        foreach ($iter as $key => $value) {
+                            yield $key => $predicate($value);
+                        }
+                    };
+                }
+
+                $res = map(function(int $i): string { return (string) $i; })([1,2,3]);
+                ',
+                'error_message' => 'InvalidArgument - src/somefile.php:13:54 - Argument 1 expects T:fn-map as mixed, int provided',
+            ],
+            'cannotInferReturnClosureWithDifferentReturnTypes' => [
+                '<?php
+                /**
+                 * @template T
+                 * @template U
+                 * @param callable(T): U $predicate
+                 * @return callable(iterable<T>): iterable<U>
+                 */
+                function map(callable $predicate): callable {
+                    return function($iter) use ($predicate): int {
+                        return 1;
+                    };
+                }
+
+                $res = map(function(int $i): string { return (string) $i; })([1,2,3]);
+                ',
+                'error_message' => 'InvalidReturnStatement - src/somefile.php:9:28 - The inferred type \'Closure(iterable<mixed, T:fn-map as mixed>):int(1)\' does not match the declared return type \'callable(iterable<mixed, T:fn-map as mixed>):iterable<mixed, U:fn-map as mixed>\' for map',
+            ],
+            'cannotInferReturnClosureWithDifferentTypes' => [
+                '<?php
+                class A {}
+                class B {}
+                /**
+                 * @return callable(A): void
+                 */
+                function map(): callable {
+                    return function(B $v): void {};
+                }
+
+                $res = map(function(int $i): string { return (string) $i; })([1,2,3]);
+                ',
+                'error_message' => 'InvalidReturnStatement - src/somefile.php:8:28 - The inferred type \'Closure(B):void\' does not match the declared return type \'callable(A):void\' for map',
+            ],
         ];
     }
 }