Improve return type inference of Collection::groupBy.

extension.neon

@@ -289,6 +289,10 @@ services:
         class: Larastan\Larastan\ReturnTypes\CollectionFilterDynamicReturnTypeExtension
         tags:
             - phpstan.broker.dynamicMethodReturnTypeExtension
+    -
+        class: Larastan\Larastan\ReturnTypes\EnumerableGroupByReturnTypeExtension
+        tags:
+            - phpstan.broker.dynamicMethodReturnTypeExtension
     -
         class: Larastan\Larastan\ReturnTypes\CollectionWhereNotNullDynamicReturnTypeExtension
         tags:

src/ReturnTypes/CollectionGenericStaticMethodDynamicMethodReturnTypeExtension.php

@@ -39,7 +39,7 @@ public function isMethodSupported(MethodReflection $methodReflection): bool
         $methods = [
             'chunk', 'chunkWhile', 'collapse', 'combine',
             'crossJoin', 'flatMap', 'flip',
-            'groupBy', 'keyBy', 'keys',
+            'keyBy', 'keys',
             'make', 'map', 'mapInto',
             'mapToDictionary', 'mapToGroups',
             'mapWithKeys', 'mergeRecursive',

src/ReturnTypes/EnumerableGroupByReturnTypeExtension.php

@@ -0,0 +1,148 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Larastan\Larastan\ReturnTypes;
+
+use Illuminate\Support\Enumerable;
+use PHPStan\Analyser\OutOfClassScope;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Type\BenevolentUnionType;
+use PHPStan\Type\Constant\ConstantStringType;
+use PHPStan\Type\DynamicMethodReturnTypeExtension;
+use PHPStan\Type\Generic\GenericObjectType;
+use PHPStan\Type\IntegerType;
+use PHPStan\Type\MixedType;
+use PHPStan\Type\StringType;
+use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
+use PhpParser\Node\Expr\MethodCall;
+
+use function array_reverse;
+use function count;
+
+class EnumerableGroupByReturnTypeExtension implements DynamicMethodReturnTypeExtension
+{
+    public function getClass(): string
+    {
+        return Enumerable::class;
+    }
+
+    public function isMethodSupported(MethodReflection $methodReflection): bool
+    {
+        return $methodReflection->getName() === 'groupBy';
+    }
+
+    public function getTypeFromMethodCall(
+        MethodReflection $methodReflection,
+        MethodCall $methodCall,
+        Scope $scope
+    ): ?Type {
+        if (count($methodCall->getArgs()) < 1) {
+            return null;
+        }
+
+        $calledOnType = $scope->getType($methodCall->var);
+        $objectClassReflections = $calledOnType->getObjectClassReflections();
+
+        if (!isset($objectClassReflections[0])) {
+            return null;
+        }
+
+        $groupByType = $scope->getType($methodCall->getArgs()[0]->value);
+        $propertyTypes = [];
+
+        $collectionName = $objectClassReflections[0]->getName();
+
+        $valueType = new MixedType();
+
+        if ($objectClassReflections[0]->isGeneric()) {
+            $tValueType = $methodReflection->getDeclaringClass()->getActiveTemplateTypeMap()->getType('TValue');
+
+            if ($tValueType !== null) {
+                $valueType = $tValueType;
+            }
+        }
+
+        if ($groupByType->isString()->yes()) {
+            $propertyTypes[] = $groupByType;
+        } elseif ($groupByType->isConstantArray()->yes()) {
+            if ($groupByType->isIterableAtLeastOnce()->no()) {
+                return $this->unknownKeysType($collectionName, $valueType);
+            }
+
+            $valuesArray = $groupByType->getConstantArrays()[0]->getValueTypes();
+
+            foreach ($valuesArray as $valuesType) {
+                $propertyTypes[] = $valuesType;
+            }
+        } else {
+            return $this->unknownKeysType($collectionName, new MixedType());
+        }
+
+
+        $innerKeyType = new IntegerType();
+
+        if (count($methodCall->getArgs()) >= 2) {
+            $preserveType = $scope->getType($methodCall->getArgs()[1]->value);
+            $tKeyType = $methodReflection->getDeclaringClass()->getActiveTemplateTypeMap()->getType('TKey');
+
+            if ($tKeyType === null) {
+                $innerKeyType = new BenevolentUnionType([new IntegerType(), new StringType()]);
+            } elseif ($preserveType->isTrue()->yes()) {
+                $innerKeyType = $tKeyType;
+            } elseif ($preserveType->isTrue()->maybe()) {
+                $innerKeyType = TypeCombinator::union($tKeyType, new IntegerType());
+            }
+        }
+
+
+        $inner = new GenericObjectType($collectionName, [$innerKeyType, $valueType]);
+
+        foreach (array_reverse($propertyTypes) as $propertyType) {
+            if (count($propertyType->getConstantStrings()) > 0) {
+                $types = [];
+                foreach ($propertyType->getConstantStrings() as $constantString) {
+                    if ($valueType->hasProperty($constantString->getValue())->yes()) {
+                        $types[] = $valueType->getProperty($constantString->getValue(), new OutOfClassScope())->getReadableType();
+                    } elseif ($valueType->hasOffsetValueType($constantString)->yes()) {
+                        $types[] = $valueType->getOffsetValueType($constantString);
+                    }
+                }
+
+                if (count($types) === 0) {
+                    $keyType = new ConstantStringType('');
+                } else {
+                    $keyType = TypeCombinator::union(...$types);
+                }
+            } elseif ($propertyType->isCallable()->yes()) {
+                $keyType = $propertyType->getCallableParametersAcceptors(new OutOfClassScope())[0]->getReturnType();
+            } else {
+                $keyType = new BenevolentUnionType([new IntegerType(), new StringType()]);
+            }
+
+            if ($keyType instanceof MixedType) {
+                $keyType = new BenevolentUnionType([new IntegerType(), new StringType()]);
+            }
+
+            $inner = new GenericObjectType($collectionName, [$keyType, $inner]);
+        }
+
+        return $inner;
+    }
+
+    /**
+     * @param class-string $name
+     */
+    private function unknownKeysType(string $name, Type $inner): Type
+    {
+        return new GenericObjectType($name, [
+            new BenevolentUnionType([new IntegerType(), new StringType()]),
+            new GenericObjectType($name, [
+                new BenevolentUnionType([new IntegerType(), new StringType()]),
+                $inner
+            ])
+        ]);
+    }
+}

tests/Type/GeneralTypeTest.php

@@ -59,6 +59,7 @@ public static function dataFileAsserts(): iterable
         yield from self::gatherAssertTypes(__DIR__.'/data/view-exists.php');
         yield from self::gatherAssertTypes(__DIR__.'/data/view.php');
         yield from self::gatherAssertTypes(__DIR__.'/data/where-relation.php');
+        yield from self::gatherAssertTypes(__DIR__.'/data/collection-group-by.php');
 
         if (version_compare(LARAVEL_VERSION, '10.15.0', '>=')) {
             yield from self::gatherAssertTypes(__DIR__.'/data/model-l10-15.php');

tests/Type/data/collection-generic-static-methods-l948.php

@@ -44,8 +44,8 @@
 assertType('Illuminate\Support\Collection<App\User, int>', $collection->flip());
 assertType('Illuminate\Support\Collection<int, string>', $items->flip());
 
-assertType('Illuminate\Database\Eloquent\Collection<(int|string), Illuminate\Database\Eloquent\Collection<(int|string), App\User>>', $collection->groupBy('id'));
-assertType('Illuminate\Support\Collection<(int|string), Illuminate\Support\Collection<(int|string), int>>', $items->groupBy('id'));
+assertType('Illuminate\Database\Eloquent\Collection<int, Illuminate\Database\Eloquent\Collection<int, App\User>>', $collection->groupBy('id'));
+assertType("Illuminate\Support\Collection<'', Illuminate\Support\Collection<int, int>>", $items->groupBy('id'));
 
 assertType('Illuminate\Database\Eloquent\Collection<(int|string), App\User>', $collection->keyBy(fn (User $user, int $key): string => $user->email));
 
@@ -183,7 +183,7 @@
     })
 );
 
-assertType('Illuminate\Support\Collection<(int|string), Illuminate\Support\Collection<(int|string), array{id: int, type: string}>>', collect([
+assertType('Illuminate\Support\Collection<string, Illuminate\Support\Collection<int, array{id: int, type: string}>>', collect([
     [
         'id'   => 1,
         'type' => 'A',

tests/Type/data/collection-group-by.php

@@ -0,0 +1,93 @@
+<?php
+
+namespace CollectionGroupBy;
+
+use App\User;
+use Illuminate\Support\Collection;
+use function PHPStan\Testing\assertType;
+
+/**
+ * @param Collection $untyped
+ * @param Collection<int, User> $users
+ * @param Collection<string, User> $stringUsers
+ * @param Collection<string, string> $stringStrings
+ * @param 'id'|'name' $key
+ * @param bool $preserve
+ * @param array<int, string> $keys
+ */
+function test(
+    Collection $untyped,
+    Collection $users,
+    Collection $stringUsers,
+    Collection $stringStrings,
+    string $key,
+    bool $preserve,
+    array $keys
+): void {
+    assertType(
+        'Illuminate\Support\Collection<(int|string), Illuminate\Support\Collection<int, mixed>>',
+        $untyped->groupBy('id')
+    );
+    assertType(
+        'Illuminate\Support\Collection<(int|string), Illuminate\Support\Collection<int, mixed>>',
+        $untyped->groupBy(['id'])
+    );
+    assertType(
+        'Illuminate\Support\Collection<(int|string), Illuminate\Support\Collection<(int|string), mixed>>',
+        $untyped->groupBy('id', $preserve)
+    );
+    assertType(
+        'Illuminate\Support\Collection<int, Illuminate\Support\Collection<int, App\User>>',
+        $users->groupBy('id')
+    );
+    assertType(
+        'Illuminate\Support\Collection<int|string, Illuminate\Support\Collection<int, App\User>>',
+        $users->groupBy($key)
+    );
+    assertType(
+        'Illuminate\Support\Collection<int, Illuminate\Support\Collection<string, Illuminate\Support\Collection<int, App\User>>>',
+        $users->groupBy(['id', 'name'])
+    );
+    assertType(
+        'Illuminate\Support\Collection<int, Illuminate\Support\Collection<int<0, 1>, Illuminate\Support\Collection<int, App\User>>>',
+        $users->groupBy(['id', fn ($user) => rand(0, 1)])
+    );
+    assertType(
+        'Illuminate\Support\Collection<int, Illuminate\Support\Collection<string, App\User>>',
+        $stringUsers->groupBy('id', true)
+    );
+    assertType(
+        'Illuminate\Support\Collection<int, Illuminate\Support\Collection<int|string, App\User>>',
+        $stringUsers->groupBy('id', $preserve)
+    );
+    assertType(
+        'Illuminate\Support\Collection<int, Illuminate\Support\Collection<string, Illuminate\Support\Collection<int|string, App\User>>>',
+        $stringUsers->groupBy(['id', 'name'], $preserve)
+    );
+    assertType(
+        "Illuminate\Support\Collection<'', Illuminate\Support\Collection<int, string>>",
+        $stringStrings->groupBy(['id'])
+    );
+    assertType(
+        'Illuminate\Support\Collection<(int|string), Illuminate\Support\Collection<(int|string), mixed>>',
+        $users->groupBy($keys)
+    );
+
+    assertType(
+        'Illuminate\Support\Collection<(int|string), Illuminate\Support\Collection<(int|string), App\User>>',
+        $users->groupBy([])
+    );
+}
+
+/**
+ * @param ?Collection<int, string> $collection
+ */
+function testNullable(?Collection $collection): void
+{
+    $res = $collection->groupBy('key');
+
+    assertType(
+        'Illuminate\Support\Collection<(int|string), Illuminate\Support\Collection<(int|string), string>>',
+        $res
+    );
+}

tests/Type/data/collection-stubs.php

@@ -12,6 +12,7 @@
 /** @var SupportCollection<string, int> $items */
 /** @var SupportCollection<int, User> $collectionOfUsers */
 /** @var User $user */
+/** @var bool $preserve */
 assertType('Illuminate\Database\Eloquent\Collection<int, App\User>', User::all()->each(function (User $user, int $key): void {
 }));
 
@@ -51,9 +52,21 @@
 ]);
 
 $foo
-    ->groupBy('type')
+    ->groupBy('type', $preserve)
     ->map(function ($grouped, $groupKey): array {
-        assertType('(int|string)', $groupKey);
+        assertType('string', $groupKey);
+    });
+
+$foo
+    ->groupBy('type', true)
+    ->map(function ($grouped, $groupKey): array {
+        assertType('string', $groupKey);
+    });
+
+$foo
+    ->groupBy('type', false)
+    ->map(function ($grouped, $groupKey): array {
+        assertType('string', $groupKey);
     });
 
 assertType('App\User|null', $collection->first());