fix multiple return type bugs

src/EventHandler/Iter/Count/FunctionReturnTypeProvider.php

@@ -19,7 +19,7 @@ final class FunctionReturnTypeProvider implements FunctionReturnTypeProviderInte
     public static function getFunctionIds(): array
     {
         return [
-            'psl\str\count',
+            'psl\iter\count',
         ];
     }
 
@@ -56,7 +56,22 @@ final class FunctionReturnTypeProvider implements FunctionReturnTypeProviderInte
 
         // array{foo: bar} -> literal-int(1)
         if ($array_argument_type instanceof Type\Atomic\TKeyedArray) {
-            return Type::getInt(false, count($array_argument_type->properties));
+            // Psalm allows extra properties in keyed arrays, so we can't return a literal integer
+            // for this.
+            //
+            // return Type::getInt(false, count($array_argument_type->properties));
+
+            if (count($array_argument_type->properties) >= 1) {
+                return Type::getPositiveInt();
+            }
+
+            return Type::getInt();
+        }
+
+        if ($array_argument_type instanceof Type\Atomic\TArray) {
+            if ($array_argument_type->type_params[0]->isEmpty() && $array_argument_type->type_params[1]->isEmpty()) {
+                return Type::getInt(false, 0);
+            }
         }
 
         return Type::getInt();

src/EventHandler/Iter/FirstKey/FunctionReturnTypeProvider.php

@@ -0,0 +1,55 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Psl\Psalm\EventHandler\Iter\FirstKey;
+
+use Psalm\Plugin\EventHandler\Event\FunctionReturnTypeProviderEvent;
+use Psalm\Plugin\EventHandler\FunctionReturnTypeProviderInterface;
+use Psalm\Type;
+use Psl\Psalm\Argument;
+
+final class FunctionReturnTypeProvider implements FunctionReturnTypeProviderInterface
+{
+    /**
+     * @return non-empty-list<lowercase-string>
+     */
+    public static function getFunctionIds(): array
+    {
+        return [
+            'psl\iter\first_key',
+        ];
+    }
+
+    public static function getFunctionReturnType(FunctionReturnTypeProviderEvent $event): ?Type\Union
+    {
+        $argument_type = Argument::getType($event->getCallArgs(), $event->getStatementsSource(), 0);
+        if (null === $argument_type) {
+            return null;
+        }
+
+        $array_argument_type = $argument_type->getAtomicTypes()['array'] ?? null;
+        if (null === $array_argument_type) {
+            return null;
+        }
+
+        if ($array_argument_type instanceof Type\Atomic\TNonEmptyArray) {
+            return clone $array_argument_type->type_params[0];
+        }
+
+        if ($array_argument_type instanceof Type\Atomic\TNonEmptyList) {
+            return Type::getInt();
+        }
+
+        if ($array_argument_type instanceof Type\Atomic\TKeyedArray) {
+            // TODO(azjezz): add support for this once psalm starts enforcing the shape order ( if ever ).
+            //
+            // foreach ($properties as $property) {
+            //     return clone $property;
+            // }
+            return clone $array_argument_type->getGenericKeyType();
+        }
+
+        return null;
+    }
+}

src/EventHandler/Iter/LastKey/FunctionReturnTypeProvider.php

@@ -0,0 +1,59 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Psl\Psalm\EventHandler\Iter\LastKey;
+
+use Psalm\Plugin\EventHandler\Event\FunctionReturnTypeProviderEvent;
+use Psalm\Plugin\EventHandler\FunctionReturnTypeProviderInterface;
+use Psalm\Type;
+use Psl\Psalm\Argument;
+
+final class FunctionReturnTypeProvider implements FunctionReturnTypeProviderInterface
+{
+    /**
+     * @return non-empty-list<lowercase-string>
+     */
+    public static function getFunctionIds(): array
+    {
+        return [
+            'psl\iter\last_key',
+        ];
+    }
+
+    public static function getFunctionReturnType(FunctionReturnTypeProviderEvent $event): ?Type\Union
+    {
+        $argument_type = Argument::getType($event->getCallArgs(), $event->getStatementsSource(), 0);
+        if (null === $argument_type) {
+            return null;
+        }
+
+        $array_argument_type = $argument_type->getAtomicTypes()['array'] ?? null;
+        if (null === $array_argument_type) {
+            return null;
+        }
+
+        if ($array_argument_type instanceof Type\Atomic\TNonEmptyArray) {
+            return clone $array_argument_type->type_params[0];
+        }
+
+        if ($array_argument_type instanceof Type\Atomic\TNonEmptyList) {
+            return Type::getInt();
+        }
+
+        if ($array_argument_type instanceof Type\Atomic\TKeyedArray) {
+            // TODO(azjezz): add support for this once psalm starts enforcing the shape order ( if ever ).
+            //
+            // $properties = $array_argument_type->properties;
+            // $last_property = null;
+            // foreach ($properties as $property) {
+            //     $last_property = $property;
+            // }
+            //
+            // return clone $last_property;
+            return clone $array_argument_type->getGenericKeyType();
+        }
+
+        return null;
+    }
+}

src/EventHandler/Str/Repeat/FunctionReturnTypeProvider.php

@@ -26,7 +26,7 @@ final class FunctionReturnTypeProvider implements FunctionReturnTypeProviderInte
     public static function getFunctionReturnType(FunctionReturnTypeProviderEvent $event): ?Type\Union
     {
         $argument_type = Argument::getType($event->getCallArgs(), $event->getStatementsSource(), 0);
-        if ($argument_type === null || !$argument_type->hasLowercaseString()) {
+        if ($argument_type === null) {
             return Type::getString();
         }
 
@@ -43,12 +43,12 @@ final class FunctionReturnTypeProvider implements FunctionReturnTypeProviderInte
             return new Type\Union([new Type\Atomic\TLowercaseString()]);
         }
 
-        if ($string_argument_type instanceof Type\Atomic\TLiteralString) {
+        if ($argument_type->hasLiteralString()) {
             $multiplier_argument_type = Argument::getType($event->getCallArgs(), $event->getStatementsSource(), 1);
             if (null !== $multiplier_argument_type && $multiplier_argument_type->hasLiteralInt()) {
                 /** @psalm-suppress MissingThrowsDocblock */
                 return Type::getString(str_repeat(
-                    $string_argument_type->value,
+                    $argument_type->getSingleStringLiteral()->value,
                     $multiplier_argument_type->getSingleIntLiteral()->value
                 ));
             }

src/Plugin.php

@@ -31,7 +31,9 @@ final class Plugin implements PluginEntryPointInterface
     {
         // Psl\Iter hooks
         yield EventHandler\Iter\First\FunctionReturnTypeProvider::class;
+        yield EventHandler\Iter\FirstKey\FunctionReturnTypeProvider::class;
         yield EventHandler\Iter\Last\FunctionReturnTypeProvider::class;
+        yield EventHandler\Iter\LastKey\FunctionReturnTypeProvider::class;
         yield EventHandler\Iter\Count\FunctionReturnTypeProvider::class;
 
         // Psl\Regex hooks