Preserve constant array when setting a union of constant scalar keys

src/Type/Constant/ConstantArrayType.php

@@ -693,6 +693,25 @@ public function getOffsetValueType(Type $offsetType): Type
 
 	public function setOffsetValueType(?Type $offsetType, Type $valueType, bool $unionValues = true): Type
 	{
+		if ($offsetType !== null) {
+			if ($offsetType instanceof IntegerRangeType) {
+				$constantScalars = $offsetType->getFiniteTypes();
+			} else {
+				$constantScalars = $offsetType->getConstantScalarTypes();
+			}
+
+			$constantScalarsCount = count($constantScalars);
+			if ($constantScalarsCount > 1 && $constantScalarsCount < ConstantArrayTypeBuilder::ARRAY_COUNT_LIMIT) {
+				$arrays = [];
+				foreach ($constantScalars as $constantScalar) {
+					$builder = ConstantArrayTypeBuilder::createFromConstantArray($this);
+					$builder->setOffsetValueType($constantScalar, $valueType);
+					$arrays[] = $builder->getArray();
+				}
+
+				return TypeCombinator::union(...$arrays);
+			}
+		}
 		$builder = ConstantArrayTypeBuilder::createFromConstantArray($this);
 		$builder->setOffsetValueType($offsetType, $valueType);
 
@@ -701,6 +720,24 @@ public function setOffsetValueType(?Type $offsetType, Type $valueType, bool $uni
 
 	public function setExistingOffsetValueType(Type $offsetType, Type $valueType): Type
 	{
+		if ($offsetType instanceof IntegerRangeType) {
+			$constantScalars = $offsetType->getFiniteTypes();
+		} else {
+			$constantScalars = $offsetType->getConstantScalarTypes();
+		}
+
+		$constantScalarsCount = count($constantScalars);
+		if ($constantScalarsCount > 1 && $constantScalarsCount < ConstantArrayTypeBuilder::ARRAY_COUNT_LIMIT) {
+			$arrays = [];
+			foreach ($constantScalars as $constantScalar) {
+				$builder = ConstantArrayTypeBuilder::createFromConstantArray($this);
+				$builder->setOffsetValueType($constantScalar, $valueType);
+				$arrays[] = $builder->getArray();
+			}
+
+			return TypeCombinator::union(...$arrays);
+		}
+
 		$builder = ConstantArrayTypeBuilder::createFromConstantArray($this);
 		$builder->setOffsetValueType($offsetType, $valueType);
 

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

@@ -22,7 +22,7 @@ public function parse(string $glue): string
 			$seenGlues[$glue] = true;
 
 			assertType("'&'|'|'", $glue);
-			assertType("array{'|': bool, '&': bool}", $seenGlues);
+			assertType("array{'|': false, '&': true}|array{'|': true, '&': false}", $seenGlues);
 		} else {
 			assertType("''", $glue);
 		}

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

@@ -21,25 +21,25 @@ public function doFoo(int $i)
 		/** @var 0|1|2 $offset */
 		$offset = doFoo();
 		$c[$offset] = true;
-		assertType('array{bool, bool, bool}', $c);
+		assertType('array{false, false, true}|array{false, true, false}|array{true, false, false}', $c);
 
 		$d = [false, false, false];
 		/** @var int<0, 2> $offset2 */
 		$offset2 = doFoo();
 		$d[$offset2] = true;
-		assertType('array{bool, bool, bool}', $d);
+		assertType('array{false, false, true}|array{false, true, false}|array{true, false, false}', $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);
+		assertType('array{false, false, false, true}|array{false, false, true}|array{false, true, false}|array{true, false, false}', $e);
 
 		$f = [false, false, false];
 		/** @var 0|1 $offset4 */
 		$offset4 = doFoo();
 		$f[$offset4] = true;
-		assertType('array{bool, bool, false}', $f);
+		assertType('array{false, true, false}|array{true, false, false}', $f);
 	}
 
 	/**
@@ -50,7 +50,7 @@ public function doBar(int $offset): void
 	{
 		$a = [false, false, false];
 		$a[$offset] = true;
-		assertType('array{bool, bool, false}', $a);
+		assertType('array{false, true, false}|array{true, false, false}', $a);
 	}
 
 	/**
@@ -94,14 +94,14 @@ public function doBar5(int $offset): void
 	{
 		$a = [false, false, false];
 		$a[$offset] = true;
-		assertType('non-empty-array<int<0, 4>, bool>', $a);
+		assertType('array{0: false, 1: false, 2: false, 4: true}|array{false, false, false, true}|array{false, false, true}|array{false, true, false}|array{true, false, false}', $a);
 	}
 
 	public function doBar6(bool $offset): void
 	{
 		$a = [false, false, false];
 		$a[$offset] = true;
-		assertType('array{bool, bool, false}', $a);
+		assertType('array{false, true, false}|array{true, false, false}', $a);
 	}
 
 	/**

tests/PHPStan/Analyser/nsrt/set-constant-union-offset-on-constant-array.php

@@ -0,0 +1,30 @@
+<?php
+
+namespace SetConstantUnionOffsetOnConstantArray;
+
+use function PHPStan\Testing\assertType;
+
+class Foo
+{
+
+	/**
+	 * @param array{foo: int}  $a
+	 */
+	public function doFoo(array $a): void
+	{
+		$k = rand(0, 1) ? 'a' : 'b';
+		$a[$k] = 256;
+		assertType('array{foo: int, a: 256}|array{foo: int, b: 256}', $a);
+	}
+
+	/**
+	 * @param array{foo: int}  $a
+	 * @param int<1,5> $intRange
+	 */
+	public function doBar(array $a, $intRange): void
+	{
+		$a[$intRange] = 256;
+		assertType('array{foo: int, 1: 256}|array{foo: int, 2: 256}|array{foo: int, 3: 256}|array{foo: int, 4: 256}|array{foo: int, 5: 256}', $a);
+	}
+
+}