From 100907b61fe2da1f82650d4aed715134096ee6e4 Mon Sep 17 00:00:00 2001
From: RobChett <robchett@gmail.com>
Date: Mon, 15 May 2023 01:17:27 +0100
Subject: [PATCH] Improve return types of pow()

---
 .../Provider/FunctionReturnTypeProvider.php   |  2 +
 .../PowReturnTypeProvider.php                 | 92 +++++++++++++++++++
 .../PowReturnTypeProviderTest.php             | 42 +++++++++
 3 files changed, 136 insertions(+)
 create mode 100644 src/Psalm/Internal/Provider/ReturnTypeProvider/PowReturnTypeProvider.php
 create mode 100644 tests/ReturnTypeProvider/PowReturnTypeProviderTest.php

diff --git a/src/Psalm/Internal/Provider/FunctionReturnTypeProvider.php b/src/Psalm/Internal/Provider/FunctionReturnTypeProvider.php
index ad1827b52c7..de1e27d7a5b 100644
--- a/src/Psalm/Internal/Provider/FunctionReturnTypeProvider.php
+++ b/src/Psalm/Internal/Provider/FunctionReturnTypeProvider.php
@@ -36,6 +36,7 @@
 use Psalm\Internal\Provider\ReturnTypeProvider\MinMaxReturnTypeProvider;
 use Psalm\Internal\Provider\ReturnTypeProvider\MktimeReturnTypeProvider;
 use Psalm\Internal\Provider\ReturnTypeProvider\ParseUrlReturnTypeProvider;
+use Psalm\Internal\Provider\ReturnTypeProvider\PowReturnTypeProvider;
 use Psalm\Internal\Provider\ReturnTypeProvider\RandReturnTypeProvider;
 use Psalm\Internal\Provider\ReturnTypeProvider\RoundReturnTypeProvider;
 use Psalm\Internal\Provider\ReturnTypeProvider\StrReplaceReturnTypeProvider;
@@ -103,6 +104,7 @@ public function __construct()
         $this->registerClass(RoundReturnTypeProvider::class);
         $this->registerClass(MbInternalEncodingReturnTypeProvider::class);
         $this->registerClass(DateReturnTypeProvider::class);
+        $this->registerClass(PowReturnTypeProvider::class);
     }
 
     /**
diff --git a/src/Psalm/Internal/Provider/ReturnTypeProvider/PowReturnTypeProvider.php b/src/Psalm/Internal/Provider/ReturnTypeProvider/PowReturnTypeProvider.php
new file mode 100644
index 00000000000..bffd07fcd3d
--- /dev/null
+++ b/src/Psalm/Internal/Provider/ReturnTypeProvider/PowReturnTypeProvider.php
@@ -0,0 +1,92 @@
+<?php
+
+namespace Psalm\Internal\Provider\ReturnTypeProvider;
+
+use Psalm\Plugin\EventHandler\Event\FunctionReturnTypeProviderEvent;
+use Psalm\Plugin\EventHandler\FunctionReturnTypeProviderInterface;
+use Psalm\Type;
+use Psalm\Type\Atomic\TFloat;
+use Psalm\Type\Atomic\TInt;
+use Psalm\Type\Atomic\TLiteralFloat;
+use Psalm\Type\Atomic\TLiteralInt;
+use Psalm\Type\Union;
+
+use function count;
+
+/**
+ * @internal
+ */
+class PowReturnTypeProvider implements FunctionReturnTypeProviderInterface
+{
+    /**
+     * @return array<lowercase-string>
+     */
+    public static function getFunctionIds(): array
+    {
+        return ['pow'];
+    }
+
+    public static function getFunctionReturnType(FunctionReturnTypeProviderEvent $event): ?Union
+    {
+        $call_args = $event->getCallArgs();
+
+        if (count($call_args) !== 2) {
+            return null;
+        }
+
+        $first_arg = $event->getStatementsSource()->getNodeTypeProvider()->getType($call_args[0]->value);
+        $second_arg = $event->getStatementsSource()->getNodeTypeProvider()->getType($call_args[1]->value);
+
+        $first_arg_literal = null;
+        $first_arg_is_int = false;
+        $first_arg_is_float = false;
+        if ($first_arg !== null && $first_arg->isSingle()) {
+            $first_atomic_type = $first_arg->getSingleAtomic();
+            if ($first_atomic_type instanceof TInt) {
+                $first_arg_is_int = true;
+            } elseif ($first_atomic_type instanceof TFloat) {
+                $first_arg_is_float = true;
+            }
+            if ($first_atomic_type instanceof TLiteralInt
+                || $first_atomic_type instanceof TLiteralFloat
+            ) {
+                $first_arg_literal = $first_atomic_type->value;
+            }
+        }
+
+        $second_arg_literal = null;
+        $second_arg_is_int = false;
+        $second_arg_is_float = false;
+        if ($second_arg !== null && $second_arg->isSingle()) {
+            $second_atomic_type = $second_arg->getSingleAtomic();
+            if ($second_atomic_type instanceof TInt) {
+                $second_arg_is_int = true;
+            } elseif ($second_atomic_type instanceof TFloat) {
+                $second_arg_is_float = true;
+            }
+            if ($second_atomic_type instanceof TLiteralInt
+                || $second_atomic_type instanceof TLiteralFloat
+            ) {
+                $second_arg_literal = $second_atomic_type->value;
+            }
+        }
+
+        if ($first_arg_literal === 0) {
+            return Type::getInt(true, 0);
+        }
+        if ($second_arg_literal === 0) {
+            return Type::getInt(true, 1);
+        }
+        if ($first_arg_literal !== null && $second_arg_literal !== null) {
+            return Type::getFloat($first_arg_literal ** $second_arg_literal);
+        }
+        if ($first_arg_is_int && $second_arg_is_int) {
+            return Type::getInt();
+        }
+        if ($first_arg_is_float || $second_arg_is_float) {
+            return Type::getFloat();
+        }
+
+        return new Union([new TInt(), new TFloat()]);
+    }
+}
diff --git a/tests/ReturnTypeProvider/PowReturnTypeProviderTest.php b/tests/ReturnTypeProvider/PowReturnTypeProviderTest.php
new file mode 100644
index 00000000000..f428518564a
--- /dev/null
+++ b/tests/ReturnTypeProvider/PowReturnTypeProviderTest.php
@@ -0,0 +1,42 @@
+<?php
+
+namespace Psalm\Tests\ReturnTypeProvider;
+
+use Psalm\Tests\TestCase;
+use Psalm\Tests\Traits\ValidCodeAnalysisTestTrait;
+
+class PowReturnTypeProviderTest extends TestCase
+{
+    use ValidCodeAnalysisTestTrait;
+
+    public function providerValidCodeParse(): iterable
+    {
+        yield 'test' => [
+            'code' => '<?php
+                function getInt(): int {
+                    return 1;
+                }
+                function getFloat(): float {
+                    return 1.0;
+                }
+                $int = getInt();
+                $float = getFloat();
+
+                $a = pow($int, $int);
+                $b = pow($int, $float);
+                $c = pow($float, $int);
+                $d = pow(1000, 1000);
+                $e = pow(0, 1000);
+                $f = pow(1000, 0);
+            ',
+            'assertions' => [
+                '$a===' => 'int',
+                '$b===' => 'float',
+                '$c===' => 'float',
+                '$d===' => 'float(INF)',
+                '$e===' => '0',
+                '$f===' => '1',
+            ],
+        ];
+    }
+}