From 22b25e6e7fae8f58ba50704f87d057392fcd8297 Mon Sep 17 00:00:00 2001
From: schlndh <schlndh@users.noreply.github.com>
Date: Sat, 27 Apr 2024 13:58:07 +0200
Subject: [PATCH] preserve large arrays with same keys through union

---
 src/Type/TypeCombinator.php                   | 73 ++++++++++++++++++-
 .../Analyser/NodeScopeResolverTest.php        |  1 +
 .../data/preserve-large-constant-array.php    | 50 +++++++++++++
 3 files changed, 121 insertions(+), 3 deletions(-)
 create mode 100644 tests/PHPStan/Analyser/data/preserve-large-constant-array.php

diff --git a/src/Type/TypeCombinator.php b/src/Type/TypeCombinator.php
index b4c52cd4cf..eb8491fcd7 100644
--- a/src/Type/TypeCombinator.php
+++ b/src/Type/TypeCombinator.php
@@ -29,6 +29,7 @@
 use function array_merge;
 use function array_slice;
 use function array_splice;
+use function array_unique;
 use function array_values;
 use function count;
 use function get_class;
@@ -643,7 +644,7 @@ private static function processArrayAccessoryTypes(array $arrayTypes): array
 	}
 
 	/**
-	 * @param Type[] $arrayTypes
+	 * @param list<Type> $arrayTypes
 	 * @return Type[]
 	 */
 	private static function processArrayTypes(array $arrayTypes): array
@@ -669,9 +670,14 @@ private static function processArrayTypes(array $arrayTypes): array
 
 		/** @var int|float $nextConstantKeyTypeIndex */
 		$nextConstantKeyTypeIndex = 1;
+		$constantArraysMap = array_map(
+			static fn (Type $t) => $t->getConstantArrays(),
+			$arrayTypes,
+		);
 
-		foreach ($arrayTypes as $arrayType) {
-			$isConstantArray = $arrayType->isConstantArray()->yes();
+		foreach ($arrayTypes as $arrayIdx => $arrayType) {
+			$constantArrays = $constantArraysMap[$arrayIdx];
+			$isConstantArray = $constantArrays !== [];
 			if (!$isConstantArray || !$arrayType->isIterableAtLeastOnce()->no()) {
 				$filledArrays++;
 			}
@@ -708,6 +714,11 @@ private static function processArrayTypes(array $arrayTypes): array
 		}
 
 		if ($generalArrayOccurred && (!$overflowed || $filledArrays > 1)) {
+			// Do this only for arrays which would be generalized, because this breaks tagged unions.
+			$singleConstantArray = self::unionConstantArrayTypesWithSameKeys($constantArraysMap);
+			if ($singleConstantArray !== null) {
+				return [$singleConstantArray];
+			}
 			$scopes = [];
 			$useTemplateArray = true;
 			foreach ($arrayTypes as $arrayType) {
@@ -748,6 +759,62 @@ private static function processArrayTypes(array $arrayTypes): array
 		);
 	}
 
+	/**
+	 * @param non-empty-list<ConstantArrayType[]> $constantArraysMap
+	 */
+	private static function unionConstantArrayTypesWithSameKeys(array $constantArraysMap): ?ConstantArrayType
+	{
+		$singleConstantArrayList = [];
+		$constantArraySizes = [];
+
+		foreach ($constantArraysMap as $constantArrays) {
+			if (count($constantArrays) !== 1) {
+				return null;
+			}
+
+			$singleConstantArrayList[] = $constantArrays[0];
+			$constantArraySizes[] = count($constantArrays[0]->getKeyTypes());
+		}
+
+		if (count(array_unique($constantArraySizes)) !== 1) {
+			return null;
+		}
+
+		$finalKeyTypes = $singleConstantArrayList[0]->getKeyTypes();
+		$finalValueTypesMap = [];
+		for ($i = 0; $i < $constantArraySizes[0]; $i++) {
+			$finalValueTypesMap[$i] = [$singleConstantArrayList[0]->getValueTypes()[$i]];
+			for ($j = 1; $j < count($singleConstantArrayList); $j++) {
+				$otherArray = $singleConstantArrayList[$j];
+				if (! $finalKeyTypes[$i]->equals($otherArray->getKeyTypes()[$i])) {
+					return null;
+				}
+
+				$finalValueTypesMap[$i][] = $otherArray->getValueTypes()[$i];
+			}
+		}
+
+		$finalValueTypes = array_map(
+			static fn (array $types) => self::union(...$types),
+			$finalValueTypesMap,
+		);
+
+		$optionalKeys = array_unique(array_merge(
+			...array_map(
+				static fn (ConstantArrayType $t) => $t->getOptionalKeys(),
+				$singleConstantArrayList,
+			),
+		));
+
+		return new ConstantArrayType(
+			$finalKeyTypes,
+			$finalValueTypes,
+			$singleConstantArrayList[0]->getNextAutoIndexes(),
+			$optionalKeys,
+			$singleConstantArrayList[0]->isList(),
+		);
+	}
+
 	/**
 	 * @param Type[] $types
 	 * @return Type[]
diff --git a/tests/PHPStan/Analyser/NodeScopeResolverTest.php b/tests/PHPStan/Analyser/NodeScopeResolverTest.php
index d214ec5906..710b292cf9 100644
--- a/tests/PHPStan/Analyser/NodeScopeResolverTest.php
+++ b/tests/PHPStan/Analyser/NodeScopeResolverTest.php
@@ -1420,6 +1420,7 @@ public function dataFileAsserts(): iterable
 			yield from $this->gatherAssertTypes(__DIR__ . '/data/trigger-error-php7.php');
 		}
 
+		yield from $this->gatherAssertTypes(__DIR__ . '/data/preserve-large-constant-array.php');
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/impure-error-log.php');
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/falsy-isset.php');
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/falsey-coalesce.php');
diff --git a/tests/PHPStan/Analyser/data/preserve-large-constant-array.php b/tests/PHPStan/Analyser/data/preserve-large-constant-array.php
new file mode 100644
index 0000000000..c16ec3c894
--- /dev/null
+++ b/tests/PHPStan/Analyser/data/preserve-large-constant-array.php
@@ -0,0 +1,50 @@
+<?php declare(strict_types=1);
+
+namespace PreserveLargeConstantArray;
+
+use function PHPStan\Testing\assertType;
+
+/**
+ * @param array{1: string|null, 2: int, 3: bool, 4: string, 5: int, 6: bool, 7: string, 8: int, 9: bool, 10: string, 11: int, 12: bool, 13: string, 14: int, 15: bool, 16: string, 17: int, 18: bool, 19: string, 20: int, 21: bool, 22: string, 23: int, 24: bool, 25: string, 26: int, 27: bool, 28: string, 29: int, 30: bool, 31: string, 32: int, 33: bool, 34: string, 35: int, 36: bool, 37: string, 38: int, 39: bool, 40: string, 41: int, 42: bool, 43: string, 44: int, 45: bool, 46: string, 47: int, 48: bool, 49: string, 50: int, 51: bool, 52: string, 53: int, 54: bool, 55: string, 56: int, 57: bool, 58: string, 59: int, 60: bool, 61: string, 62: int, 63: bool, 64: float} $arr
+ */
+function simpleUnion(array $arr): void
+{
+	$val = $arr[1] !== null
+		? $arr[1]
+		: null;
+	assertType('array{1: string|null, 2: int, 3: bool, 4: string, 5: int, 6: bool, 7: string, 8: int, 9: bool, 10: string, 11: int, 12: bool, 13: string, 14: int, 15: bool, 16: string, 17: int, 18: bool, 19: string, 20: int, 21: bool, 22: string, 23: int, 24: bool, 25: string, 26: int, 27: bool, 28: string, 29: int, 30: bool, 31: string, 32: int, 33: bool, 34: string, 35: int, 36: bool, 37: string, 38: int, 39: bool, 40: string, 41: int, 42: bool, 43: string, 44: int, 45: bool, 46: string, 47: int, 48: bool, 49: string, 50: int, 51: bool, 52: string, 53: int, 54: bool, 55: string, 56: int, 57: bool, 58: string, 59: int, 60: bool, 61: string, 62: int, 63: bool, 64: float}', $arr);
+	echo 1;
+}
+
+/**
+ * @param array{1?: string, 2: int, 3: bool, 4: string, 5: int, 6: bool, 7: string, 8: int, 9: bool, 10: string, 11: int, 12: bool, 13: string, 14: int, 15: bool, 16: string, 17: int, 18: bool, 19: string, 20: int, 21: bool, 22: string, 23: int, 24: bool, 25: string, 26: int, 27: bool, 28: string, 29: int, 30: bool, 31: string, 32: int, 33: bool, 34: string, 35: int, 36: bool, 37: string, 38: int, 39: bool, 40: string, 41: int, 42: bool, 43: string, 44: int, 45: bool, 46: string, 47: int, 48: bool, 49: string, 50: int, 51: bool, 52: string, 53: int, 54: bool, 55: string, 56: int, 57: bool, 58: string, 59: int, 60: bool, 61: string, 62: int, 63: bool, 64: float} $arr
+ */
+function optionalKey(array $arr): void
+{
+	$val = isset($arr[1])
+		? $arr[1]
+		: null;
+	assertType('array{1?: string, 2: int, 3: bool, 4: string, 5: int, 6: bool, 7: string, 8: int, 9: bool, 10: string, 11: int, 12: bool, 13: string, 14: int, 15: bool, 16: string, 17: int, 18: bool, 19: string, 20: int, 21: bool, 22: string, 23: int, 24: bool, 25: string, 26: int, 27: bool, 28: string, 29: int, 30: bool, 31: string, 32: int, 33: bool, 34: string, 35: int, 36: bool, 37: string, 38: int, 39: bool, 40: string, 41: int, 42: bool, 43: string, 44: int, 45: bool, 46: string, 47: int, 48: bool, 49: string, 50: int, 51: bool, 52: string, 53: int, 54: bool, 55: string, 56: int, 57: bool, 58: string, 59: int, 60: bool, 61: string, 62: int, 63: bool, 64: float}', $arr);
+	echo 1;
+}
+
+/**
+ * @param array{1: string, 2: int, 3: bool, 4: string, 5: int, 6: bool, 7: string, 8: int, 9: bool, 10: string, 11: int, 12: bool, 13: string, 14: int, 15: bool, 16: string, 17: int, 18: bool, 19: string, 20: int, 21: bool, 22: string, 23: int, 24: bool, 25: string, 26: int, 27: bool, 28: string, 29: int, 30: bool, 31: string, 32: int, 33: bool, 34: string, 35: int, 36: bool, 37: string, 38: int, 39: bool, 40: string, 41: int, 42: bool, 43: string, 44: int, 45: bool, 46: string, 47: int, 48: bool, 49: string, 50: int, 51: bool, 52: string, 53: int, 54: bool, 55: string, 56: int, 57: bool, 58: string, 59: int, 60: bool, 61: string, 62: int, 63: bool, 64: float} $arr
+ */
+function multipleOptions(array $arr): void
+{
+	if ($arr[1] === 'a') {
+		$brr = $arr;
+		$brr[1] = 'b';
+	} elseif ($arr[1] === 'b') {
+		$brr = $arr;
+		$brr[1] = 'c';
+	} elseif ($arr[1] === 'c') {
+		$brr = $arr;
+		$brr[1] = 'd';
+	} else {
+		$brr = $arr;
+	}
+	assertType('array{1: string, 2: int, 3: bool, 4: string, 5: int, 6: bool, 7: string, 8: int, 9: bool, 10: string, 11: int, 12: bool, 13: string, 14: int, 15: bool, 16: string, 17: int, 18: bool, 19: string, 20: int, 21: bool, 22: string, 23: int, 24: bool, 25: string, 26: int, 27: bool, 28: string, 29: int, 30: bool, 31: string, 32: int, 33: bool, 34: string, 35: int, 36: bool, 37: string, 38: int, 39: bool, 40: string, 41: int, 42: bool, 43: string, 44: int, 45: bool, 46: string, 47: int, 48: bool, 49: string, 50: int, 51: bool, 52: string, 53: int, 54: bool, 55: string, 56: int, 57: bool, 58: string, 59: int, 60: bool, 61: string, 62: int, 63: bool, 64: float}', $brr);
+	echo 1;
+}