Support mysqli_result->fetch_assoc() and mysqli_result->fetch_all()

config/extensions.neon

@@ -69,6 +69,11 @@ services:
         tags:
         	- phpstan.broker.dynamicFunctionReturnTypeExtension
 
+    -
+        class: staabm\PHPStanDba\Extensions\MysqliResultDynamicReturnTypeExtension
+        tags:
+        	- phpstan.broker.dynamicMethodReturnTypeExtension
+
     -
         class: staabm\PHPStanDba\Extensions\MysqliQueryDynamicReturnTypeExtension
         tags:

src/Extensions/MysqliQueryDynamicReturnTypeExtension.php

@@ -130,10 +130,7 @@ private function inferResultType(Expr $queryExpr, Scope $scope): ?Type
                 return null;
             }
 
-            $resultObjectType = new MysqliResultObjectType(\mysqli_result::class);
-            $resultObjectType->setRowType($resultType);
-
-            $objects[] = $resultObjectType;
+            $objects[] = MysqliResultObjectType::newWithRowType($resultType);
         }
 
         if (0 === \count($objects)) {

src/Extensions/MysqliResultDynamicReturnTypeExtension.php

@@ -0,0 +1,164 @@
+<?php
+
+declare(strict_types=1);
+
+namespace staabm\PHPStanDba\Extensions;
+
+use mysqli_result;
+use PhpParser\Node\Expr\MethodCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Type\Accessory\AccessoryArrayListType;
+use PHPStan\Type\ArrayType;
+use PHPStan\Type\Constant\ConstantArrayType;
+use PHPStan\Type\Constant\ConstantArrayTypeBuilder;
+use PHPStan\Type\Constant\ConstantIntegerType;
+use PHPStan\Type\DynamicMethodReturnTypeExtension;
+use PHPStan\Type\IntegerType;
+use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
+use PHPStan\Type\UnionType;
+use staabm\PHPStanDba\MysqliReflection\MysqliResultObjectType;
+
+final class MysqliResultDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
+{
+    public function getClass(): string
+    {
+        return mysqli_result::class;
+    }
+
+    public function isMethodSupported(MethodReflection $methodReflection): bool
+    {
+        return \in_array(strtolower($methodReflection->getName()), ['fetch_assoc', 'fetch_all'], true);
+    }
+
+    public function getTypeFromMethodCall(MethodReflection $methodReflection, MethodCall $methodCall, Scope $scope): ?Type
+    {
+        $resultType = $scope->getType($methodCall->var);
+        $methodName = strtolower($methodReflection->getName());
+        if ('fetch_assoc' === $methodName) {
+            if (count($methodCall->getArgs()) === 0) {
+                return $this->getRowType($resultType);
+            }
+
+            return null;
+        }
+
+        $modes = [\MYSQLI_NUM];
+        if (count($methodCall->getArgs()) > 0) {
+            $modes = $this->getModes($scope->getType($methodCall->getArgs()[0]->value));
+            if ($modes === null) {
+                return null;
+            }
+        }
+
+        $rowType = $this->getRowType($resultType);
+        if ($rowType === null) {
+            return null;
+        }
+
+        $rows = [];
+        foreach ($modes as $mode) {
+            $modeRowType = $this->getRowTypeForMode($rowType, $mode);
+            if ($modeRowType === null) {
+                return null;
+            }
+
+            $rows[] = $modeRowType;
+        }
+
+        return TypeCombinator::intersect(
+            new ArrayType(new IntegerType(), TypeCombinator::union(...$rows)),
+            new AccessoryArrayListType()
+        );
+    }
+
+    private function getRowType(Type $resultType): ?Type
+    {
+        if ($resultType instanceof MysqliResultObjectType) {
+            return $resultType->getRowType();
+        }
+
+        if (! $resultType instanceof UnionType) {
+            return null;
+        }
+
+        $rowTypes = [];
+        foreach ($resultType->getTypes() as $type) {
+            $rowType = $this->getRowType($type);
+            if ($rowType === null) {
+                return null;
+            }
+
+            $rowTypes[] = $rowType;
+        }
+
+        return TypeCombinator::union(...$rowTypes);
+    }
+
+    /**
+     * @return list<int>|null
+     */
+    private function getModes(Type $modeType): ?array
+    {
+        if ($modeType instanceof ConstantIntegerType) {
+            return [$modeType->getValue()];
+        }
+
+        if (! $modeType instanceof UnionType) {
+            return null;
+        }
+
+        $modes = [];
+        foreach ($modeType->getTypes() as $type) {
+            $typeModes = $this->getModes($type);
+            if ($typeModes === null) {
+                return null;
+            }
+
+            foreach ($typeModes as $mode) {
+                $modes[] = $mode;
+            }
+        }
+
+        return $modes;
+    }
+
+    private function getRowTypeForMode(Type $rowType, int $mode): ?Type
+    {
+        if (! \in_array($mode, [\MYSQLI_NUM, \MYSQLI_ASSOC, \MYSQLI_BOTH], true)) {
+            return null;
+        }
+
+        $rowTypes = [];
+        foreach ($rowType->getConstantArrays() as $arrayType) {
+            $rowTypes[] = $this->convertRowType($arrayType, $mode);
+        }
+
+        if (count($rowTypes) === 0) {
+            return null;
+        }
+
+        return TypeCombinator::union(...$rowTypes);
+    }
+
+    private function convertRowType(ConstantArrayType $rowType, int $mode): Type
+    {
+        $builder = ConstantArrayTypeBuilder::createEmpty();
+        $numericIndex = 0;
+
+        foreach ($rowType->getKeyTypes() as $index => $keyType) {
+            $valueType = $rowType->getValueTypes()[$index];
+            if ($mode === \MYSQLI_ASSOC || $mode === \MYSQLI_BOTH) {
+                $builder->setOffsetValueType($keyType, $valueType);
+            }
+
+            if ($mode === \MYSQLI_NUM || $mode === \MYSQLI_BOTH) {
+                $builder->setOffsetValueType(new ConstantIntegerType($numericIndex), $valueType);
+                $numericIndex++;
+            }
+        }
+
+        return $builder->getArray();
+    }
+}

src/MysqliReflection/MysqliResultObjectType.php

@@ -4,6 +4,7 @@
 
 namespace staabm\PHPStanDba\MysqliReflection;
 
+use PHPStan\ShouldNotHappenException;
 use PHPStan\Type\IsSuperTypeOfResult;
 use PHPStan\Type\ObjectType;
 use PHPStan\Type\Type;
@@ -12,18 +13,25 @@ final class MysqliResultObjectType extends ObjectType
 {
     private ?Type $rowType = null;
 
-    public function setRowType(Type $rowType): void
+    public static function newWithRowType(Type $rowType): self
     {
-        $this->rowType = $rowType;
+        $new = new self(\mysqli_result::class);
+        $new->rowType = $rowType;
+        return $new;
     }
 
-    public function getIterableValueType(): \PHPStan\Type\Type
+    public function getRowType(): Type
     {
-        if ($this->rowType !== null) {
-            return $this->rowType;
+        if ($this->rowType === null) {
+            throw new ShouldNotHappenException();
         }
 
-        return parent::getIterableValueType();
+        return $this->rowType;
+    }
+
+    public function getIterableValueType(): Type
+    {
+        return $this->getRowType();
     }
 
     // differentiate objects based on the local properties,

tests/default/data/mysqli.php

@@ -57,4 +57,53 @@ public function unionResult(mysqli $mysqli)
             }
         }
     }
+
+    public function mysqliResultfetch(mysqli $mysqli)
+    {
+        $result = mysqli_query($mysqli, 'SELECT adaid FROM ada');
+        assertType('array{adaid: int<-32768, 32767>}', $result->fetch_assoc());
+
+        $result = mysqli_query($mysqli, 'SELECT adaid FROM ada');
+        assertType('list<array{int<-32768, 32767>}>', $result->fetch_all());
+        assertType('list<array{adaid: int<-32768, 32767>}>', $result->fetch_all(MYSQLI_ASSOC));
+        assertType('list<array{int<-32768, 32767>}>', $result->fetch_all(MYSQLI_NUM));
+        assertType('list<array{adaid: int<-32768, 32767>, 0: int<-32768, 32767>}>', $result->fetch_all(MYSQLI_BOTH));
+    }
+
+    public function fetchAllColumns(mysqli $mysqli)
+    {
+        $result = $mysqli->query('SELECT email, adaid FROM ada');
+        assertType('list<array{string, int<-32768, 32767>}>', $result->fetch_all());
+        assertType('list<array{email: string, adaid: int<-32768, 32767>}>', $result->fetch_all(MYSQLI_ASSOC));
+        assertType('list<array{string, int<-32768, 32767>}>', $result->fetch_all(MYSQLI_NUM));
+        assertType('list<array{email: string, 0: string, adaid: int<-32768, 32767>, 1: int<-32768, 32767>}>', $result->fetch_all(MYSQLI_BOTH));
+        assertType('array{email: string, adaid: int<-32768, 32767>}', $result->fetch_assoc());
+        foreach ($result as $row) {
+            assertType('array{email: string, adaid: int<-32768, 32767>}', $row);
+        }
+    }
+
+    public function fetchAllUnionResult(mysqli $mysqli)
+    {
+        $queries = ['SELECT adaid FROM ada', 'SELECT email FROM ada'];
+
+        foreach ($queries as $query) {
+            $result = $mysqli->query($query);
+            assertType('list<array{int<-32768, 32767>}|array{string}>', $result->fetch_all());
+            assertType('list<array{adaid: int<-32768, 32767>}|array{email: string}>', $result->fetch_all(MYSQLI_ASSOC));
+            assertType('list<array{int<-32768, 32767>}|array{string}>', $result->fetch_all(MYSQLI_NUM));
+            assertType('list<array{adaid: int<-32768, 32767>, 0: int<-32768, 32767>}|array{email: string, 0: string}>', $result->fetch_all(MYSQLI_BOTH));
+            assertType('array{adaid: int<-32768, 32767>}|array{email: string}', $result->fetch_assoc());
+        }
+    }
+
+    public function fetchAllUnionMode(mysqli $mysqli)
+    {
+        $result = $mysqli->query('SELECT adaid FROM ada');
+        $mode = rand(0, 1) ? MYSQLI_ASSOC : MYSQLI_NUM;
+        assertType('list<array{adaid: int<-32768, 32767>}|array{int<-32768, 32767>}>', $result->fetch_all($mode));
+        assertType('list<array{adaid: int<-32768, 32767>, 0: int<-32768, 32767>}>', $result->fetch_all(MYSQLI_ASSOC | MYSQLI_NUM));
+        assertType('list<array{adaid: int<-32768, 32767>}>', $result->fetch_all(MYSQLI_ASSOC));
+    }
+
 }