More precise array_key_first/array_key_last inference

src/Analyser/TypeSpecifier.php

@@ -725,7 +725,44 @@
 
 			if ($context->null()) {
 				$specifiedTypes = $this->specifyTypesInCondition($scope->exitFirstLevelStatements(), $expr->expr, $context)->setRootExpr($expr);
+			} else {
+				$specifiedTypes = $this->specifyTypesInCondition($scope->exitFirstLevelStatements(), $expr->var, $context)->setRootExpr($expr);
+			}
+
+			// infer $arr[$key] after $key = array_key_first/last($arr)
+			if (
+				$expr->expr instanceof FuncCall
+				&& $expr->expr->name instanceof Name
+				&& in_array($expr->expr->name->toLowerString(), ['array_key_first', 'array_key_last'], true)
+				&& count($expr->expr->getArgs()) >= 1
+			) {
+				$arrayArg = $expr->expr->getArgs()[0]->value;
+				$arrayType = $scope->getType($arrayArg);
+
+				if ($arrayType->isArray()->yes()) {
+					if ($context->true()) {
+						$specifiedTypes = $specifiedTypes->unionWith(
+							$this->create($arrayArg, new NonEmptyArrayType(), TypeSpecifierContext::createTrue(), $scope),
+						);
+						$isNonEmpty = true;
+					} else {
+						$isNonEmpty = $arrayType->isIterableAtLeastOnce()->yes();
+					}
+
+					if ($isNonEmpty) {
+						$dimFetch = new ArrayDimFetch($arrayArg, $expr->var);
+						$iterableValueType = $expr->expr->name->toLowerString() === 'array_key_first'
+							? $arrayType->getIterableValueType()
+							: $arrayType->getIterableValueType();
 
+						$specifiedTypes = $specifiedTypes->unionWith(
+							$this->create($dimFetch, $iterableValueType, TypeSpecifierContext::createTrue(), $scope),
+						);
+					}
+				}
+			}
+
+			if ($context->null()) {
 				// infer $arr[$key] after $key = array_rand($arr)
 				if (
 					$expr->expr instanceof FuncCall
@@ -755,30 +792,6 @@
 					}
 				}
 
-				// infer $arr[$key] after $key = array_key_first/last($arr)
-				if (
-					$expr->expr instanceof FuncCall
-					&& $expr->expr->name instanceof Name
-					&& in_array($expr->expr->name->toLowerString(), ['array_key_first', 'array_key_last'], true)
-					&& count($expr->expr->getArgs()) >= 1
-				) {
-					$arrayArg = $expr->expr->getArgs()[0]->value;
-					$arrayType = $scope->getType($arrayArg);
-					if (
-						$arrayType->isArray()->yes()
-						&& $arrayType->isIterableAtLeastOnce()->yes()
-					) {
-						$dimFetch = new ArrayDimFetch($arrayArg, $expr->var);
-						$iterableValueType = $expr->expr->name->toLowerString() === 'array_key_first'
-							? $arrayType->getIterableValueType()
-							: $arrayType->getIterableValueType();
-
-						return $specifiedTypes->unionWith(
-							$this->create($dimFetch, $iterableValueType, TypeSpecifierContext::createTrue(), $scope),
-						);
-					}
-				}
-
 				// infer $list[$count] after $count = count($list) - 1
 				if (
 					$expr->expr instanceof Expr\BinaryOp\Minus
@@ -806,8 +819,6 @@
 				return $specifiedTypes;
 			}
 
-			$specifiedTypes = $this->specifyTypesInCondition($scope->exitFirstLevelStatements(), $expr->var, $context)->setRootExpr($expr);
-
 			if ($context->true()) {
 				// infer $arr[$key] after $key = array_search($needle, $arr)
 				if (

tests/PHPStan/Analyser/nsrt/bug-14081.php

@@ -0,0 +1,68 @@
+<?php
+
+declare(strict_types = 1);
+
+namespace Bug14081;
+
+use function PHPStan\Testing\assertType;
+
+/** @param list<string> $array */
+function first(array $array): mixed
+{
+	if (($key = array_key_first($array))) {
+		assertType('int<1, max>', $key);
+		assertType('non-empty-list<string>', $array);
+		assertType('string', $array[$key]);
+		return $array[$key];
+	}
+	return null;
+}
+
+/** @param list<string> $array */
+function last(array $array): mixed
+{
+	if (($key = array_key_last($array))) {
+		assertType('int<1, max>', $key);
+		assertType('non-empty-list<string>', $array);
+		assertType('string', $array[$key]);
+		return $array[$key];
+	}
+	return null;
+}
+
+function maybeNonEmpty(): void
+{
+	if (rand(0,1)) {
+		$array = ['one', 'two'];
+	} else {
+		$array = [];
+	}
+	assertType("array{}|array{'one', 'two'}", $array);
+	$key = array_key_last($array);
+	assertType('0|1|null', $key);
+	assertType("'one'|'two'", $array[$key]);
+}
+
+/** @param list<string> $array */
+function firstNotNull(array $array): mixed
+{
+	if (($key = array_key_first($array)) !== null) {
+		assertType('int<0, max>', $key); // could be int<1, max>
+		assertType('list<string>', $array); // could be non-empty-list<string>
+		assertType('string', $array[$key]);
+		return $array[$key];
+	}
+	return null;
+}
+
+/** @param list<string> $array */
+function lastNotNull(array $array): mixed
+{
+	if (($key = array_key_last($array)) !== null) {
+		assertType('int<0, max>', $key); // could be int<1, max>
+		assertType('list<string>', $array); // could be non-empty-list<string>
+		assertType('string', $array[$key]);
+		return $array[$key];
+	}
+	return null;
+}