Reimplemented contextual inference

psalm-baseline.xml

@@ -315,12 +315,6 @@
       <code><![CDATA[$method_tree->children[1]]]></code>
     </PossiblyUndefinedIntArrayOffset>
   </file>
-  <file src="src/Psalm/Internal/PhpVisitor/Reflector/ClassLikeNodeScanner.php">
-    <PossiblyUndefinedIntArrayOffset>
-      <code>$l[4]</code>
-      <code>$r[4]</code>
-    </PossiblyUndefinedIntArrayOffset>
-  </file>
   <file src="src/Psalm/Internal/PhpVisitor/Reflector/ExpressionScanner.php">
     <PossiblyUndefinedIntArrayOffset>
       <code><![CDATA[$node->getArgs()[0]]]></code>

src/Psalm/Context.php

@@ -422,6 +422,11 @@ final class Context
      */
     public $parent_remove_vars = [];
 
+    /**
+     * @var ContextualTypeResolver|null
+     */
+    public $contextual_type_resolver = null;
+
     /** @internal */
     public function __construct(?string $self = null)
     {
@@ -947,4 +952,28 @@ public function insideUse(): bool
             || $this->inside_throw
             || $this->inside_isset;
     }
+
+    /**
+     * Returns all possible template definers in context (class, function or method).
+     *
+     * @return array<string, true>
+     */
+    public function getPossibleTemplateDefiners(): array
+    {
+        $who_can_define_templates = [];
+
+        if ($this->calling_function_id !== null) {
+            $who_can_define_templates['fn-' . $this->calling_function_id] = true;
+        }
+
+        if ($this->calling_method_id !== null) {
+            $who_can_define_templates['fn-' . $this->calling_method_id] = true;
+        }
+
+        if ($this->self !== null) {
+            $who_can_define_templates[$this->self] = true;
+        }
+
+        return $who_can_define_templates;
+    }
 }

src/Psalm/ContextualTypeResolver.php

@@ -0,0 +1,60 @@
+<?php
+
+namespace Psalm;
+
+use Psalm\Internal\Type\TemplateInferredTypeReplacer;
+use Psalm\Internal\Type\TemplateResult;
+use Psalm\Internal\Type\TemplateStandinTypeReplacer;
+use Psalm\Type\Union;
+
+final class ContextualTypeResolver
+{
+    private Union $contextual_type;
+    private TemplateResult $template_result;
+    private Codebase $codebase;
+
+    public function __construct(
+        Union $contextual_type,
+        TemplateResult $template_result,
+        Codebase $codebase
+    ) {
+        $this->contextual_type = $contextual_type;
+        $this->template_result = $template_result;
+        $this->codebase = $codebase;
+    }
+
+    public function getCodebase(): Codebase
+    {
+        return $this->codebase;
+    }
+
+    /**
+     * @return ($type is Union ? self : null)
+     */
+    public function withContextualType(?Union $type): ?self
+    {
+        return $type !== null
+            ? new self($type, $this->template_result, $this->codebase)
+            : null;
+    }
+
+    public function resolve(): Union
+    {
+        return TemplateInferredTypeReplacer::replace(
+            $this->contextual_type,
+            $this->template_result,
+            $this->codebase,
+        );
+    }
+
+    public function fillTemplateResult(Union $input_type): void
+    {
+        TemplateStandinTypeReplacer::fillTemplateResult(
+            $this->contextual_type,
+            $this->template_result,
+            $this->codebase,
+            null,
+            $input_type,
+        );
+    }
+}

src/Psalm/Internal/Analyzer/ClosureAnalyzer.php

@@ -9,14 +9,14 @@
 use Psalm\Internal\Codebase\VariableUseGraph;
 use Psalm\Internal\DataFlow\DataFlowNode;
 use Psalm\Internal\PhpVisitor\ShortClosureVisitor;
+use Psalm\Internal\Type\Comparator\TypeComparisonResult;
 use Psalm\Internal\Type\Comparator\UnionTypeComparator;
+use Psalm\Internal\Type\TypeExpander;
 use Psalm\Issue\DuplicateParam;
 use Psalm\Issue\PossiblyUndefinedVariable;
 use Psalm\Issue\UndefinedVariable;
 use Psalm\IssueBuffer;
 use Psalm\Type;
-use Psalm\Type\Atomic\TCallable;
-use Psalm\Type\Atomic\TClosure;
 use Psalm\Type\Atomic\TMixed;
 use Psalm\Type\Atomic\TNamedObject;
 use Psalm\Type\Union;
@@ -33,6 +33,8 @@
  */
 class ClosureAnalyzer extends FunctionLikeAnalyzer
 {
+    public ?Union $possibly_return_type = null;
+
     /**
      * @param PhpParser\Node\Expr\Closure|PhpParser\Node\Expr\ArrowFunction $function
      */
@@ -78,13 +80,6 @@ public static function analyzeExpression(
     ): bool {
         $closure_analyzer = new ClosureAnalyzer($stmt, $statements_analyzer);
 
-        if ($context->inside_return) {
-            self::potentiallyInferTypesOnClosureFromParentReturnType(
-                $statements_analyzer,
-                $closure_analyzer,
-            );
-        }
-
         if ($stmt instanceof PhpParser\Node\Expr\Closure
             && self::analyzeClosureUses($statements_analyzer, $stmt, $context) === false
         ) {
@@ -216,6 +211,12 @@ public static function analyzeExpression(
         $use_context->calling_method_id = $context->calling_method_id;
         $use_context->phantom_classes = $context->phantom_classes;
 
+        $closure_analyzer->possibly_return_type = self::inferParamTypesFromContextualTypeAndGetPossibleReturnType(
+            $context,
+            $statements_analyzer,
+            $closure_analyzer,
+        );
+
         $closure_analyzer->analyze($use_context, $statements_analyzer->node_data, $context, false);
 
         if ($closure_analyzer->inferred_impure
@@ -360,82 +361,94 @@ public static function analyzeClosureUses(
      *
      * @see \Psalm\Tests\ReturnTypeTest:756
      */
-    private static function potentiallyInferTypesOnClosureFromParentReturnType(
+    private static function inferParamTypesFromContextualTypeAndGetPossibleReturnType(
+        Context $context,
         StatementsAnalyzer $statements_analyzer,
         ClosureAnalyzer $closure_analyzer
-    ): void {
-        $parent_source = $statements_analyzer->getSource();
-
-        // if not returning from inside a function, return
-        if (!$parent_source instanceof ClosureAnalyzer && !$parent_source instanceof FunctionAnalyzer) {
-            return;
+    ): ?Union {
+        if (!$context->contextual_type_resolver) {
+            return null;
         }
 
-        $closure_id = $closure_analyzer->getClosureId();
-        $closure_storage = $statements_analyzer
-            ->getCodebase()
-            ->getFunctionLikeStorage($statements_analyzer, $closure_id);
-
-        $parent_fn_storage = $parent_source->getFunctionLikeStorage($statements_analyzer);
+        $contextual_callable_type = ClosureAnalyzerContextualTypeExtractor::extract(
+            $context->contextual_type_resolver,
+        );
 
-        if ($parent_fn_storage->return_type === null) {
-            return;
+        if ($contextual_callable_type === null) {
+            return null;
         }
 
-        // can't infer returned closure if the parent doesn't have a callable return type
-        if (!$parent_fn_storage->return_type->hasCallableType()) {
-            return;
+        if ($contextual_callable_type->params === null) {
+            return null;
         }
 
-        // cannot infer if we have union/intersection types
-        if (!$parent_fn_storage->return_type->isSingle()) {
-            return;
-        }
+        $codebase = $statements_analyzer->getCodebase();
 
-        /** @var TClosure|TCallable $parent_callable_return_type */
-        $parent_callable_return_type = $parent_fn_storage->return_type->getSingleAtomic();
+        $calling_closure_storage = $codebase->getFunctionLikeStorage(
+            $statements_analyzer,
+            $closure_analyzer->getClosureId(),
+        );
 
-        if ($parent_callable_return_type->params === null && $parent_callable_return_type->return_type === null) {
-            return;
-        }
+        foreach ($calling_closure_storage->params as $param_index => $calling_param) {
+            $contextual_param_type = $contextual_callable_type->params[$param_index]->type ?? null;
 
-        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->type ?? null,
-            );
-        }
+            // No contextual type info. Don't infer.
+            if ($contextual_param_type === null) {
+                continue;
+            }
 
-        $closure_storage->return_type = self::inferInnerClosureTypeFromParent(
-            $statements_analyzer->getCodebase(),
-            $closure_storage->return_type,
-            $parent_callable_return_type->return_type,
-        );
+            // Explicit docblock type. Don't infer.
+            if ($calling_param->type !== $calling_param->signature_type) {
+                continue;
+            }
 
-        if (!$closure_storage->template_types && $parent_callable_return_type->templates) {
-            $closure_storage->template_types = $parent_callable_return_type->getTemplateMap();
-        }
-    }
+            $contextual_param_type = self::expandContextualType($codebase, $context, $contextual_param_type);
+            $type_comparison_result = new TypeComparisonResult();
+
+            if ($calling_param->type === null
+                || UnionTypeComparator::isContainedBy(
+                    $codebase,
+                    $contextual_param_type,
+                    $calling_param->type,
+                    false,
+                    false,
+                    $type_comparison_result,
+                )
+            ) {
+                if ($type_comparison_result->to_string_cast) {
+                    continue;
+                }
 
-    /**
-     * - 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,
-        ?Union $return_type,
-        ?Union $parent_return_type
-    ): ?Union {
-        if (!$parent_return_type) {
-            return $return_type;
+                $calling_param->type = $contextual_param_type;
+                $calling_param->type_inferred = true;
+            }
         }
-        if (!$return_type || UnionTypeComparator::isContainedBy($codebase, $parent_return_type, $return_type)) {
-            return $parent_return_type;
+
+        if (!$calling_closure_storage->template_types && $contextual_callable_type->templates) {
+            $calling_closure_storage->template_types = $contextual_callable_type->getTemplateMap();
         }
-        return $return_type;
+
+        return $contextual_callable_type->return_type;
+    }
+
+    private static function expandContextualType(Codebase $codebase, Context $context, Union $contextual_type): Union
+    {
+        $expand_templates = true;
+        $do_not_expand_template_defined_at = $context->getPossibleTemplateDefiners();
+
+        return TypeExpander::expandUnion(
+            $codebase,
+            $contextual_type,
+            null,
+            null,
+            null,
+            true,
+            false,
+            false,
+            false,
+            $expand_templates,
+            false,
+            $do_not_expand_template_defined_at,
+        );
     }
 }

src/Psalm/Internal/Analyzer/ClosureAnalyzerContextualTypeExtractor.php

@@ -0,0 +1,36 @@
+<?php
+
+namespace Psalm\Internal\Analyzer;
+
+use Psalm\ContextualTypeResolver;
+use Psalm\Type\Atomic;
+use Psalm\Type\Atomic\TCallable;
+use Psalm\Type\Atomic\TClosure;
+
+use function count;
+
+/**
+ * @internal
+ */
+final class ClosureAnalyzerContextualTypeExtractor
+{
+    /**
+     * @return null|TClosure|TCallable
+     */
+    public static function extract(ContextualTypeResolver $contextual_type_resolver): ?Atomic
+    {
+        $candidates = [];
+
+        $atomics = $contextual_type_resolver
+            ->resolve()
+            ->getAtomicTypes();
+
+        foreach ($atomics as $atomic) {
+            if ($atomic instanceof TClosure || $atomic instanceof TCallable) {
+                $candidates[] = $atomic;
+            }
+        }
+
+        return count($candidates) === 1 ? $candidates[0] : null;
+    }
+}

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

@@ -364,10 +364,23 @@ private static function analyzeArrayItem(
             $key_type = new Union([$key_atomic_type]);
         }
 
+        $was_contextual_type_resolver = $context->contextual_type_resolver;
+
+        $context->contextual_type_resolver = $context->contextual_type_resolver !== null
+            ? ArrayAnalyzerContextualTypeExtractor::extract(
+                $item_key_type ?? Type::getInt(false, $array_creation_info->int_offset),
+                $context->contextual_type_resolver,
+            )
+            : null;
+
         if (ExpressionAnalyzer::analyze($statements_analyzer, $item->value, $context) === false) {
+            $context->contextual_type_resolver = $was_contextual_type_resolver;
+
             return;
         }
 
+        $context->contextual_type_resolver = $was_contextual_type_resolver;
+
         $array_creation_info->all_list = $array_creation_info->all_list && $item_is_list_item;
 
         if ($item_key_value !== null) {