Track by-reference value modifications in constant array foreach unrolling

src/Analyser/MutatingScope.php

@@ -2359,7 +2359,7 @@
 		// corresponding array dim fetch without being confused by a reassignment
 		// ($type = 'foo' invalidates this expression, same as OriginalForeachKeyExpr).
 		$scope = $scope->assignExpression(new OriginalForeachValueExpr($valueName), $valueType, $nativeValueType);
-		if ($valueByRef && $iterateeType->isArray()->yes() && $iterateeType->isConstantArray()->no()) {
+		if ($valueByRef && $iterateeType->isArray()->yes()) {
 			$scope = $scope->assignExpression(
 				new IntertwinedVariableByReferenceWithExpr($valueName, $iteratee, new SetExistingOffsetValueTypeExpr(
 					$iteratee,
@@ -2373,7 +2373,7 @@
 		if ($keyName !== null) {
 			$scope = $scope->enterForeachKey($originalScope, $iteratee, $keyName);
 
-			if ($valueByRef && $iterateeType->isArray()->yes() && $iterateeType->isConstantArray()->no()) {
+			if ($valueByRef && $iterateeType->isArray()->yes()) {
 				$scope = $scope->assignExpression(
 					new IntertwinedVariableByReferenceWithExpr($valueName, new Expr\ArrayDimFetch($iteratee, new Variable($keyName)), new Variable($valueName)),
 					$valueType,

src/Analyser/NodeScopeResolver.php

@@ -3885,9 +3885,6 @@ private function tryProcessUnrolledConstantArrayForeach(
 		StatementContext $context,
 	): ?array
 	{
-		if ($stmt->byRef) {
-			return null;
-		}
 		if (!($stmt->valueVar instanceof Variable && is_string($stmt->valueVar->name))) {
 			return null;
 		}
@@ -3982,6 +3979,32 @@ private function tryProcessUnrolledConstantArrayForeach(
 				$breakScopes[] = $breakExitPoint->getScope();
 			}
 
+			if ($stmt->byRef) {
+				$newValueType = $iterEndScope->getType(new Variable($valueVarName));
+				$newNativeValueType = $iterEndScope->getNativeType(new Variable($valueVarName));
+
+				$currentArrayType = $iterEndScope->getType($stmt->expr);
+				$currentNativeArrayType = $iterEndScope->getNativeType($stmt->expr);
+
+				$newArrayType = $currentArrayType->setExistingOffsetValueType($keyType, $newValueType);
+				$newNativeArrayType = $currentNativeArrayType->setExistingOffsetValueType($nativeKeyType, $newNativeValueType);
+
+				if ($stmt->expr instanceof Variable && is_string($stmt->expr->name)) {
+					$iterEndScope = $iterEndScope->assignVariable(
+						$stmt->expr->name,
+						$newArrayType,
+						$newNativeArrayType,
+						TrinaryLogic::createYes(),
+					);
+				} else {
+					$iterEndScope = $iterEndScope->assignExpression(
+						$stmt->expr,
+						$newArrayType,
+						$newNativeArrayType,
+					);
+				}
+			}
+
 			if ($isOptional) {
 				$chainScope = $iterEndScope->mergeWith($chainScope);
 			} else {

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

@@ -0,0 +1,105 @@
+<?php declare(strict_types = 1);
+
+namespace Bug1311;
+
+use function PHPStan\Testing\assertType;
+
+class HelloWorld
+{
+	/**
+	 * @param array<int, array{a: string}> $sets
+	 *
+	 * @return array<int, array{a: string, b: bool}>
+	 */
+	public function sayHello(array $sets): array
+	{
+		foreach ($sets as &$set) {
+			$set['b'] = false;
+		}
+
+		assertType('array<int, array{a: string, b: false}>', $sets);
+
+		return $sets;
+	}
+}
+
+function foreachByRefConstantArray(): void
+{
+	$temp = [1, 2, 3];
+
+	foreach ($temp as &$item) {
+		$item = (string) $item;
+	}
+
+	assertType("array{'1', '2', '3'}", $temp);
+}
+
+function foreachByRefConstantArrayWithKey(): void
+{
+	$temp = [1, 2, 3];
+
+	foreach ($temp as $key => &$item) {
+		$item = (string) $item;
+	}
+
+	assertType("array{'1', '2', '3'}", $temp);
+}
+
+function foreachByRefShapedArray(): void
+{
+	/** @var array{a: int, b: int} $data */
+	$data = ['a' => 1, 'b' => 2];
+
+	foreach ($data as &$val) {
+		$val = (string) $val;
+	}
+
+	assertType("array{a: lowercase-string&numeric-string&uppercase-string, b: lowercase-string&numeric-string&uppercase-string}", $data);
+}
+
+function foreachByRefListPreservation(): void
+{
+	/** @var list<int> $list */
+	$list = [1, 2, 3];
+
+	foreach ($list as &$item) {
+		$item = $item * 2;
+	}
+
+	assertType("list<int>", $list);
+}
+
+function foreachByRefConditionalModification(): void
+{
+	$temp = [1, 2, 3];
+
+	foreach ($temp as &$item) {
+		if ($item > 1) {
+			$item = (string) $item;
+		}
+	}
+
+	assertType("array{1, '2', '3'}", $temp);
+}
+
+function foreachByRefAppend(): void
+{
+	$data = ['a' => 1, 'b' => 2];
+
+	foreach ($data as &$val) {
+		$val = [$val, 'extra'];
+	}
+
+	assertType("array{a: array{1, 'extra'}, b: array{2, 'extra'}}", $data);
+}
+
+function forLoop(): void
+{
+	$temp = [1, 2, 3];
+
+	for ($i = 0; $i < count($temp); $i++) {
+		$temp[$i] = (string) $temp[$i];
+	}
+
+	assertType("array{1|(literal-string&lowercase-string&non-falsy-string&numeric-string&uppercase-string), 2|(literal-string&lowercase-string&non-falsy-string&numeric-string&uppercase-string), 3|(literal-string&lowercase-string&non-falsy-string&numeric-string&uppercase-string)}", $temp);
+}