feat: narrow expectation subject types

extension.neon

@@ -26,6 +26,30 @@ services:
     -
         class: Pest\PHPStan\Analysis\Expectation\ExpectationSemanticAnalyzer
 
+    -
+        class: Pest\PHPStan\Analysis\Expectation\ExpectationNarrowingResolver
+
+    -
+        class: Pest\PHPStan\Type\Pest\ExpectationTypeSpecifyingExtension
+        arguments:
+            className: Pest\Expectation
+        tags:
+            - phpstan.typeSpecifier.methodTypeSpecifyingExtension
+
+    -
+        class: Pest\PHPStan\Type\Pest\ExpectationTypeSpecifyingExtension
+        arguments:
+            className: Pest\Mixins\Expectation
+        tags:
+            - phpstan.typeSpecifier.methodTypeSpecifyingExtension
+
+    -
+        class: Pest\PHPStan\Type\Pest\ExpectationTypeSpecifyingExtension
+        arguments:
+            className: Pest\Expectations\OppositeExpectation
+        tags:
+            - phpstan.typeSpecifier.methodTypeSpecifyingExtension
+
     -
         class: Pest\PHPStan\Type\Pest\PestFunctionReturnTypeExtension
         tags:

src/Analysis/Expectation/ExpectationMatcherRegistry.php

@@ -28,6 +28,11 @@ public function assertedTypeFor(string $methodName, MethodCall $methodCall, Scop
         return $this->assertionRegistry->assertedTypeFor($methodName, $methodCall, $scope);
     }
 
+    public function assertsExactTypeFor(string $methodName, MethodCall $methodCall, Scope $scope): bool
+    {
+        return $this->assertionRegistry->assertsExactTypeFor($methodName, $methodCall, $scope);
+    }
+
     public function metadataFor(string $methodName): ?MatcherSemanticMetadata
     {
         if (array_key_exists($methodName, $this->metadataCache)) {

src/Analysis/Expectation/ExpectationNarrowing.php

@@ -0,0 +1,29 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Pest\PHPStan\Analysis\Expectation;
+
+use PhpParser\Node\Expr;
+use PHPStan\Type\Type;
+
+final readonly class ExpectationNarrowing
+{
+    private function __construct(
+        public Expr $subject,
+        public ?Type $assertedType,
+        public ?Expr $comparedExpr,
+        public bool $loose,
+        public bool $negated,
+    ) {}
+
+    public static function type(Expr $subject, Type $assertedType, bool $negated = false): self
+    {
+        return new self($subject, $assertedType, null, false, $negated);
+    }
+
+    public static function comparison(Expr $subject, Expr $comparedExpr, bool $loose, bool $negated): self
+    {
+        return new self($subject, null, $comparedExpr, $loose, $negated);
+    }
+}

src/Analysis/Expectation/ExpectationNarrowingResolver.php

@@ -0,0 +1,161 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Pest\PHPStan\Analysis\Expectation;
+
+use Pest\Expectation;
+use Pest\Mixins\Expectation as MixinsExpectation;
+use PhpParser\Node\Expr;
+use PhpParser\Node\Expr\FuncCall;
+use PhpParser\Node\Expr\MethodCall;
+use PhpParser\Node\Expr\PropertyFetch;
+use PhpParser\Node\Identifier;
+use PhpParser\Node\Name;
+use PHPStan\Analyser\Scope;
+use PHPStan\Type\ObjectType;
+use PHPStan\Type\Type;
+
+final class ExpectationNarrowingResolver
+{
+    private const string REBIND_METHOD = 'and';
+
+    private const string REBIND_PARAMETER = 'value';
+
+    private const string NEGATE_METHOD = 'not';
+
+    private const array PASSTHROUGH_METHODS = ['when', 'unless', 'sequence', 'match', 'ray'];
+
+    /** @var array<string, bool> Matcher name => uses loose (==) comparison */
+    private const array COMPARISON_METHODS = ['toBe' => false, 'toEqual' => true];
+
+    private const string COMPARISON_PARAMETER = 'expected';
+
+    public function __construct(
+        private readonly ExpectationMatcherRegistry $matcherRegistry,
+    ) {}
+
+    /**
+     * @return list<ExpectationNarrowing>
+     */
+    public function resolve(MethodCall $methodCall, Scope $scope): array
+    {
+        $links = [];
+        $current = $methodCall;
+
+        while (($current instanceof MethodCall || $current instanceof PropertyFetch) && $current->name instanceof Identifier) {
+            $links[] = $current;
+            $current = $current->var;
+        }
+
+        $subject = $this->resolveSubject($current, $scope);
+        if (! $subject instanceof Expr) {
+            return [];
+        }
+
+        $narrowings = [];
+        $negated = false;
+
+        foreach (array_reverse($links) as $link) {
+            if ($link instanceof PropertyFetch) {
+                /** @var Identifier $name */
+                $name = $link->name;
+                if ($name->name === self::NEGATE_METHOD) {
+                    $negated = ! $negated;
+
+                    continue;
+                }
+
+                return $narrowings;
+            }
+
+            if ($link->isFirstClassCallable()) {
+                return $narrowings;
+            }
+
+            /** @var Identifier $name */
+            $name = $link->name;
+            $methodName = $name->name;
+
+            if ($methodName === self::NEGATE_METHOD && $link->getArgs() === []) {
+                $negated = ! $negated;
+
+                continue;
+            }
+
+            if ($methodName === self::REBIND_METHOD) {
+                $rebound = MatcherArgument::first($link, self::REBIND_PARAMETER);
+                if (! $rebound instanceof Expr) {
+                    return $narrowings;
+                }
+
+                if ($this->mayBeExpectation($rebound, $scope)) {
+                    return $narrowings;
+                }
+
+                $subject = $rebound;
+                $negated = false;
+
+                continue;
+            }
+
+            if (in_array($methodName, self::PASSTHROUGH_METHODS, true)) {
+                continue;
+            }
+
+            if (! method_exists(MixinsExpectation::class, $methodName)) {
+                return $narrowings;
+            }
+
+            if (isset(self::COMPARISON_METHODS[$methodName])) {
+                $compared = MatcherArgument::first($link, self::COMPARISON_PARAMETER);
+                if ($compared instanceof Expr) {
+                    $narrowings[] = ExpectationNarrowing::comparison($subject, $compared, self::COMPARISON_METHODS[$methodName], $negated);
+                }
+
+                $negated = false;
+
+                continue;
+            }
+
+            if ($negated && ! $this->matcherRegistry->assertsExactTypeFor($methodName, $link, $scope)) {
+                $negated = false;
+
+                continue;
+            }
+
+            $assertedType = $this->matcherRegistry->assertedTypeFor($methodName, $link, $scope);
+
+            if ($assertedType instanceof Type) {
+                $narrowings[] = ExpectationNarrowing::type($subject, $assertedType, $negated);
+            }
+
+            $negated = false;
+        }
+
+        return $narrowings;
+    }
+
+    /** @return bool True when and() may unwrap the argument to an inner value we cannot track */
+    private function mayBeExpectation(Expr $expr, Scope $scope): bool
+    {
+        return ! new ObjectType(Expectation::class)->isSuperTypeOf($scope->getType($expr))->no();
+    }
+
+    private function resolveSubject(Expr $root, Scope $scope): ?Expr
+    {
+        if (! $root instanceof FuncCall || ! $root->name instanceof Name) {
+            return null;
+        }
+
+        if ($root->name->toLowerString() !== 'expect') {
+            return null;
+        }
+
+        if (! new ObjectType(Expectation::class)->isSuperTypeOf($scope->getType($root))->yes()) {
+            return null;
+        }
+
+        return $root->getArgs()[0]->value ?? null;
+    }
+}

src/Analysis/Expectation/MatcherArgument.php

@@ -0,0 +1,43 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Pest\PHPStan\Analysis\Expectation;
+
+use PhpParser\Node\Expr;
+use PhpParser\Node\Expr\MethodCall;
+use PhpParser\Node\Identifier;
+
+final class MatcherArgument
+{
+    /**
+     * @param  string  $parameterName  Declared name of the first parameter, for named arguments
+     * @return Expr|null Null whenever the argument cannot be read with certainty
+     */
+    public static function first(MethodCall $methodCall, string $parameterName): ?Expr
+    {
+        if ($methodCall->isFirstClassCallable()) {
+            return null;
+        }
+
+        foreach ($methodCall->getArgs() as $position => $argument) {
+            if ($argument->unpack) {
+                return null;
+            }
+
+            if ($argument->name instanceof Identifier) {
+                if ($argument->name->name === $parameterName) {
+                    return $argument->value;
+                }
+
+                continue;
+            }
+
+            if ($position === 0) {
+                return $argument->value;
+            }
+        }
+
+        return null;
+    }
+}

src/Analysis/Expectation/MatcherAssertionRegistry.php

@@ -4,6 +4,7 @@
 
 namespace Pest\PHPStan\Analysis\Expectation;
 
+use PhpParser\Node\Expr;
 use PhpParser\Node\Expr\MethodCall;
 use PHPStan\Analyser\Scope;
 use PHPStan\Type\Accessory\AccessoryArrayListType;
@@ -12,6 +13,7 @@
 use PHPStan\Type\BooleanType;
 use PHPStan\Type\CallableType;
 use PHPStan\Type\Constant\ConstantBooleanType;
+use PHPStan\Type\Constant\ConstantStringType;
 use PHPStan\Type\FloatType;
 use PHPStan\Type\IntegerType;
 use PHPStan\Type\IntersectionType;
@@ -80,6 +82,8 @@ final class MatcherAssertionRegistry
         'toBeResource' => self::RESOURCE,
     ];
 
+    private const string INSTANCE_OF_PARAMETER = 'class';
+
     /** @var array<string, Type> */
     private array $staticAssertedTypeCache = [];
 
@@ -146,26 +150,43 @@ public function assertedTypeFor(string $methodName, MethodCall $methodCall, Scop
         return $assertedType;
     }
 
+    /** @return bool True when the asserted type mirrors the matcher exactly, so it may also be removed */
+    public function assertsExactTypeFor(string $methodName, MethodCall $methodCall, Scope $scope): bool
+    {
+        if ($this->assertionFor($methodName) !== self::INSTANCE_OF) {
+            return true;
+        }
+
+        return count($this->constantClassNames($methodCall, $scope)) === 1;
+    }
+
     private function resolveToBeInstanceOf(MethodCall $methodCall, Scope $scope): Type
     {
-        $args = $methodCall->getArgs();
+        $classNames = $this->constantClassNames($methodCall, $scope);
 
-        if ($args === []) {
+        if ($classNames === []) {
             return new ObjectWithoutClassType;
         }
 
-        $classType = $scope->getType($args[0]->value);
-        $classNames = $classType->getConstantStrings();
+        $objectTypes = array_map(
+            static fn (ConstantStringType $name): ObjectType => new ObjectType($name->getValue()),
+            $classNames
+        );
+
+        return TypeCombinator::union(...$objectTypes);
+    }
 
-        if ($classNames !== []) {
-            $objectTypes = array_map(
-                static fn ($name): ObjectType => new ObjectType($name->getValue()),
-                $classNames
-            );
+    /**
+     * @return list<ConstantStringType>
+     */
+    private function constantClassNames(MethodCall $methodCall, Scope $scope): array
+    {
+        $class = MatcherArgument::first($methodCall, self::INSTANCE_OF_PARAMETER);
 
-            return TypeCombinator::union(...$objectTypes);
+        if (! $class instanceof Expr) {
+            return [];
         }
 
-        return new ObjectWithoutClassType;
+        return $scope->getType($class)->getConstantStrings();
     }
 }