Fix bug with prefetching on different nesting levels

src/Parameters/PrefetchDataParameter.php

@@ -52,24 +52,31 @@ public function resolve(object|null $source, array $args, mixed $context, Resolv
         // So we record all of these ->resolve() calls, collect them together and when a value is actually
         // needed, GraphQL calls the callback of Deferred below. That's when we call the prefetch method,
         // already knowing all the requested fields (source-arguments combinations).
-        return new Deferred(function () use ($info, $context, $args, $prefetchBuffer) {
-            if (! $prefetchBuffer->hasResult($args, $info)) {
-                $prefetchResult = $this->computePrefetch($args, $context, $info, $prefetchBuffer);
-
-                $prefetchBuffer->storeResult($prefetchResult, $args, $info);
+        return new Deferred(function () use ($source, $info, $context, $args, $prefetchBuffer) {
+            if (! $prefetchBuffer->hasResult($source)) {
+                $this->computePrefetch($args, $context, $info, $prefetchBuffer);
             }
 
-            return $prefetchResult ?? $prefetchBuffer->getResult($args, $info);
+            $result = $prefetchBuffer->getResult($source);
+            // clear internal storage
+            $prefetchBuffer->purgeResult($source);
+            return $result;
         });
     }
 
     /** @param array<string, mixed> $args */
-    private function computePrefetch(array $args, mixed $context, ResolveInfo $info, PrefetchBuffer $prefetchBuffer): mixed
+    private function computePrefetch(array $args, mixed $context, ResolveInfo $info, PrefetchBuffer $prefetchBuffer): void
     {
         $sources = $prefetchBuffer->getObjectsByArguments($args, $info);
+        $prefetchBuffer->purge($args, $info);
         $toPassPrefetchArgs = QueryField::paramsToArguments($this->fieldName, $this->parameters, null, $args, $context, $info, $this->resolver);
 
-        return ($this->resolver)($sources, ...$toPassPrefetchArgs);
+        $resolvedValues = ($this->resolver)($sources, ...$toPassPrefetchArgs);
+
+        foreach ($sources as $source) {
+            // map results to each source to support old prefetch behavior
+            $prefetchBuffer->storeResult($source, $resolvedValues);
+        }
     }
 
     /** @inheritDoc */

src/PrefetchBuffer.php

@@ -5,8 +5,8 @@
 namespace TheCodingMachine\GraphQLite;
 
 use GraphQL\Type\Definition\ResolveInfo;
+use SplObjectStorage;
 
-use function array_key_exists;
 use function md5;
 use function serialize;
 
@@ -18,8 +18,13 @@ class PrefetchBuffer
     /** @var array<string, array<int, object>> An array of buffered, indexed by hash of arguments. */
     private array $objects = [];
 
-    /** @var array<string, mixed> An array of prefetch method results, indexed by hash of arguments. */
-    private array $results = [];
+    /** @var SplObjectStorage A Storage of prefetch method results, holds source to resolved values. */
+    private SplObjectStorage $results;
+
+    public function __construct()
+    {
+        $this->results = new SplObjectStorage();
+    }
 
     /** @param array<array-key, mixed> $arguments The input arguments passed from GraphQL to the field. */
     public function register(
@@ -66,28 +71,28 @@ public function purge(
         unset($this->objects[$this->computeHash($arguments, $info)]);
     }
 
-    /** @param array<array-key, mixed> $arguments The input arguments passed from GraphQL to the field. */
     public function storeResult(
+        object $source,
         mixed $result,
-        array $arguments,
-        ResolveInfo|null $info = null,
     ): void {
-        $this->results[$this->computeHash($arguments, $info)] = $result;
+        $this->results->offsetSet($source, $result);
     }
 
-    /** @param array<array-key, mixed> $arguments The input arguments passed from GraphQL to the field. */
     public function hasResult(
-        array $arguments,
-        ResolveInfo|null $info = null,
+        object $source,
     ): bool {
-        return array_key_exists($this->computeHash($arguments, $info), $this->results);
+        return $this->results->offsetExists($source);
     }
 
-    /** @param array<array-key, mixed> $arguments The input arguments passed from GraphQL to the field. */
     public function getResult(
-        array $arguments,
-        ResolveInfo|null $info = null,
+        object $source,
     ): mixed {
-        return $this->results[$this->computeHash($arguments, $info)];
+        return $this->results->offsetGet($source);
+    }
+
+    public function purgeResult(
+        object $source,
+    ): void {
+        $this->results->offsetUnset($source);
     }
 }

src/QueryField.php

@@ -4,11 +4,9 @@
 
 namespace TheCodingMachine\GraphQLite;
 
-use GraphQL\Deferred;
+use Closure;
 use GraphQL\Error\ClientAware;
 use GraphQL\Executor\Promise\Adapter\SyncPromise;
-use GraphQL\Executor\Promise\Adapter\SyncPromiseAdapter;
-use GraphQL\Executor\Promise\Promise;
 use GraphQL\Language\AST\FieldDefinitionNode;
 use GraphQL\Type\Definition\FieldDefinition;
 use GraphQL\Type\Definition\ListOfType;
@@ -22,9 +20,6 @@
 use TheCodingMachine\GraphQLite\Parameters\ParameterInterface;
 use TheCodingMachine\GraphQLite\Parameters\SourceParameter;
 
-use function array_filter;
-use function array_map;
-
 /**
  * A GraphQL field that maps to a PHP method automatically.
  *
@@ -79,42 +74,39 @@ public function __construct(
             $callResolver = function (...$args) use ($originalResolver, $source, $resolver) {
                 $result = $resolver($source, ...$args);
 
-                try {
-                    $this->assertReturnType($result);
-                } catch (TypeMismatchRuntimeException $e) {
-                    $e->addInfo($this->name, $originalResolver->toString());
-
-                    throw $e;
-                }
-
-                return $result;
+                return $this->resolveWithPromise($result, $originalResolver);
             };
 
-            $deferred = (bool) array_filter($toPassArgs, static fn (mixed $value) => $value instanceof SyncPromise);
-
             // GraphQL allows deferring resolving the field's value using promises, i.e. they call the resolve
             // function ahead of time for all of the fields (allowing us to gather all calls and do something
             // in batch, like prefetch) and then resolve the promises as needed. To support that for prefetch,
             // we're checking if any of the resolved parameters returned a promise. If they did, we know
             // that the value should also be resolved using a promise, so we're wrapping it in one.
-            return $deferred ? new Deferred(static function () use ($toPassArgs, $callResolver) {
-                $syncPromiseAdapter = new SyncPromiseAdapter();
-
-                // Wait for every deferred parameter.
-                $toPassArgs = array_map(
-                    static fn (mixed $value) => $value instanceof SyncPromise ? $syncPromiseAdapter->wait(new Promise($value, $syncPromiseAdapter)) : $value,
-                    $toPassArgs,
-                );
-
-                return $callResolver(...$toPassArgs);
-            }) : $callResolver(...$toPassArgs);
+            return $this->deferred($toPassArgs, $callResolver);
         };
 
         $config += $additionalConfig;
 
         parent::__construct($config);
     }
 
+    private function resolveWithPromise(mixed $result, ResolverInterface $originalResolver): mixed
+    {
+        if ($result instanceof SyncPromise) {
+            return $result->then(fn ($resolvedValue) => $this->resolveWithPromise($resolvedValue, $originalResolver));
+        }
+
+        try {
+            $this->assertReturnType($result);
+        } catch (TypeMismatchRuntimeException $e) {
+            $e->addInfo($this->name, $originalResolver->toString());
+
+            throw $e;
+        }
+
+        return $result;
+    }
+
     /**
      * This method checks the returned value of the resolver to be sure it matches the documented return type.
      * We are sure the returned value is of the correct type... except if the return type is type-hinted as an array.
@@ -204,4 +196,24 @@ public static function paramsToArguments(string $name, array $parameters, object
 
         return $toPassArgs;
     }
+
+    /**
+     * @param array<int, mixed> $toPassArgs
+     * Create Deferred if any of arguments contain promise
+     */
+    public static function deferred(array $toPassArgs, Closure $callResolver): mixed
+    {
+        $deferredArgument = null;
+        foreach ($toPassArgs as $position => $toPassArg) {
+            if ($toPassArg instanceof SyncPromise) {
+                $deferredArgument = $toPassArg->then(static function ($resolvedValue) use ($toPassArgs, $position, $callResolver) {
+                    $toPassArgs[$position] = $resolvedValue;
+                    return self::deferred($toPassArgs, $callResolver);
+                });
+                break;
+            }
+        }
+
+        return $deferredArgument ?? $callResolver(...$toPassArgs);
+    }
 }

tests/Fixtures/Integration/Controllers/BlogController.php

@@ -0,0 +1,21 @@
+<?php
+
+declare(strict_types=1);
+
+namespace TheCodingMachine\GraphQLite\Fixtures\Integration\Controllers;
+
+use TheCodingMachine\GraphQLite\Annotations\Query;
+use TheCodingMachine\GraphQLite\Fixtures\Integration\Models\Blog;
+
+class BlogController
+{
+    /** @return Blog[] */
+    #[Query]
+    public function getBlogs(): array
+    {
+        return [
+            new Blog(1),
+            new Blog(2),
+        ];
+    }
+}

tests/Fixtures/Integration/Models/Blog.php

@@ -0,0 +1,83 @@
+<?php
+
+declare(strict_types=1);
+
+namespace TheCodingMachine\GraphQLite\Fixtures\Integration\Models;
+
+use GraphQL\Deferred;
+use TheCodingMachine\GraphQLite\Annotations\Field;
+use TheCodingMachine\GraphQLite\Annotations\Prefetch;
+use TheCodingMachine\GraphQLite\Annotations\Type;
+
+#[Type]
+class Blog
+{
+    public function __construct(
+        private readonly int $id,
+    ) {
+    }
+
+    #[Field(outputType: 'ID!')]
+    public function getId(): int
+    {
+        return $this->id;
+    }
+
+    /**
+     * @param Post[][] $prefetchedPosts
+     *
+     * @return Post[]
+     */
+    #[Field]
+    public function getPosts(
+        #[Prefetch('prefetchPosts')]
+        array $prefetchedPosts,
+    ): array {
+        return $prefetchedPosts[$this->id];
+    }
+
+    /** @param Blog[][] $prefetchedSubBlogs */
+    #[Field(outputType: '[Blog!]!')]
+    public function getSubBlogs(
+        #[Prefetch('prefetchSubBlogs')]
+        array $prefetchedSubBlogs,
+    ): Deferred {
+        return new Deferred(fn () => $prefetchedSubBlogs[$this->id]);
+    }
+
+    /**
+     * @param Blog[] $blogs
+     *
+     * @return Post[][]
+     */
+    public static function prefetchPosts(iterable $blogs): array
+    {
+        $posts = [];
+        foreach ($blogs as $blog) {
+            $blogId = $blog->getId();
+            $posts[$blog->getId()] = [
+                new Post('post-' . $blogId . '.1'),
+                new Post('post-' . $blogId . '.2'),
+            ];
+        }
+
+        return $posts;
+    }
+
+    /**
+     * @param Blog[] $blogs
+     *
+     * @return Blog[][]
+     */
+    public static function prefetchSubBlogs(iterable $blogs): array
+    {
+        $subBlogs = [];
+        foreach ($blogs as $blog) {
+            $blogId = $blog->getId();
+            $subBlogId = $blogId * 10;
+            $subBlogs[$blog->id] = [new Blog($subBlogId)];
+        }
+
+        return $subBlogs;
+    }
+}

tests/Fixtures/Integration/Models/Comment.php

@@ -0,0 +1,23 @@
+<?php
+
+declare(strict_types=1);
+
+namespace TheCodingMachine\GraphQLite\Fixtures\Integration\Models;
+
+use TheCodingMachine\GraphQLite\Annotations\Field;
+use TheCodingMachine\GraphQLite\Annotations\Type;
+
+#[Type]
+class Comment
+{
+    public function __construct(
+        private readonly string $text,
+    ) {
+    }
+
+    #[Field]
+    public function getText(): string
+    {
+        return $this->text;
+    }
+}

tests/Fixtures/Integration/Types/CompanyType.php

@@ -13,22 +13,27 @@
 #[ExtendType(class:Company::class)]
 class CompanyType
 {
-
     #[Field]
     public function getName(Company $company): string
     {
         return $company->name;
     }
 
+    /** @param Contact[] $contacts */
     #[Field]
     public function getContact(
         Company $company,
         #[Prefetch('prefetchContacts')]
-        array $contacts
-    ): ?Contact {
+        array $contacts,
+    ): Contact|null {
         return $contacts[$company->name] ?? null;
     }
 
+    /**
+     * @param Company[] $companies
+     *
+     * @return Contact[]
+     */
     public static function prefetchContacts(array $companies): array
     {
         $contacts = [];