Preserve constant array when setting a union of constant scalar keys

Author: ondrejmirtes

src/Type/Constant/ConstantArrayType.php

@@ -693,6 +693,20 @@ public function getOffsetValueType(Type $offsetType): Type
 
 	public function setOffsetValueType(?Type $offsetType, Type $valueType, bool $unionValues = true): Type
 	{
+		if ($offsetType !== null) {
+			$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($this, ...$arrays);
+			}
+		}
 		$builder = ConstantArrayTypeBuilder::createFromConstantArray($this);
 		$builder->setOffsetValueType($offsetType, $valueType);
 
@@ -701,6 +715,19 @@ public function setOffsetValueType(?Type $offsetType, Type $valueType, bool $uni
 
 	public function setExistingOffsetValueType(Type $offsetType, Type $valueType): Type
 	{
+		$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

@@ -33,7 +33,7 @@ public function doFoo(int $i)
 		/** @var 0|1|2|3 $offset3 */
 		$offset3 = doFoo();
 		$e[$offset3] = true;
-		assertType('non-empty-array<0|1|2|3, bool>', $e);
+		assertType('array{0: bool, 1: bool, 2: bool, 3?: true}', $e);
 
 		$f = [false, false, false];
 		/** @var 0|1 $offset4 */

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

@@ -0,0 +1,20 @@
+<?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);
+	}
+
+}