Fix ArrayColumnFunctionReturnTypeExtension producing array key types

Fixes phpstan/phpstan#6682

src/Type/Php/ArrayColumnFunctionReturnTypeExtension.php

@@ -90,7 +90,7 @@ private function handleAnyArray(Type $arrayType, Type $columnType, ?Type $indexT
 			$returnKeyType = new IntegerType();
 		}
 
-		$returnType = new ArrayType($returnKeyType, $returnValueType);
+		$returnType = new ArrayType($this->castToArrayKeyType($returnKeyType), $returnValueType);
 
 		if ($iterableAtLeastOnce->yes()) {
 			$returnType = TypeCombinator::intersect($returnType, new NonEmptyArrayType());
@@ -125,6 +125,9 @@ private function handleConstantArray(ConstantArrayType $arrayType, Type $columnT
 				$keyType = null;
 			}
 
+			if ($keyType !== null) {
+				$keyType = $this->castToArrayKeyType($keyType);
+			}
 			$builder->setOffsetValueType($keyType, $valueType);
 		}
 
@@ -180,4 +183,19 @@ private function getOffsetOrProperty(Type $type, Type $offsetOrProperty, Scope $
 		return TypeCombinator::union(...$returnTypes);
 	}
 
+	private function castToArrayKeyType(Type $type): Type
+	{
+		$isArray = $type->isArray();
+		if ($isArray->yes()) {
+			return new IntegerType();
+		}
+		if ($isArray->no()) {
+			return ArrayType::castToArrayKeyType($type);
+		}
+		return TypeCombinator::union(
+			ArrayType::castToArrayKeyType(TypeCombinator::remove($type, new ArrayType(new MixedType(), new MixedType()))),
+			new IntegerType(),
+		);
+	}
+
 }

tests/PHPStan/Analyser/NodeScopeResolverTest.php

@@ -750,6 +750,7 @@ public function dataFileAsserts(): iterable
 
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/bug-6699.php');
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/bug-6715.php');
+		yield from $this->gatherAssertTypes(__DIR__ . '/data/bug-6682.php');
 	}
 
 	/**

tests/PHPStan/Analyser/data/array-column.php

@@ -20,6 +20,17 @@ function testArrays(array $array): void
 	assertType('array{}', array_column($array, 'column'));
 	assertType('array{}', array_column($array, 'column', 'key'));
 	assertType('array{}', array_column($array, null, 'key'));
+
+	/** @var array<int, array<string, float>> $array */
+	assertType('array<int, float>', array_column($array, 'column', 'key'));
+	/** @var array<int, array<string, bool>> $array */
+	assertType('array<int, bool>', array_column($array, 'column', 'key'));
+	/** @var array<int, array<string, true>> $array */
+	assertType('array<int, true>', array_column($array, 'column', 'key'));
+	/** @var array<int, array<string, null>> $array */
+	assertType('array<\'\'|int, null>', array_column($array, 'column', 'key'));
+	/** @var array<int, array<string, array>> $array */
+	assertType('array<int, array>', array_column($array, 'column', 'key'));
 }
 
 function testConstantArrays(array $array): void
@@ -54,6 +65,17 @@ function testConstantArrays(array $array): void
 	assertType('non-empty-array<int, string>', array_column($array, 'column'));
 	assertType('non-empty-array<string, string>', array_column($array, 'column', 'key'));
 	assertType('non-empty-array<string, array{column: string, key: string}>', array_column($array, null, 'key'));
+
+	/** @var array<int, array{column: string, key: float}> $array */
+	assertType('array<int, string>', array_column($array, 'column', 'key'));
+	/** @var array<int, array{column: string, key: bool}> $array */
+	assertType('array<0|1, string>', array_column($array, 'column', 'key'));
+	/** @var array<int, array{column: string, key: true}> $array */
+	assertType('array<1, string>', array_column($array, 'column', 'key'));
+	/** @var array<int, array{column: string, key: null}> $array */
+	assertType('array<\'\', string>', array_column($array, 'column', 'key'));
+	/** @var array<int, array{column: string, key: array}> $array */
+	assertType('array<int, string>', array_column($array, 'column', 'key'));
 }
 
 function testImprecise(array $array): void {
@@ -84,17 +106,17 @@ function testObjects(array $array): void {
 	assertType('array<string, DOMElement>', array_column($array, null, 'tagName'));
 	assertType('array<int, mixed>', array_column($array, 'foo'));
 	assertType('array<string, mixed>', array_column($array, 'foo', 'tagName'));
-	assertType('array<string>', array_column($array, 'nodeName', 'foo'));
-	assertType('array<DOMElement>', array_column($array, null, 'foo'));
+	assertType('array<int|string, string>', array_column($array, 'nodeName', 'foo'));
+	assertType('array<int|string, DOMElement>', array_column($array, null, 'foo'));
 
 	/** @var non-empty-array<int, DOMElement> $array */
 	assertType('non-empty-array<int, string>', array_column($array, 'nodeName'));
 	assertType('non-empty-array<string, string>', array_column($array, 'nodeName', 'tagName'));
 	assertType('non-empty-array<string, DOMElement>', array_column($array, null, 'tagName'));
 	assertType('array<int, mixed>', array_column($array, 'foo'));
 	assertType('array<string, mixed>', array_column($array, 'foo', 'tagName'));
-	assertType('non-empty-array<string>', array_column($array, 'nodeName', 'foo'));
-	assertType('non-empty-array<DOMElement>', array_column($array, null, 'foo'));
+	assertType('non-empty-array<int|string, string>', array_column($array, 'nodeName', 'foo'));
+	assertType('non-empty-array<int|string, DOMElement>', array_column($array, null, 'foo'));
 
 	/** @var array{DOMElement} $array */
 	assertType('array{string}', array_column($array, 'nodeName'));

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

@@ -0,0 +1,19 @@
+<?php
+
+declare(strict_types = 1);
+
+namespace Bug6682;
+
+use function PHPStan\Testing\assertType;
+
+class Example
+{
+    /**
+     * @param array<int, array<string, string|array<string,string>|null>> $data
+     */
+    public function __construct(array $data)
+    {
+        $x = array_column($data, null, 'type');
+		assertType('array<int|string, array<string, array<string, string>|string|null>>', $x);
+    }
+}