Improve partitionT inference

Author: klimick

src/Fp/Psalm/Hook/FunctionReturnTypeProvider/PartitionFunctionReturnTypeProvider.php

@@ -56,16 +56,28 @@ public static function getFunctionReturnType(FunctionReturnTypeProviderEvent $ev
                         fn() => proveOf($event->getStatementsSource(), StatementsAnalyzer::class),
                         fn() => Option::some($type_params),
                     ))
-                    ->map(fn(ArrayList $args) => $args
-                        ->mapN(ctor(RefinementContext::class))
-                        ->map(RefineByPredicate::for(...))
-                        ->map(fn(CollectionTypeParams $params) => new Union(
-                            [new TList($params->val_type)],
-                        ))
-                        ->appended(new Union(
-                            [new TList($type_params->val_type)],
-                        ))
-                        ->toList())
+                    ->map(function(ArrayList $args) use ($type_params) {
+                        $init_types = $args
+                            ->mapN(ctor(RefinementContext::class))
+                            ->map(RefineByPredicate::for(...))
+                            ->map(fn(CollectionTypeParams $params) => $params->val_type);
+
+                        $last_type = $init_types
+                            ->fold($type_params->val_type)(
+                                function(Union $last, Union $current) {
+                                    $cloned = clone $last;
+                                    $cloned->removeType($current->getId());
+
+                                    return $cloned;
+                                },
+                            );
+
+                        return $init_types
+                            ->appended($last_type)
+                            ->map(fn(Union $type) => [new TList($type)])
+                            ->map(ctor(Union::class))
+                            ->toList();
+                    })
             )
             ->flatMap(fn(array $partitions) => asNonEmptyArray($partitions))
             ->map(function(array $non_empty_partitions) {

tests/Static/Functions/Collection/PartitionStaticTest.php

@@ -4,6 +4,10 @@
 
 namespace Tests\Static\Functions\Collection;
 
+use Tests\Mock\Bar;
+use Tests\Mock\Baz;
+use Tests\Mock\Foo;
+
 use function Fp\Collection\partitionT;
 
 final class PartitionStaticTest
@@ -29,4 +33,13 @@ public function testPartitionWithTwoPredicates(array $nums): array
             fn(int $v) => $v % 2 === 1,
         );
     }
+
+    /**
+     * @param list<Foo|Bar|Baz> $list
+     * @return array{list<Foo>, list<Bar>, list<Baz>}
+     */
+    public function testExhaustiveInference(array $list): array
+    {
+        return partitionT($list, fn($i) => $i instanceof Foo, fn($i) => $i instanceof Bar);
+    }
 }