Track $arr[$key] existence across array_search/array_find_key via conditional expression holders

src/Analyser/TypeSpecifier.php

@@ -849,34 +849,77 @@ public function specifyTypesInCondition(
 
 					if ($isNonEmpty) {
 						$dimFetch = new ArrayDimFetch($arrayArg, $expr->var);
-
 						$specifiedTypes = $specifiedTypes->unionWith(
 							$this->create($dimFetch, $arrayType->getIterableValueType(), TypeSpecifierContext::createTrue(), $scope),
 						);
 					} elseif ($expr->var instanceof Expr\Variable && is_string($expr->var->name)) {
-						// The array might be empty here, so we cannot register
-						// $arr[$key] unconditionally. Attach a conditional holder
-						// that fires once the user narrows $key to non-null
-						// (e.g. `if ($key !== null)`), giving the deep-write
-						// path the same shape information `isset($arr[$key])`
-						// would have provided.
 						$keyType = $scope->getType($expr->expr);
 						$nonNullKeyType = TypeCombinator::removeNull($keyType);
-						if (!$nonNullKeyType instanceof NeverType && !$keyType->isNull()->yes()) {
-							$dimFetch = new ArrayDimFetch($arrayArg, $expr->var);
-							$dimFetchString = $this->exprPrinter->printExpr($dimFetch);
-							$keyExprString = $this->exprPrinter->printExpr($expr->var);
+						if (!$nonNullKeyType instanceof NeverType) {
+							$specifiedTypes = $specifiedTypes->unionWith(
+								$this->createArrayDimFetchConditionalExpressionHolder($expr->var, $arrayArg, $nonNullKeyType, $arrayType->getIterableValueType()),
+							);
+						}
+					}
+				}
+			}
+
+			// infer $arr[$key] after $key = array_search($needle, $arr) or $key = array_find_key($arr, $callback)
+			if (
+				$expr->expr instanceof FuncCall
+				&& $expr->expr->name instanceof Name
+				&& !$expr->expr->isFirstClassCallable()
+				&& count($expr->expr->getArgs()) >= 2
+			) {
+				$funcName = $expr->expr->name->toLowerString();
+				$arrayArg = null;
+				$sentinelType = null;
+				$isStrictArraySearch = false;
+
+				if ($funcName === 'array_search') {
+					$arrayArg = $expr->expr->getArgs()[1]->value;
+					$sentinelType = new ConstantBooleanType(false);
+					$isStrictArraySearch = count($expr->expr->getArgs()) >= 3 && $scope->getType($expr->expr->getArgs()[2]->value)->isTrue()->yes();
+				} elseif ($funcName === 'array_find_key') {
+					$arrayArg = $expr->expr->getArgs()[0]->value;
+					$sentinelType = new NullType();
+				}
 
-							$holder = new ConditionalExpressionHolder(
-								[$keyExprString => ExpressionTypeHolder::createYes($expr->var, $nonNullKeyType)],
-								ExpressionTypeHolder::createYes($dimFetch, $arrayType->getIterableValueType()),
+				if ($arrayArg !== null) {
+					$arrayType = $scope->getType($arrayArg);
+
+					if ($arrayType->isArray()->yes()) {
+						if ($context->true()) {
+							$specifiedTypes = $specifiedTypes->unionWith(
+								$this->create($arrayArg, new NonEmptyArrayType(), TypeSpecifierContext::createTrue(), $scope),
 							);
 
+							$dimFetch = new ArrayDimFetch($arrayArg, $expr->var);
+
+							if ($isStrictArraySearch) {
+								$needleType = $scope->getType($expr->expr->getArgs()[0]->value);
+								$dimFetchType = TypeCombinator::intersect($needleType, $arrayType->getIterableValueType());
+							} else {
+								$dimFetchType = $arrayType->getIterableValueType();
+							}
+
 							$specifiedTypes = $specifiedTypes->unionWith(
-								(new SpecifiedTypes([], []))->setNewConditionalExpressionHolders([
-									$dimFetchString => [$holder->getKey() => $holder],
-								]),
+								$this->create($dimFetch, $dimFetchType, TypeSpecifierContext::createTrue(), $scope),
 							);
+						} elseif ($expr->var instanceof Expr\Variable && is_string($expr->var->name)) {
+							$keyType = $scope->getType($expr->expr);
+							$narrowedKeyType = TypeCombinator::remove($keyType, $sentinelType);
+							if (!$narrowedKeyType instanceof NeverType) {
+								if ($isStrictArraySearch) {
+									$needleType = $scope->getType($expr->expr->getArgs()[0]->value);
+									$dimFetchType = TypeCombinator::intersect($needleType, $arrayType->getIterableValueType());
+								} else {
+									$dimFetchType = $arrayType->getIterableValueType();
+								}
+								$specifiedTypes = $specifiedTypes->unionWith(
+									$this->createArrayDimFetchConditionalExpressionHolder($expr->var, $arrayArg, $narrowedKeyType, $dimFetchType),
+								);
+							}
 						}
 					}
 				}
@@ -941,28 +984,6 @@ public function specifyTypesInCondition(
 				return $specifiedTypes;
 			}
 
-			if ($context->true()) {
-				// infer $arr[$key] after $key = array_search($needle, $arr)
-				if (
-					$expr->expr instanceof FuncCall
-					&& $expr->expr->name instanceof Name
-					&& !$expr->expr->isFirstClassCallable()
-					&& $expr->expr->name->toLowerString() === 'array_search'
-					&& count($expr->expr->getArgs()) >= 2
-				) {
-					$arrayArg = $expr->expr->getArgs()[1]->value;
-					$arrayType = $scope->getType($arrayArg);
-
-					if ($arrayType->isArray()->yes()) {
-						$dimFetch = new ArrayDimFetch($arrayArg, $expr->var);
-						$iterableValueType = $arrayType->getIterableValueType();
-
-						return $specifiedTypes->unionWith(
-							$this->create($dimFetch, $iterableValueType, TypeSpecifierContext::createTrue(), $scope),
-						);
-					}
-				}
-			}
 			return $specifiedTypes;
 		} elseif (
 			$expr instanceof Expr\Isset_
@@ -2404,6 +2425,27 @@ public function create(
 		return $types;
 	}
 
+	private function createArrayDimFetchConditionalExpressionHolder(
+		Expr\Variable $keyVar,
+		Expr $arrayArg,
+		Type $narrowedKeyType,
+		Type $dimFetchType,
+	): SpecifiedTypes
+	{
+		$dimFetch = new ArrayDimFetch($arrayArg, $keyVar);
+		$dimFetchString = $this->exprPrinter->printExpr($dimFetch);
+		$keyExprString = $this->exprPrinter->printExpr($keyVar);
+
+		$holder = new ConditionalExpressionHolder(
+			[$keyExprString => ExpressionTypeHolder::createYes($keyVar, $narrowedKeyType)],
+			ExpressionTypeHolder::createYes($dimFetch, $dimFetchType),
+		);
+
+		return (new SpecifiedTypes([], []))->setNewConditionalExpressionHolders([
+			$dimFetchString => [$holder->getKey() => $holder],
+		]);
+	}
+
 	private function createForExpr(
 		Expr $expr,
 		Type $type,
@@ -3032,11 +3074,12 @@ private function resolveNormalizedIdentical(Expr\BinaryOp\Identical $expr, Scope
 
 		// array_key_first($a) !== null
 		// array_key_last($a) !== null
+		// array_find_key($a, $cb) !== null
 		if (
 			$unwrappedLeftExpr instanceof FuncCall
 			&& $unwrappedLeftExpr->name instanceof Name
 			&& !$unwrappedLeftExpr->isFirstClassCallable()
-			&& in_array($unwrappedLeftExpr->name->toLowerString(), ['array_key_first', 'array_key_last'], true)
+			&& in_array($unwrappedLeftExpr->name->toLowerString(), ['array_key_first', 'array_key_last', 'array_find_key'], true)
 			&& isset($unwrappedLeftExpr->getArgs()[0])
 			&& $rightType->isNull()->yes()
 		) {

tests/PHPStan/Analyser/nsrt/array-find-key-existing.php

@@ -0,0 +1,49 @@
+<?php // lint >= 8.4
+
+declare(strict_types=1);
+
+namespace ArrayFindKeyExisting;
+
+use function PHPStan\Testing\assertType;
+
+/**
+ * @param list<string> $list
+ */
+function arrayFindKeyNotNull(array $list, string $s): void
+{
+	$key = array_find_key($list, fn (string $v) => $v === $s);
+	if ($key !== null) {
+		assertType('non-empty-list<string>', $list);
+		assertType('int<0, max>', $key);
+		assertType('string', $list[$key]);
+	} else {
+		assertType('array{}', $list);
+		assertType('null', $key);
+		assertType('*ERROR*', $list[$key]);
+	}
+	assertType('list<string>', $list);
+	assertType('int<0, max>|null', $key);
+	assertType('string', $list[$key]);
+}
+
+/**
+ * @param array<string, int> $map
+ */
+function arrayFindKeyStringKey(array $map): void
+{
+	$key = array_find_key($map, fn (int $v) => $v > 10);
+	if ($key !== null) {
+		assertType('int', $map[$key]);
+	}
+}
+
+/**
+ * @param list<string> $list
+ */
+function arrayFindKeyReversedComparison(array $list, string $s): void
+{
+	$key = array_find_key($list, fn (string $v) => $v === $s);
+	if (null !== $key) {
+		assertType('string', $list[$key]);
+	}
+}

tests/PHPStan/Analyser/nsrt/array-search-existing.php

@@ -0,0 +1,81 @@
+<?php declare(strict_types=1);
+
+namespace ArraySearchExisting;
+
+use function PHPStan\Testing\assertType;
+
+/**
+ * @param list<string> $list
+ */
+function arraySearchNotFalse(array $list, string $s): void
+{
+	$key = array_search($s, $list);
+	if ($key !== false) {
+		assertType('non-empty-list<string>', $list);
+		assertType('string', $list[$key]);
+	}
+}
+
+/**
+ * @param array<string, int> $map
+ */
+function arraySearchStringKey(array $map, int $needle): void
+{
+	$key = array_search($needle, $map);
+	if ($key !== false) {
+		assertType('int', $map[$key]);
+	}
+}
+
+/**
+ * @param list<string> $list
+ */
+function arraySearchReversedComparison(array $list, string $s): void
+{
+	$key = array_search($s, $list);
+	if (false !== $key) {
+		assertType('string', $list[$key]);
+	}
+}
+
+/**
+ * @param array<string, int|string> $arr
+ */
+function arraySearchStrictNarrowsToNeedle(array $arr, int $needle): void
+{
+	$key = array_search($needle, $arr, true);
+	if ($key !== false) {
+		assertType('non-empty-array<string, int|string>', $arr);
+		assertType('string', $key);
+		assertType('int', $arr[$key]);
+	} else {
+		assertType('array<string, int|string>', $arr);
+		assertType('false', $key);
+		assertType('*ERROR*', $arr[$key]);
+	}
+	assertType('array<string, int|string>', $arr);
+	assertType('string|false', $key);
+	assertType('int|string', $arr[$key]);
+
+}
+
+/**
+ * @param array<string, int|string> $arr
+ */
+function arraySearchLooseKeepsValueType(array $arr, int $needle): void
+{
+	$key = array_search($needle, $arr);
+	if ($key !== false) {
+		assertType('int|string', $arr[$key]);
+	}
+}
+
+/**
+ * @param array<string, int|string> $arr
+ */
+function arraySearchStrictInlineAssign(array $arr, int $needle): void
+{
+	if (($key = array_search($needle, $arr, true)) !== false) {
+		assertType('int', $arr[$key]);
+	}
+}

tests/PHPStan/Rules/Arrays/NonexistentOffsetInArrayDimFetchRuleTest.php

@@ -873,12 +873,7 @@ public function testArrayDimFetchAfterArraySearch(): void
 	{
 		$this->reportPossiblyNonexistentGeneralArrayOffset = true;
 
-		$this->analyse([__DIR__ . '/data/array-dim-after-array-search.php'], [
-			[
-				'Offset int|string might not exist on non-empty-array.',
-				20,
-			],
-		]);
+		$this->analyse([__DIR__ . '/data/array-dim-after-array-search.php'], []);
 	}
 
 	public function testArrayDimFetchOnArrayKeyFirsOrLastOrCount(): void
@@ -1310,4 +1305,37 @@ public function testBug11218(): void
 		$this->analyse([__DIR__ . '/data/bug-11218.php'], []);
 	}
 
+	public function testArraySearchExisting(): void
+	{
+		$this->reportPossiblyNonexistentGeneralArrayOffset = true;
+
+		$this->analyse([__DIR__ . '/../../Analyser/nsrt/array-search-existing.php'], [
+			[
+				'Offset false does not exist on array<string, int|string>.',
+				54,
+			],
+			[
+				'Offset string|false might not exist on array<string, int|string>.',
+				58,
+			],
+		]);
+	}
+
+	#[RequiresPhp('>= 8.4.0')]
+	public function testArrayFindKeyExisting(): void
+	{
+		$this->reportPossiblyNonexistentGeneralArrayOffset = true;
+
+		$this->analyse([__DIR__ . '/../../Analyser/nsrt/array-find-key-existing.php'], [
+			[
+				'Offset null does not exist on array{}.',
+				22,
+			],
+			[
+				'Offset int<0, max>|null might not exist on list<string>.',
+				26,
+			],
+		]);
+	}
+
 }