From 2a4fd1932850775b0a4be39f940fa8ed4c61c612 Mon Sep 17 00:00:00 2001
From: Martin Herndl <martin@herndl.org>
Date: Mon, 24 Jan 2022 23:06:19 +0100
Subject: [PATCH] Support all array_filter flags for narrowing down types

---
 ...rFunctionReturnTypeReturnTypeExtension.php | 130 ++++++++++++------
 .../data/array-filter-arrow-functions.php     |  22 +--
 .../Analyser/data/array-filter-callables.php  |  24 ++--
 tests/PHPStan/Analyser/data/array-filter.php  |   4 +-
 4 files changed, 116 insertions(+), 64 deletions(-)

diff --git a/src/Type/Php/ArrayFilterFunctionReturnTypeReturnTypeExtension.php b/src/Type/Php/ArrayFilterFunctionReturnTypeReturnTypeExtension.php
index c0ba83e062..a8fa5864c7 100644
--- a/src/Type/Php/ArrayFilterFunctionReturnTypeReturnTypeExtension.php
+++ b/src/Type/Php/ArrayFilterFunctionReturnTypeReturnTypeExtension.php
@@ -2,8 +2,11 @@
 
 namespace PHPStan\Type\Php;
 
+use PhpParser\Node\Expr;
 use PhpParser\Node\Expr\ArrowFunction;
 use PhpParser\Node\Expr\Closure;
+use PhpParser\Node\Expr\ConstFetch;
+use PhpParser\Node\Expr\Error;
 use PhpParser\Node\Expr\FuncCall;
 use PhpParser\Node\Expr\Variable;
 use PhpParser\Node\Stmt\Return_;
@@ -26,6 +29,7 @@
 use function array_map;
 use function count;
 use function is_string;
+use function strtolower;
 
 class ArrayFilterFunctionReturnTypeReturnTypeExtension implements DynamicFunctionReturnTypeExtension
 {
@@ -41,57 +45,62 @@ public function getTypeFromFunctionCall(FunctionReflection $functionReflection,
 		$callbackArg = $functionCall->getArgs()[1]->value ?? null;
 		$flagArg = $functionCall->getArgs()[2]->value ?? null;
 
-		if ($arrayArg !== null) {
-			$arrayArgType = $scope->getType($arrayArg);
-			$keyType = $arrayArgType->getIterableKeyType();
-			$itemType = $arrayArgType->getIterableValueType();
+		if ($arrayArg === null) {
+			return new ArrayType(new MixedType(), new MixedType());
+		}
 
-			if ($arrayArgType instanceof MixedType) {
-				return new BenevolentUnionType([
-					new ArrayType(new MixedType(), new MixedType()),
-					new NullType(),
-				]);
-			}
+		$arrayArgType = $scope->getType($arrayArg);
+		$keyType = $arrayArgType->getIterableKeyType();
+		$itemType = $arrayArgType->getIterableValueType();
+
+		if ($arrayArgType instanceof MixedType) {
+			return new BenevolentUnionType([
+				new ArrayType(new MixedType(), new MixedType()),
+				new NullType(),
+			]);
+		}
 
-			if ($callbackArg === null) {
-				return TypeCombinator::union(
-					...array_map([$this, 'removeFalsey'], TypeUtils::getArrays($arrayArgType)),
-				);
+		if ($callbackArg === null || ($callbackArg instanceof ConstFetch && strtolower($callbackArg->name->parts[0]) === 'null')) {
+			return TypeCombinator::union(
+				...array_map([$this, 'removeFalsey'], TypeUtils::getArrays($arrayArgType)),
+			);
+		}
+
+		if ($flagArg === null) {
+			if ($callbackArg instanceof Closure && count($callbackArg->stmts) === 1 && count($callbackArg->params) > 0) {
+				$statement = $callbackArg->stmts[0];
+				if ($statement instanceof Return_ && $statement->expr !== null) {
+					[$itemType, $keyType] = $this->filterByTruthyValue($scope, $callbackArg->params[0]->var, $itemType, null, $keyType, $statement->expr);
+				}
+			} elseif ($callbackArg instanceof ArrowFunction && count($callbackArg->params) > 0) {
+				[$itemType, $keyType] = $this->filterByTruthyValue($scope, $callbackArg->params[0]->var, $itemType, null, $keyType, $callbackArg->expr);
 			}
+		}
 
-			if ($flagArg === null) {
-				$var = null;
-				$expr = null;
-				if ($callbackArg instanceof Closure && count($callbackArg->stmts) === 1 && count($callbackArg->params) > 0) {
-					$statement = $callbackArg->stmts[0];
-					if ($statement instanceof Return_ && $statement->expr !== null) {
-						$var = $callbackArg->params[0]->var;
-						$expr = $statement->expr;
-					}
-				} elseif ($callbackArg instanceof ArrowFunction && count($callbackArg->params) > 0) {
-					$var = $callbackArg->params[0]->var;
-					$expr = $callbackArg->expr;
+		if ($flagArg instanceof ConstFetch && $flagArg->name->parts[0] === 'ARRAY_FILTER_USE_KEY') {
+			if ($callbackArg instanceof Closure && count($callbackArg->stmts) === 1 && count($callbackArg->params) > 0) {
+				$statement = $callbackArg->stmts[0];
+				if ($statement instanceof Return_ && $statement->expr !== null) {
+					[$itemType, $keyType] = $this->filterByTruthyValue($scope, null, $itemType, $callbackArg->params[0]->var, $keyType, $statement->expr);
 				}
-				if ($var !== null && $expr !== null) {
-					if (!$var instanceof Variable || !is_string($var->name)) {
-						throw new ShouldNotHappenException();
-					}
-					$itemVariableName = $var->name;
-					if (!$scope instanceof MutatingScope) {
-						throw new ShouldNotHappenException();
-					}
-					$scope = $scope->assignVariable($itemVariableName, $itemType);
-					$scope = $scope->filterByTruthyValue($expr);
-					$itemType = $scope->getVariableType($itemVariableName);
-					if ($itemType instanceof NeverType) {
-						return new ConstantArrayType([], []);
-					}
+			} elseif ($callbackArg instanceof ArrowFunction && count($callbackArg->params) > 0) {
+				[$itemType, $keyType] = $this->filterByTruthyValue($scope, null, $itemType, $callbackArg->params[0]->var, $keyType, $callbackArg->expr);
+			}
+		}
+
+		if ($flagArg instanceof ConstFetch && $flagArg->name->parts[0] === 'ARRAY_FILTER_USE_BOTH') {
+			if ($callbackArg instanceof Closure && count($callbackArg->stmts) === 1 && count($callbackArg->params) > 0) {
+				$statement = $callbackArg->stmts[0];
+				if ($statement instanceof Return_ && $statement->expr !== null) {
+					[$itemType, $keyType] = $this->filterByTruthyValue($scope, $callbackArg->params[0]->var, $itemType, $callbackArg->params[1]->var ?? null, $keyType, $statement->expr);
 				}
+			} elseif ($callbackArg instanceof ArrowFunction && count($callbackArg->params) > 0) {
+				[$itemType, $keyType] = $this->filterByTruthyValue($scope, $callbackArg->params[0]->var, $itemType, $callbackArg->params[1]->var ?? null, $keyType, $callbackArg->expr);
 			}
+		}
 
-		} else {
-			$keyType = new MixedType();
-			$itemType = new MixedType();
+		if ($itemType instanceof NeverType || $keyType instanceof NeverType) {
+			return new ConstantArrayType([], []);
 		}
 
 		return new ArrayType($keyType, $itemType);
@@ -132,4 +141,39 @@ public function removeFalsey(Type $type): Type
 		return new ArrayType($keyType, $valueType);
 	}
 
+	/**
+	 * @return array{Type, Type}
+	 */
+	private function filterByTruthyValue(Scope $scope, Error|Variable|null $itemVar, Type $itemType, Error|Variable|null $keyVar, Type $keyType, Expr $expr): array
+	{
+		if (!$scope instanceof MutatingScope) {
+			throw new ShouldNotHappenException();
+		}
+
+		$itemVarName = null;
+		if ($itemVar !== null) {
+			if (!$itemVar instanceof Variable || !is_string($itemVar->name)) {
+				throw new ShouldNotHappenException();
+			}
+			$itemVarName = $itemVar->name;
+			$scope = $scope->assignVariable($itemVarName, $itemType);
+		}
+
+		$keyVarName = null;
+		if ($keyVar !== null) {
+			if (!$keyVar instanceof Variable || !is_string($keyVar->name)) {
+				throw new ShouldNotHappenException();
+			}
+			$keyVarName = $keyVar->name;
+			$scope = $scope->assignVariable($keyVarName, $keyType);
+		}
+
+		$scope = $scope->filterByTruthyValue($expr);
+
+		return [
+			$itemVarName !== null ? $scope->getVariableType($itemVarName) : $itemType,
+			$keyVarName !== null ? $scope->getVariableType($keyVarName) : $keyType,
+		];
+	}
+
 }
diff --git a/tests/PHPStan/Analyser/data/array-filter-arrow-functions.php b/tests/PHPStan/Analyser/data/array-filter-arrow-functions.php
index f49e1fe3a0..0c39643fc9 100644
--- a/tests/PHPStan/Analyser/data/array-filter-arrow-functions.php
+++ b/tests/PHPStan/Analyser/data/array-filter-arrow-functions.php
@@ -5,35 +5,39 @@
 use function PHPStan\Testing\assertType;
 
 /**
- * @param int[] $list1
- * @param int[] $list2
- * @param int[] $list3
+ * @param array<int, int> $list1
+ * @param array<int, int> $list2
+ * @param array<int, int> $list3
  */
 function alwaysEvaluatesToFalse(array $list1, array $list2, array $list3): void
 {
 	$filtered1 = array_filter($list1, static fn($item): bool => is_string($item));
 	assertType('array{}', $filtered1);
 
-	$filtered2 = array_filter($list2, static fn($item): bool => is_string($item), ARRAY_FILTER_USE_KEY);
-	assertType('array<int>', $filtered2); // not supported yet
+	$filtered2 = array_filter($list2, static fn($key): bool => is_string($key), ARRAY_FILTER_USE_KEY);
+	assertType('array{}', $filtered2);
 
 	$filtered3 = array_filter($list3, static fn($item, $key): bool => is_string($item) && is_string($key), ARRAY_FILTER_USE_BOTH);
-	assertType('array<int>', $filtered3); // not supported yet
+	assertType('array{}', $filtered3);
 }
 
 /**
  * @param array<int|string, int|string> $map1
  * @param array<int|string, int|string> $map2
  * @param array<int|string, int|string> $map3
+ * @param array<int|string, int|string> $map4
  */
 function filtersString(array $map1, array $map2, array $map3, array $map4): void
 {
 	$filtered1 = array_filter($map1, static fn($item): bool => is_string($item));
 	assertType('array<int|string, string>', $filtered1);
 
-	$filtered2 = array_filter($map2, static fn($item): bool => is_string($item), ARRAY_FILTER_USE_KEY);
-	assertType('array<int|string, int|string>', $filtered2); // not supported yet
+	$filtered2 = array_filter($map2, static fn($key): bool => is_string($key), ARRAY_FILTER_USE_KEY);
+	assertType('array<string, int|string>', $filtered2);
 
 	$filtered3 = array_filter($map3, static fn($item, $key): bool => is_string($item) && is_string($key), ARRAY_FILTER_USE_BOTH);
-	assertType('array<int|string, int|string>', $filtered3); // not supported yet
+	assertType('array<string, string>', $filtered3);
+
+	$filtered4 = array_filter($map4, static fn($item): bool => is_string($item), ARRAY_FILTER_USE_BOTH);
+	assertType('array<int|string, string>', $filtered4);
 }
diff --git a/tests/PHPStan/Analyser/data/array-filter-callables.php b/tests/PHPStan/Analyser/data/array-filter-callables.php
index 2f32fbc528..6855d1fe92 100644
--- a/tests/PHPStan/Analyser/data/array-filter-callables.php
+++ b/tests/PHPStan/Analyser/data/array-filter-callables.php
@@ -5,35 +5,39 @@
 use function PHPStan\Testing\assertType;
 
 /**
- * @param int[] $list1
- * @param int[] $list2
- * @param int[] $list3
+ * @param array<int, int> $list1
+ * @param array<int, int> $list2
+ * @param array<int, int> $list3
  */
 function alwaysEvaluatesToFalse(array $list1, array $list2, array $list3): void
 {
 	$filtered1 = array_filter($list1, static function ($item): bool { return is_string($item); });
 	assertType('array{}', $filtered1);
 
-	$filtered2 = array_filter($list2, static function ($item): bool { return is_string($item); }, ARRAY_FILTER_USE_KEY);
-	assertType('array<int>', $filtered2); // not supported yet
+	$filtered2 = array_filter($list2, static function ($key): bool { return is_string($key); }, ARRAY_FILTER_USE_KEY);
+	assertType('array{}', $filtered2);
 
 	$filtered3 = array_filter($list3, static function ($item, $key): bool { return is_string($item) && is_string($key); }, ARRAY_FILTER_USE_BOTH);
-	assertType('array<int>', $filtered3); // not supported yet
+	assertType('array{}', $filtered3);
 }
 
 /**
  * @param array<int|string, int|string> $map1
  * @param array<int|string, int|string> $map2
  * @param array<int|string, int|string> $map3
+ * @param array<int|string, int|string> $map4
  */
-function filtersString(array $map1, array $map2, array $map3): void
+function filtersString(array $map1, array $map2, array $map3, array $map4): void
 {
 	$filtered1 = array_filter($map1, static function ($item): bool { return is_string($item); });
 	assertType('array<int|string, string>', $filtered1);
 
-	$filtered2 = array_filter($map2, static function ($item): bool { return is_string($item); }, ARRAY_FILTER_USE_KEY);
-	assertType('array<int|string, int|string>', $filtered2); // not supported yet
+	$filtered2 = array_filter($map2, static function ($key): bool { return is_string($key); }, ARRAY_FILTER_USE_KEY);
+	assertType('array<string, int|string>', $filtered2);
 
 	$filtered3 = array_filter($map3, static function ($item, $key): bool { return is_string($item) && is_string($key); }, ARRAY_FILTER_USE_BOTH);
-	assertType('array<int|string, int|string>', $filtered3); // not supported yet
+	assertType('array<string, string>', $filtered3);
+
+	$filtered4 = array_filter($map4, static function ($item): bool { return is_string($item); }, ARRAY_FILTER_USE_BOTH);
+	assertType('array<int|string, string>', $filtered4);
 }
diff --git a/tests/PHPStan/Analyser/data/array-filter.php b/tests/PHPStan/Analyser/data/array-filter.php
index 06ea76f861..f5d1e181fe 100644
--- a/tests/PHPStan/Analyser/data/array-filter.php
+++ b/tests/PHPStan/Analyser/data/array-filter.php
@@ -30,8 +30,8 @@ function withoutCallback(array $map1, array $map2, array $map3): void
 	assertType('array<string, float|int<min, -1>|int<1, max>|non-empty-string|true>', $filtered1);
 
 	$filtered2 = array_filter($map2, null, ARRAY_FILTER_USE_KEY);
-	assertType('array<string, bool|float|int|string>', $filtered2); // not supported yet
+	assertType('array<string, float|int<min, -1>|int<1, max>|non-empty-string|true>', $filtered2);
 
 	$filtered3 = array_filter($map3, null, ARRAY_FILTER_USE_BOTH);
-	assertType('array<string, bool|float|int|string>', $filtered3); // not supported yet
+	assertType('array<string, float|int<min, -1>|int<1, max>|non-empty-string|true>', $filtered3);
 }