From 480626ecb52d2e98cc28cee8a18dfb86112b7f8f Mon Sep 17 00:00:00 2001
From: Ondrej Mirtes <ondrej@mirtes.cz>
Date: Wed, 4 Jan 2023 10:49:28 +0100
Subject: [PATCH] UnionType::pickTypes overriden in BenevolentUnionType for a
 more benevolent behaviour

---
 src/Type/BenevolentUnionType.php     | 13 ++++++
 src/Type/UnionType.php               | 66 +++++++++++-----------------
 tests/PHPStan/Type/UnionTypeTest.php |  9 ++++
 3 files changed, 47 insertions(+), 41 deletions(-)

diff --git a/src/Type/BenevolentUnionType.php b/src/Type/BenevolentUnionType.php
index 5a29676f29..a557981894 100644
--- a/src/Type/BenevolentUnionType.php
+++ b/src/Type/BenevolentUnionType.php
@@ -43,6 +43,19 @@ protected function unionTypes(callable $getType): Type
 		return TypeUtils::toBenevolentUnion(TypeCombinator::union(...$resultTypes));
 	}
 
+	protected function pickTypes(callable $getTypes): array
+	{
+		$types = [];
+		foreach ($this->getTypes() as $type) {
+			$innerTypes = $getTypes($type);
+			foreach ($innerTypes as $innerType) {
+				$types[] = $innerType;
+			}
+		}
+
+		return $types;
+	}
+
 	public function getOffsetValueType(Type $offsetType): Type
 	{
 		$types = [];
diff --git a/src/Type/UnionType.php b/src/Type/UnionType.php
index 9630c6bbc3..5e3913ee0e 100644
--- a/src/Type/UnionType.php
+++ b/src/Type/UnionType.php
@@ -114,55 +114,17 @@ public function getReferencedClasses(): array
 
 	public function getArrays(): array
 	{
-		$arrays = [];
-		foreach ($this->types as $type) {
-			$innerTypeArrays = $type->getArrays();
-			if ($innerTypeArrays === []) {
-				return [];
-			}
-
-			foreach ($innerTypeArrays as $array) {
-				$arrays[] = $array;
-			}
-		}
-
-		return $arrays;
+		return $this->pickTypes(static fn (Type $type) => $type->getArrays());
 	}
 
 	public function getConstantArrays(): array
 	{
-		$constantArrays = [];
-		foreach ($this->types as $type) {
-			$typeAsConstantArrays = $type->getConstantArrays();
-
-			if ($typeAsConstantArrays === []) {
-				return [];
-			}
-
-			foreach ($typeAsConstantArrays as $constantArray) {
-				$constantArrays[] = $constantArray;
-			}
-		}
-
-		return $constantArrays;
+		return $this->pickTypes(static fn (Type $type) => $type->getConstantArrays());
 	}
 
 	public function getConstantStrings(): array
 	{
-		$strings = [];
-		foreach ($this->types as $type) {
-			$constantStrings = $type->getConstantStrings();
-
-			if ($constantStrings === []) {
-				return [];
-			}
-
-			foreach ($constantStrings as $string) {
-				$strings[] = $string;
-			}
-		}
-
-		return $strings;
+		return $this->pickTypes(static fn (Type $type) => $type->getConstantStrings());
 	}
 
 	public function accepts(Type $type, bool $strictTypes): TrinaryLogic
@@ -929,4 +891,26 @@ protected function unionTypes(callable $getType): Type
 		return TypeCombinator::union(...array_map($getType, $this->types));
 	}
 
+	/**
+	 * @template T of Type
+	 * @param callable(Type $type): list<T> $getTypes
+	 * @return list<T>
+	 */
+	protected function pickTypes(callable $getTypes): array
+	{
+		$types = [];
+		foreach ($this->types as $type) {
+			$innerTypes = $getTypes($type);
+			if ($innerTypes === []) {
+				return [];
+			}
+
+			foreach ($innerTypes as $innerType) {
+				$types[] = $innerType;
+			}
+		}
+
+		return $types;
+	}
+
 }
diff --git a/tests/PHPStan/Type/UnionTypeTest.php b/tests/PHPStan/Type/UnionTypeTest.php
index 61d4768534..d3aab21a64 100644
--- a/tests/PHPStan/Type/UnionTypeTest.php
+++ b/tests/PHPStan/Type/UnionTypeTest.php
@@ -1407,6 +1407,15 @@ public function dataGetConstantStrings(): iterable
 					"'bar'",
 				],
 			],
+			[
+				new BenevolentUnionType([
+					new ConstantStringType('foo'),
+					new NullType(),
+				]),
+				[
+					"'foo'",
+				],
+			],
 		];
 	}