From f76505b435110b6f126f2a520ead909ac233cca0 Mon Sep 17 00:00:00 2001
From: adrew <klimichkartorgnusov@gmail.com>
Date: Sun, 11 Sep 2022 20:47:17 +0300
Subject: [PATCH] Add AfterFunctionCallAnalysisEvent and
 AfterMethodCallAnalysisEvent support for Args::getCallArgs

---
 src/Toolkit/Args.php  | 34 +++++++++++++++++++---------------
 src/Toolkit/Types.php |  3 +++
 2 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/src/Toolkit/Args.php b/src/Toolkit/Args.php
index a20011d..bdb90e4 100644
--- a/src/Toolkit/Args.php
+++ b/src/Toolkit/Args.php
@@ -5,18 +5,19 @@
 namespace Fp\PsalmToolkit\Toolkit;
 
 use Fp\Collections\ArrayList;
-use Fp\Collections\NonEmptyArrayList;
 use PhpParser\Node;
 use Fp\Functional\Option\Option;
 use Psalm\CodeLocation;
 use Psalm\NodeTypeProvider;
 use Psalm\Plugin\EventHandler\Event\AfterExpressionAnalysisEvent;
+use Psalm\Plugin\EventHandler\Event\AfterFunctionCallAnalysisEvent;
 use Psalm\Plugin\EventHandler\Event\AfterStatementAnalysisEvent;
 use Psalm\Plugin\EventHandler\Event\AfterMethodCallAnalysisEvent;
 use Psalm\Plugin\EventHandler\Event\FunctionReturnTypeProviderEvent;
 use Psalm\Plugin\EventHandler\Event\MethodReturnTypeProviderEvent;
 use Psalm\StatementsSource;
 use Psalm\Type\Union;
+
 use function Fp\Evidence\proveOf;
 
 final class Args
@@ -29,12 +30,12 @@ public function getArgType(
         NodeTypeProvider |
         AfterStatementAnalysisEvent |
         AfterMethodCallAnalysisEvent |
+        AfterFunctionCallAnalysisEvent |
         MethodReturnTypeProviderEvent |
         FunctionReturnTypeProviderEvent |
         AfterExpressionAnalysisEvent $from,
         Node\Arg | Node\VariadicPlaceholder $for,
-    ): Option
-    {
+    ): Option {
         return proveOf($for, Node\Arg::class)
             ->flatMap(fn(Node\Arg $arg) => PsalmApi::$types->getType($from, $arg->value));
     }
@@ -42,25 +43,28 @@ public function getArgType(
     /**
      * @return Option<ArrayList<CallArg>>
      */
-    public function getCallArgs(MethodReturnTypeProviderEvent | FunctionReturnTypeProviderEvent $from): Option
-    {
+    public function getCallArgs(
+        MethodReturnTypeProviderEvent |
+        FunctionReturnTypeProviderEvent |
+        AfterFunctionCallAnalysisEvent |
+        AfterMethodCallAnalysisEvent $from,
+    ): Option {
         $source = match (true) {
             $from instanceof MethodReturnTypeProviderEvent => $from->getSource(),
             $from instanceof FunctionReturnTypeProviderEvent => $from->getStatementsSource(),
+            $from instanceof AfterFunctionCallAnalysisEvent => $from->getStatementsSource(),
+            $from instanceof AfterMethodCallAnalysisEvent => $from->getStatementsSource(),
+        };
+
+        $args = match (true) {
+            $from instanceof AfterFunctionCallAnalysisEvent => $from->getExpr()->getArgs(),
+            $from instanceof AfterMethodCallAnalysisEvent => $from->getExpr()->getArgs(),
+            default => $from->getCallArgs(),
         };
 
-        return ArrayList::collect($from->getCallArgs())
+        return ArrayList::collect($args)
             ->traverseOption(fn($arg) => $this->getArgType($from, $arg)->map(
                 fn(Union $type) => new CallArg($arg, new CodeLocation($source, $arg), $type)
             ));
     }
-
-    /**
-     * @return Option<NonEmptyArrayList<CallArg>>
-     */
-    public function getNonEmptyCallArgs(MethodReturnTypeProviderEvent | FunctionReturnTypeProviderEvent $from): Option
-    {
-        return $this->getCallArgs($from)
-            ->flatMap(fn($args) => $args->toNonEmptyArrayList());
-    }
 }
diff --git a/src/Toolkit/Types.php b/src/Toolkit/Types.php
index 9ef6cfc..09aafc7 100644
--- a/src/Toolkit/Types.php
+++ b/src/Toolkit/Types.php
@@ -12,6 +12,7 @@
 use Psalm\Internal\Type\TypeExpander;
 use Psalm\NodeTypeProvider;
 use Psalm\Plugin\EventHandler\Event\AfterExpressionAnalysisEvent;
+use Psalm\Plugin\EventHandler\Event\AfterFunctionCallAnalysisEvent;
 use Psalm\Plugin\EventHandler\Event\AfterMethodCallAnalysisEvent;
 use Psalm\Plugin\EventHandler\Event\FunctionReturnTypeProviderEvent;
 use Psalm\Plugin\EventHandler\Event\MethodReturnTypeProviderEvent;
@@ -140,6 +141,7 @@ public function getType(
         StatementsSource |
         NodeTypeProvider |
         AfterMethodCallAnalysisEvent |
+        AfterFunctionCallAnalysisEvent |
         MethodReturnTypeProviderEvent |
         AfterStatementAnalysisEvent |
         FunctionReturnTypeProviderEvent |
@@ -155,6 +157,7 @@ public function getType(
             $from instanceof FunctionReturnTypeProviderEvent => $from->getStatementsSource()->getNodeTypeProvider(),
             $from instanceof AfterExpressionAnalysisEvent => $from->getStatementsSource()->getNodeTypeProvider(),
             $from instanceof AfterMethodCallAnalysisEvent => $from->getStatementsSource()->getNodeTypeProvider(),
+            $from instanceof AfterFunctionCallAnalysisEvent => $from->getStatementsSource()->getNodeTypeProvider(),
         };
 
         return Option::fromNullable($provider->getType($for));