From b54e909986751319ca8a43744c12a1deb84de321 Mon Sep 17 00:00:00 2001
From: Roman Lytvynenko <litvinenko95@gmail.com>
Date: Wed, 22 May 2024 09:20:58 +0300
Subject: [PATCH] method calls inference improvements

---
 src/Infer/Handler/ArrayItemHandler.php       |  2 +-
 src/Infer/Services/ReferenceTypeResolver.php | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/Infer/Handler/ArrayItemHandler.php b/src/Infer/Handler/ArrayItemHandler.php
index 60bdcad5..09a5638a 100644
--- a/src/Infer/Handler/ArrayItemHandler.php
+++ b/src/Infer/Handler/ArrayItemHandler.php
@@ -18,7 +18,7 @@ public function leave(Node\Expr\ArrayItem $node, Scope $scope)
         $scope->setType(
             $node,
             new ArrayItemType_(
-                $node->key->value ?? null,
+                $node->key ? ($scope->getType($node->key)->value ?? null) : null, // @todo handle cases when key is something dynamic
                 $scope->getType($node->value),
                 isOptional: false,
                 shouldUnpack: $node->unpack,
diff --git a/src/Infer/Services/ReferenceTypeResolver.php b/src/Infer/Services/ReferenceTypeResolver.php
index 30df1eef..91342e42 100644
--- a/src/Infer/Services/ReferenceTypeResolver.php
+++ b/src/Infer/Services/ReferenceTypeResolver.php
@@ -3,8 +3,10 @@
 namespace Dedoc\Scramble\Infer\Services;
 
 use Dedoc\Scramble\Infer\Analyzer\ClassAnalyzer;
+use Dedoc\Scramble\Infer\Context;
 use Dedoc\Scramble\Infer\Definition\ClassDefinition;
 use Dedoc\Scramble\Infer\Definition\FunctionLikeDefinition;
+use Dedoc\Scramble\Infer\Extensions\Event\MethodCallEvent;
 use Dedoc\Scramble\Infer\Scope\Index;
 use Dedoc\Scramble\Infer\Scope\Scope;
 use Dedoc\Scramble\Support\Type\CallableStringType;
@@ -169,6 +171,22 @@ private function resolveMethodCallReferenceType(Scope $scope, MethodCallReferenc
             throw new \LogicException('Should not happen.');
         }
 
+        // Attempting extensions broker before potentially giving up on type inference
+        if (($calleeType instanceof TemplateType || $calleeType instanceof ObjectType)) {
+            $unwrappedType = $calleeType instanceof TemplateType && $calleeType->is
+                ? $calleeType->is
+                : $calleeType;
+
+            if ($unwrappedType instanceof ObjectType && $returnType = Context::getInstance()->extensionsBroker->getMethodReturnType(new MethodCallEvent(
+                instance: $unwrappedType,
+                name: $type->methodName,
+                scope: $scope,
+                arguments: $type->arguments,
+            ))) {
+                return $returnType;
+            }
+        }
+
         // (#TName).listTableDetails()
         if ($calleeType instanceof TemplateType) {
             // This maybe is not a good idea as it make references bleed into the fully analyzed