Merge similar dynamic return type extensions into a single class

src/Type/Php/ArrayFirstLastDynamicReturnTypeExtension.php

@@ -6,6 +6,7 @@
 use PHPStan\Analyser\Scope;
 use PHPStan\DependencyInjection\AutowiredService;
 use PHPStan\Reflection\FunctionReflection;
+use PHPStan\ShouldNotHappenException;
 use PHPStan\Type\DynamicFunctionReturnTypeExtension;
 use PHPStan\Type\NullType;
 use PHPStan\Type\Type;
@@ -19,7 +20,12 @@ final class ArrayFirstLastDynamicReturnTypeExtension implements DynamicFunctionR
 
 	public function isFunctionSupported(FunctionReflection $functionReflection): bool
 	{
-		return in_array($functionReflection->getName(), ['array_first', 'array_last'], true);
+		return in_array($functionReflection->getName(), [
+			'array_key_first',
+			'array_key_last',
+			'array_first',
+			'array_last',
+		], true);
 	}
 
 	public function getTypeFromFunctionCall(FunctionReflection $functionReflection, FuncCall $functionCall, Scope $scope): ?Type
@@ -37,13 +43,24 @@ public function getTypeFromFunctionCall(FunctionReflection $functionReflection,
 			return new NullType();
 		}
 
-		$valueType = $argType->getIterableValueType();
+		switch ($functionReflection->getName()) {
+			case 'array_key_first':
+			case 'array_key_last':
+				$resultType = $argType->getIterableKeyType();
+				break;
+			case 'array_first':
+			case 'array_last':
+				$resultType = $argType->getIterableValueType();
+				break;
+			default:
+				throw new ShouldNotHappenException();
+		}
 
 		if ($iterableAtLeastOnce->yes()) {
-			return $valueType;
+			return $resultType;
 		}
 
-		return TypeCombinator::union($valueType, new NullType());
+		return TypeCombinator::union($resultType, new NullType());
 	}
 
 }