From 12b6bf7834aa91b7cb78675c783699a462fa3c37 Mon Sep 17 00:00:00 2001
From: Ondrej Mirtes <ondrej@mirtes.cz>
Date: Wed, 12 Jan 2022 10:26:36 +0100
Subject: [PATCH] Improve for loop handling

---
 src/Analyser/NodeScopeResolver.php            | 27 +++++++++++++++++--
 .../PHPStan/Analyser/data/for-loop-i-type.php | 25 +++++++++++++++++
 2 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/src/Analyser/NodeScopeResolver.php b/src/Analyser/NodeScopeResolver.php
index 1165c81b72..13e6788e40 100644
--- a/src/Analyser/NodeScopeResolver.php
+++ b/src/Analyser/NodeScopeResolver.php
@@ -978,9 +978,17 @@ private function processStmtNode(
 			}
 
 			$bodyScope = $initScope;
+			$isIterableAtLeastOnce = TrinaryLogic::createYes();
 			foreach ($stmt->cond as $condExpr) {
 				$condResult = $this->processExprNode($condExpr, $bodyScope, static function (): void {
 				}, ExpressionContext::createDeep());
+				$condTruthiness = $condResult->getScope()->getType($condExpr)->toBoolean();
+				if ($condTruthiness instanceof ConstantBooleanType) {
+					$condTruthinessTrinary = TrinaryLogic::createFromBoolean($condTruthiness->getValue());
+				} else {
+					$condTruthinessTrinary = TrinaryLogic::createMaybe();
+				}
+				$isIterableAtLeastOnce = $isIterableAtLeastOnce->and($condTruthinessTrinary);
 				$hasYield = $hasYield || $condResult->hasYield();
 				$throwPoints = array_merge($throwPoints, $condResult->getThrowPoints());
 				$bodyScope = $condResult->getTruthyScope();
@@ -1043,8 +1051,23 @@ private function processStmtNode(
 				$finalScope = $breakExitPoint->getScope()->mergeWith($finalScope);
 			}
 
-			if (!$this->polluteScopeWithLoopInitialAssignments) {
-				$finalScope = $finalScope->mergeWith($scope);
+			if ($isIterableAtLeastOnce->no() || $finalScopeResult->isAlwaysTerminating()) {
+				if ($this->polluteScopeWithLoopInitialAssignments) {
+					$finalScope = $initScope;
+				} else {
+					$finalScope = $scope;
+				}
+
+			} elseif ($isIterableAtLeastOnce->maybe()) {
+				if ($this->polluteScopeWithLoopInitialAssignments) {
+					$finalScope = $finalScope->mergeWith($initScope);
+				} else {
+					$finalScope = $finalScope->mergeWith($scope);
+				}
+			} else {
+				if (!$this->polluteScopeWithLoopInitialAssignments) {
+					$finalScope = $finalScope->mergeWith($scope);
+				}
 			}
 
 			return new StatementResult(
diff --git a/tests/PHPStan/Analyser/data/for-loop-i-type.php b/tests/PHPStan/Analyser/data/for-loop-i-type.php
index 692c30585d..f3cf8a1a8e 100644
--- a/tests/PHPStan/Analyser/data/for-loop-i-type.php
+++ b/tests/PHPStan/Analyser/data/for-loop-i-type.php
@@ -8,11 +8,14 @@ class Foo
 {
 
 	public function doBar() {
+		$foo = null;
 		for($i = 1; $i < 50; $i++) {
+			$foo = new \stdClass();
 			assertType('int<1, 49>', $i);
 		}
 
 		assertType('50', $i);
+		assertType(\stdClass::class, $foo);
 
 		for($i = 50; $i > 0; $i--) {
 			assertType('int<1, max>', $i); // could be int<1, 50>
@@ -21,6 +24,28 @@ public function doBar() {
 		assertType('0', $i);
 	}
 
+	public function doCount(array $a) {
+		$foo = null;
+		for($i = 1; $i < count($a); $i++) {
+			$foo = new \stdClass();
+			assertType('int<1, max>', $i);
+		}
+
+		assertType('int<1, max>', $i);
+		assertType(\stdClass::class . '|null', $foo);
+	}
+
+	public function doCount2() {
+		$foo = null;
+		for($i = 1; $i < count([]); $i++) {
+			$foo = new \stdClass();
+			assertType('string', $i); // should be *NEVER*
+		}
+
+		assertType('1', $i);
+		assertType('null', $foo);
+	}
+
 	public function doBaz() {
 		for($i = 1; $i < 50; $i += 2) {
 			assertType('int<1, 49>', $i);