From 481cffa093d00799a4aefc192013eb8773ce5a61 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Tamarelle?=
 <jerome.tamarelle@mongodb.com>
Date: Wed, 26 Feb 2025 15:48:15 +0100
Subject: [PATCH] Pre-init all sub-encoders of the builder encoder (#1613)

Use weak reference to prevent circular reference
Avoid having the same encoder initialized multiple times
---
 src/Builder/BuilderEncoder.php          | 50 +++++++++++--------------
 src/Builder/Encoder/QueryEncoder.php    |  2 +-
 src/Builder/Encoder/RecursiveEncode.php | 16 ++++++--
 3 files changed, 35 insertions(+), 33 deletions(-)

diff --git a/src/Builder/BuilderEncoder.php b/src/Builder/BuilderEncoder.php
index 695d3fb14..833409b07 100644
--- a/src/Builder/BuilderEncoder.php
+++ b/src/Builder/BuilderEncoder.php
@@ -29,10 +29,10 @@
 use MongoDB\Codec\Encoder;
 use MongoDB\Exception\UnsupportedValueException;
 use stdClass;
+use WeakReference;
 
 use function array_key_exists;
 use function is_object;
-use function is_string;
 
 /** @template-implements Encoder<Type|stdClass|array|string|int, Pipeline|StageInterface|ExpressionInterface|QueryInterface> */
 final class BuilderEncoder implements Encoder
@@ -40,25 +40,28 @@ final class BuilderEncoder implements Encoder
     /** @template-use EncodeIfSupported<Type|stdClass|array|string|int, Pipeline|StageInterface|ExpressionInterface|QueryInterface> */
     use EncodeIfSupported;
 
-    /** @var array<class-string, class-string<Encoder>> */
-    private array $defaultEncoders = [
-        Pipeline::class => PipelineEncoder::class,
-        Variable::class => VariableEncoder::class,
-        DictionaryInterface::class => DictionaryEncoder::class,
-        FieldPathInterface::class => FieldPathEncoder::class,
-        CombinedFieldQuery::class => CombinedFieldQueryEncoder::class,
-        QueryObject::class => QueryEncoder::class,
-        OutputWindow::class => OutputWindowEncoder::class,
-        OperatorInterface::class => OperatorEncoder::class,
-        DateTimeInterface::class => DateTimeEncoder::class,
-    ];
+    /** @var array<class-string, Encoder> */
+    private array $encoders;
 
     /** @var array<class-string, Encoder|null> */
     private array $cachedEncoders = [];
 
-    /** @param array<class-string, Encoder> $customEncoders */
-    public function __construct(private readonly array $customEncoders = [])
+    /** @param array<class-string, Encoder> $encoders */
+    public function __construct(array $encoders = [])
     {
+        $self = WeakReference::create($this);
+
+        $this->encoders = $encoders + [
+            Pipeline::class => new PipelineEncoder($self),
+            Variable::class => new VariableEncoder(),
+            DictionaryInterface::class => new DictionaryEncoder(),
+            FieldPathInterface::class => new FieldPathEncoder(),
+            CombinedFieldQuery::class => new CombinedFieldQueryEncoder($self),
+            QueryObject::class => new QueryEncoder($self),
+            OutputWindow::class => new OutputWindowEncoder($self),
+            OperatorInterface::class => new OperatorEncoder($self),
+            DateTimeInterface::class => new DateTimeEncoder(),
+        ];
     }
 
     /** @psalm-assert-if-true object $value */
@@ -89,25 +92,14 @@ private function getEncoderFor(object $value): Encoder|null
             return $this->cachedEncoders[$valueClass];
         }
 
-        $encoderList = $this->customEncoders + $this->defaultEncoders;
-
         // First attempt: match class name exactly
-        if (isset($encoderList[$valueClass])) {
-            $encoder = $encoderList[$valueClass];
-            if (is_string($encoder)) {
-                $encoder = new $encoder($this);
-            }
-
-            return $this->cachedEncoders[$valueClass] = $encoder;
+        if (isset($this->encoders[$valueClass])) {
+            return $this->cachedEncoders[$valueClass] = $this->encoders[$valueClass];
         }
 
         // Second attempt: catch child classes
-        foreach ($encoderList as $className => $encoder) {
+        foreach ($this->encoders as $className => $encoder) {
             if ($value instanceof $className) {
-                if (is_string($encoder)) {
-                    $encoder = new $encoder($this);
-                }
-
                 return $this->cachedEncoders[$valueClass] = $encoder;
             }
         }
diff --git a/src/Builder/Encoder/QueryEncoder.php b/src/Builder/Encoder/QueryEncoder.php
index 6a43f37c0..19f21897c 100644
--- a/src/Builder/Encoder/QueryEncoder.php
+++ b/src/Builder/Encoder/QueryEncoder.php
@@ -53,7 +53,7 @@ public function encode(mixed $value): stdClass
                     throw new LogicException(sprintf('Duplicate key "%s" in query object', $key));
                 }
 
-                $result->{$key} = $this->encoder->encodeIfSupported($value);
+                $result->{$key} = $this->recursiveEncode($value);
             }
         }
 
diff --git a/src/Builder/Encoder/RecursiveEncode.php b/src/Builder/Encoder/RecursiveEncode.php
index d93fdd998..e83af8a22 100644
--- a/src/Builder/Encoder/RecursiveEncode.php
+++ b/src/Builder/Encoder/RecursiveEncode.php
@@ -4,15 +4,18 @@
 
 namespace MongoDB\Builder\Encoder;
 
-use MongoDB\Builder\BuilderEncoder;
+use MongoDB\Codec\Encoder;
 use stdClass;
+use WeakReference;
 
 use function get_object_vars;
 use function is_array;
 
+/** @internal */
 trait RecursiveEncode
 {
-    final public function __construct(protected readonly BuilderEncoder $encoder)
+    /** @param WeakReference<Encoder> $encoder */
+    final public function __construct(private readonly WeakReference $encoder)
     {
     }
 
@@ -44,6 +47,13 @@ private function recursiveEncode(mixed $value): mixed
             return $value;
         }
 
-        return $this->encoder->encodeIfSupported($value);
+        /**
+         * If the BuilderEncoder instance is removed from the memory, the
+         * instances of the classes using this trait will be removed as well.
+         * Therefore, the weak reference will never return null.
+         *
+         * @psalm-suppress PossiblyNullReference
+         */
+        return $this->encoder->get()->encodeIfSupported($value);
     }
 }