Narrow array key type after type-checking the key variable inside a foreach loop

src/Analyser/NodeScopeResolver.php

@@ -1315,21 +1315,30 @@ public function processStmtNode(
 			) {
 				$arrayExprDimFetch = new ArrayDimFetch($stmt->expr, $stmt->keyVar);
 				$arrayDimFetchLoopTypes = [];
+				$keyLoopTypes = [];
 				foreach ($scopesWithIterableValueType as $scopeWithIterableValueType) {
 					$arrayDimFetchLoopTypes[] = $scopeWithIterableValueType->getType($arrayExprDimFetch);
+					$keyLoopTypes[] = $scopeWithIterableValueType->getType($stmt->keyVar);
 				}
 
 				$arrayDimFetchLoopType = TypeCombinator::union(...$arrayDimFetchLoopTypes);
+				$keyLoopType = TypeCombinator::union(...$keyLoopTypes);
 
 				$arrayDimFetchLoopNativeTypes = [];
+				$keyLoopNativeTypes = [];
 				foreach ($scopesWithIterableValueType as $scopeWithIterableValueType) {
 					$arrayDimFetchLoopNativeTypes[] = $scopeWithIterableValueType->getNativeType($arrayExprDimFetch);
+					$keyLoopNativeTypes[] = $scopeWithIterableValueType->getNativeType($stmt->keyVar);
 				}
 
 				$arrayDimFetchLoopNativeType = TypeCombinator::union(...$arrayDimFetchLoopNativeTypes);
+				$keyLoopNativeType = TypeCombinator::union(...$keyLoopNativeTypes);
 
-				if (!$arrayDimFetchLoopType->equals($exprType->getIterableValueType())) {
-					$newExprType = TypeTraverser::map($exprType, static function (Type $type, callable $traverse) use ($arrayDimFetchLoopType): Type {
+				$valueTypeChanged = !$arrayDimFetchLoopType->equals($exprType->getIterableValueType());
+				$keyTypeChanged = !$keyLoopType->equals($exprType->getIterableKeyType());
+
+				if ($valueTypeChanged || $keyTypeChanged) {
+					$newExprType = TypeTraverser::map($exprType, static function (Type $type, callable $traverse) use ($arrayDimFetchLoopType, $keyLoopType, $valueTypeChanged, $keyTypeChanged): Type {
 						if ($type instanceof UnionType || $type instanceof IntersectionType) {
 							return $traverse($type);
 						}
@@ -1338,9 +1347,13 @@ public function processStmtNode(
 							return $type;
 						}
 
-						return new ArrayType($type->getKeyType(), $arrayDimFetchLoopType);
+						return new ArrayType(
+							$keyTypeChanged ? $keyLoopType : $type->getKeyType(),
+							$valueTypeChanged ? $arrayDimFetchLoopType : $type->getIterableValueType(),
+						);
 					});
-					$newExprNativeType = TypeTraverser::map($scope->getNativeType($stmt->expr), static function (Type $type, callable $traverse) use ($arrayDimFetchLoopNativeType): Type {
+					$nativeExprType = $scope->getNativeType($stmt->expr);
+					$newExprNativeType = TypeTraverser::map($nativeExprType, static function (Type $type, callable $traverse) use ($arrayDimFetchLoopNativeType, $keyLoopNativeType, $valueTypeChanged, $keyTypeChanged): Type {
 						if ($type instanceof UnionType || $type instanceof IntersectionType) {
 							return $traverse($type);
 						}
@@ -1349,7 +1362,10 @@ public function processStmtNode(
 							return $type;
 						}
 
-						return new ArrayType($type->getKeyType(), $arrayDimFetchLoopNativeType);
+						return new ArrayType(
+							$keyTypeChanged ? $keyLoopNativeType : $type->getKeyType(),
+							$valueTypeChanged ? $arrayDimFetchLoopNativeType : $type->getIterableValueType(),
+						);
 					});
 
 					if ($stmt->expr instanceof Variable && is_string($stmt->expr->name)) {

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

@@ -0,0 +1,202 @@
+<?php
+
+namespace Bug7076;
+
+use function PHPStan\Testing\assertType;
+
+/**
+ * @param array<int|string, mixed> $arguments
+ * @return array<string, mixed>
+ */
+function narrowWithIsString(array $arguments): array
+{
+	foreach ($arguments as $key => $argument) {
+		if (!is_string($key)) {
+			throw new \Exception('Key must be a string');
+		}
+	}
+
+	assertType('array<string, mixed>', $arguments);
+
+	return $arguments;
+}
+
+/**
+ * @param array<int|string, mixed> $arguments
+ * @return array<string, mixed>
+ */
+function narrowWithIsInt(array $arguments): array
+{
+	foreach ($arguments as $key => $argument) {
+		if (is_int($key)) {
+			throw new \Exception('Key must be a string');
+		}
+	}
+
+	assertType('array<string, mixed>', $arguments);
+
+	return $arguments;
+}
+
+/**
+ * @param array<int|string, mixed> $arguments
+ */
+function narrowToIntKeys(array $arguments): void
+{
+	foreach ($arguments as $key => $argument) {
+		if (!is_int($key)) {
+			throw new \Exception('Key must be an int');
+		}
+	}
+
+	assertType('array<int, mixed>', $arguments);
+}
+
+/**
+ * @param array<int|string, mixed> $arguments
+ */
+function narrowWithReturn(array $arguments): void
+{
+	foreach ($arguments as $key => $argument) {
+		if (!is_string($key)) {
+			return;
+		}
+	}
+
+	assertType('array<string, mixed>', $arguments);
+}
+
+/**
+ * @param array<int|string, mixed> $arguments
+ */
+function continueDoesNotNarrow(array $arguments): void
+{
+	foreach ($arguments as $key => $argument) {
+		if (!is_string($key)) {
+			continue;
+		}
+	}
+
+	assertType('array<int|string, mixed>', $arguments);
+}
+
+/**
+ * @param array<int|string, mixed> $arguments
+ */
+function breakPreventsNarrowing(array $arguments): void
+{
+	foreach ($arguments as $key => $argument) {
+		if (!is_string($key)) {
+			throw new \Exception();
+		}
+		if (rand(0, 1)) {
+			break;
+		}
+	}
+
+	assertType('array<int|string, mixed>', $arguments);
+}
+
+/**
+ * @param array<int|string, string|null> $arguments
+ */
+function keyAndValueNarrowing(array $arguments): void
+{
+	foreach ($arguments as $key => $argument) {
+		if (!is_string($key)) {
+			throw new \Exception();
+		}
+		$arguments[$key] = $argument ?? '';
+	}
+
+	assertType('array<string, string>', $arguments);
+}
+
+/**
+ * @param array<int|string, mixed> $arguments
+ */
+function noKeyVar(array $arguments): void
+{
+	foreach ($arguments as $argument) {
+		if (!is_string($argument)) {
+			throw new \Exception();
+		}
+	}
+
+	assertType('array<int|string, mixed>', $arguments);
+}
+
+/**
+ * @param array<int|string, mixed> $arguments
+ */
+function keyReassignedPreventsNarrowing(array $arguments): void
+{
+	foreach ($arguments as $key => $argument) {
+		$key = 'test';
+		if (!is_string($key)) {
+			throw new \Exception();
+		}
+	}
+
+	assertType('array<int|string, mixed>', $arguments);
+}
+
+/**
+ * @param array<int|string, mixed> $arguments
+ */
+function narrowWithAssert(array $arguments): void
+{
+	foreach ($arguments as $key => $argument) {
+		assert(is_string($key));
+	}
+
+	assertType('array<string, mixed>', $arguments);
+}
+
+/**
+ * @param non-empty-array<int|string, mixed> $arguments
+ */
+function narrowNonEmptyArray(array $arguments): void
+{
+	foreach ($arguments as $key => $argument) {
+		if (!is_string($key)) {
+			throw new \Exception();
+		}
+	}
+
+	assertType('non-empty-array<string, mixed>', $arguments);
+}
+
+class Foo
+{
+	/** @var array<int|string, mixed> */
+	private array $prop;
+
+	public function narrowPropertyKey(): void
+	{
+		foreach ($this->prop as $k => $v) {
+			if (!is_string($k)) {
+				throw new \Exception();
+			}
+		}
+
+		assertType('array<string, mixed>', $this->prop);
+	}
+}
+
+/**
+ * @param array<int|string, mixed> $arguments
+ */
+function partialContinueNarrowingDoesNotApply(array $arguments): void
+{
+	foreach ($arguments as $key => $argument) {
+		if (rand(0, 1)) {
+			continue;
+		}
+		if (!is_string($key)) {
+			throw new \Exception();
+		}
+	}
+
+	assertType('array<int|string, mixed>', $arguments);
+}