Be smarter about new array keys after assignment

Author: ondrejmirtes

src/Type/ArrayType.php

@@ -28,6 +28,7 @@
 use PHPStan\Type\Traits\UndecidedBooleanTypeTrait;
 use PHPStan\Type\Traits\UndecidedComparisonTypeTrait;
 use function array_merge;
+use function count;
 use function sprintf;
 
 /** @api */
@@ -419,7 +420,31 @@ public function getOffsetValueType(Type $offsetType): Type
 	public function setOffsetValueType(?Type $offsetType, Type $valueType, bool $unionValues = true): Type
 	{
 		if ($offsetType === null) {
-			$offsetType = new IntegerType();
+			$isKeyTypeInteger = $this->keyType->isInteger();
+			if ($isKeyTypeInteger->no()) {
+				$offsetType = new IntegerType();
+			} elseif ($isKeyTypeInteger->yes()) {
+				$offsetType = $this->keyType;
+			} else {
+				$integerTypes = [];
+				TypeTraverser::map($this->keyType, static function (Type $type, callable $traverse) use (&$integerTypes): Type {
+					if ($type instanceof UnionType) {
+						return $traverse($type);
+					}
+
+					$isInteger = $type->isInteger();
+					if ($isInteger->yes()) {
+						$integerTypes[] = $type;
+					}
+
+					return $type;
+				});
+				if (count($integerTypes) === 0) {
+					$offsetType = $this->keyType;
+				} else {
+					$offsetType = TypeCombinator::union(...$integerTypes);
+				}
+			}
 		} else {
 			$offsetType = $offsetType->toArrayKey();
 		}

tests/PHPStan/Analyser/NodeScopeResolverTest.php

@@ -1219,6 +1219,7 @@ public function dataFileAsserts(): iterable
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/trait-type-alias.php');
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/bug-8609.php');
 		yield from $this->gatherAssertTypes(__DIR__ . '/../Rules/PhpDoc/data/bug-8609-function.php');
+		yield from $this->gatherAssertTypes(__DIR__ . '/data/bug-9131.php');
 	}
 
 	/**

tests/PHPStan/Analyser/data/bug-9131.php

@@ -0,0 +1,37 @@
+<?php
+
+namespace Bug9131TypeInference;
+
+use function PHPStan\Testing\assertType;
+
+class Foo
+{
+
+	/**
+	 * @param string[] $a
+	 * @param array<string, string> $b
+	 * @param array<int<0, max>, string> $c
+	 * @param array<int<0, max>|string, string> $d
+	 * @return void
+	 */
+	public function doFoo(
+		array $a,
+		array $b,
+		array $c,
+		array $d
+	): void
+	{
+		$a[] = 'foo';
+		assertType('non-empty-array<string>', $a);
+
+		$b[] = 'foo';
+		assertType('non-empty-array<int|string, string>', $b);
+
+		$c[] = 'foo';
+		assertType('non-empty-array<int<0, max>, string>', $c);
+
+		$d[] = 'foo';
+		assertType('non-empty-array<int<0, max>|string, string>', $d);
+	}
+
+}

tests/PHPStan/Rules/Properties/TypesAssignedToPropertiesRuleTest.php

@@ -526,4 +526,10 @@ public function testBug7789(): void
 		$this->analyse([__DIR__ . '/data/bug-7789.php'], []);
 	}
 
+	public function testBug9131(): void
+	{
+		$this->checkExplicitMixed = true;
+		$this->analyse([__DIR__ . '/data/bug-9131.php'], []);
+	}
+
 }

tests/PHPStan/Rules/Properties/data/bug-9131.php

@@ -0,0 +1,13 @@
+<?php
+
+namespace Bug9131;
+
+class A
+{
+	/** @var array<int<0, max>, string> */
+	public array $l = [];
+
+	public function add(string $s): void {
+		$this->l[] = $s;
+	}
+}