Improve preserving ConstantArrayType after setting new offset

src/Type/Constant/ConstantArrayTypeBuilder.php

@@ -2,10 +2,12 @@
 
 namespace PHPStan\Type\Constant;
 
+use PHPStan\ShouldNotHappenException;
 use PHPStan\Type\Accessory\NonEmptyArrayType;
 use PHPStan\Type\ArrayType;
 use PHPStan\Type\Type;
 use PHPStan\Type\TypeCombinator;
+use PHPStan\Type\TypeUtils;
 use function array_filter;
 use function array_values;
 use function count;
@@ -87,6 +89,40 @@ public function setOffsetValueType(?Type $offsetType, Type $valueType, bool $opt
 			return;
 		}
 
+		$scalarTypes = TypeUtils::getConstantScalars($offsetType);
+		if (!$this->degradeToGeneralArray && count($scalarTypes) > 0) {
+			$match = true;
+			$valueTypes = $this->valueTypes;
+			foreach ($scalarTypes as $scalarType) {
+				$scalarOffsetType = ArrayType::castToArrayKeyType($scalarType);
+				if (!$scalarOffsetType instanceof ConstantIntegerType && !$scalarOffsetType instanceof ConstantStringType) {
+					throw new ShouldNotHappenException();
+				}
+				$offsetMatch = false;
+
+				/** @var ConstantIntegerType|ConstantStringType $keyType */
+				foreach ($this->keyTypes as $i => $keyType) {
+					if ($keyType->getValue() !== $scalarOffsetType->getValue()) {
+						continue;
+					}
+
+					$valueTypes[$i] = TypeCombinator::union($valueTypes[$i], $valueType);
+					$offsetMatch = true;
+				}
+
+				if ($offsetMatch) {
+					continue;
+				}
+
+				$match = false;
+			}
+
+			if ($match) {
+				$this->valueTypes = $valueTypes;
+				return;
+			}
+		}
+
 		$this->keyTypes[] = $offsetType;
 		$this->valueTypes[] = $valueType;
 		if ($optional) {

tests/PHPStan/Analyser/LegacyNodeScopeResolverTest.php

@@ -8107,7 +8107,7 @@ public function dataArrayKeysInBranches(): array
 				'$arrayAppendedInForeach',
 			],
 			[
-				'array<int, \'bar\'|\'baz\'|\'foo\'>',
+				'array<int, literal-string&non-empty-string>', // could be 'array<int, \'bar\'|\'baz\'|\'foo\'>'
 				'$anotherArrayAppendedInForeach',
 			],
 			[

tests/PHPStan/Analyser/NodeScopeResolverTest.php

@@ -595,6 +595,7 @@ public function dataFileAsserts(): iterable
 
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/array-destructuring-types.php');
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/pdo-prepare.php');
+		yield from $this->gatherAssertTypes(__DIR__ . '/data/constant-array-type-set.php');
 	}
 
 	/**

tests/PHPStan/Analyser/data/constant-array-type-set.php

@@ -0,0 +1,39 @@
+<?php
+
+namespace ConstantArrayTypeSet;
+
+use function PHPStan\Testing\assertType;
+
+class Foo
+{
+
+	public function doFoo(int $i)
+	{
+		$a = [1, 2, 3];
+		$a[$i] = 4;
+		assertType('non-empty-array<int, 1|2|3|4>', $a);
+
+		$b = [1, 2, 3];
+		$b[3] = 4;
+		assertType('array{1, 2, 3, 4}', $b);
+
+		$c = [false, false, false];
+		/** @var 0|1|2 $offset */
+		$offset = doFoo();
+		$c[$offset] = true;
+		assertType('array{bool, bool, bool}', $c);
+
+		$d = [false, false, false];
+		/** @var int<0, 2> $offset2 */
+		$offset2 = doFoo();
+		$d[$offset2] = true;
+		//assertType('array{bool, bool, bool}', $d);
+
+		$e = [false, false, false];
+		/** @var 0|1|2|3 $offset3 */
+		$offset3 = doFoo();
+		$e[$offset3] = true;
+		assertType('non-empty-array<0|1|2|3, bool>', $e);
+	}
+
+}