From 7c86ac8ea6f737ebe416178402320996085430e9 Mon Sep 17 00:00:00 2001
From: Ion Bazan <ion.bazan@gmail.com>
Date: Tue, 14 Sep 2021 16:43:43 +0800
Subject: [PATCH] Allow injecting tagged iterator as service locator arguments

---
 .../Compiler/ServiceLocatorTagPass.php        | 20 +++++++++++-------
 .../Compiler/ServiceLocatorTagPassTest.php    | 21 +++++++++++++++++++
 2 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/src/Symfony/Component/DependencyInjection/Compiler/ServiceLocatorTagPass.php b/src/Symfony/Component/DependencyInjection/Compiler/ServiceLocatorTagPass.php
index 0d541183606fd..faa7b57e45232 100644
--- a/src/Symfony/Component/DependencyInjection/Compiler/ServiceLocatorTagPass.php
+++ b/src/Symfony/Component/DependencyInjection/Compiler/ServiceLocatorTagPass.php
@@ -14,6 +14,7 @@
 use Symfony\Component\DependencyInjection\Alias;
 use Symfony\Component\DependencyInjection\Argument\ServiceClosureArgument;
 use Symfony\Component\DependencyInjection\Argument\ServiceLocatorArgument;
+use Symfony\Component\DependencyInjection\Argument\TaggedIteratorArgument;
 use Symfony\Component\DependencyInjection\ContainerBuilder;
 use Symfony\Component\DependencyInjection\Definition;
 use Symfony\Component\DependencyInjection\Exception\InvalidArgumentException;
@@ -47,14 +48,19 @@ protected function processValue($value, bool $isRoot = false)
             $value->setClass(ServiceLocator::class);
         }
 
-        $arguments = $value->getArguments();
-        if (!isset($arguments[0]) || !\is_array($arguments[0])) {
+        $services = $value->getArguments()[0] ?? null;
+
+        if ($services instanceof TaggedIteratorArgument) {
+            $services = $this->findAndSortTaggedServices($services, $this->container);
+        }
+
+        if (!\is_array($services)) {
             throw new InvalidArgumentException(sprintf('Invalid definition for service "%s": an array of references is expected as first argument when the "container.service_locator" tag is set.', $this->currentId));
         }
 
         $i = 0;
 
-        foreach ($arguments[0] as $k => $v) {
+        foreach ($services as $k => $v) {
             if ($v instanceof ServiceClosureArgument) {
                 continue;
             }
@@ -63,18 +69,18 @@ protected function processValue($value, bool $isRoot = false)
             }
 
             if ($i === $k) {
-                unset($arguments[0][$k]);
+                unset($services[$k]);
 
                 $k = (string) $v;
                 ++$i;
             } elseif (\is_int($k)) {
                 $i = null;
             }
-            $arguments[0][$k] = new ServiceClosureArgument($v);
+            $services[$k] = new ServiceClosureArgument($v);
         }
-        ksort($arguments[0]);
+        ksort($services);
 
-        $value->setArguments($arguments);
+        $value->setArgument(0, $services);
 
         $id = '.service_locator.'.ContainerBuilder::hash($value);
 
diff --git a/src/Symfony/Component/DependencyInjection/Tests/Compiler/ServiceLocatorTagPassTest.php b/src/Symfony/Component/DependencyInjection/Tests/Compiler/ServiceLocatorTagPassTest.php
index be63ff751f69b..702137961d789 100644
--- a/src/Symfony/Component/DependencyInjection/Tests/Compiler/ServiceLocatorTagPassTest.php
+++ b/src/Symfony/Component/DependencyInjection/Tests/Compiler/ServiceLocatorTagPassTest.php
@@ -147,6 +147,27 @@ public function testBindingsAreCopied()
         $this->assertInstanceOf(BoundArgument::class, $locator->getBindings()['foo']);
     }
 
+    public function testTaggedServices()
+    {
+        $container = new ContainerBuilder();
+
+        $container->register('bar', TestDefinition1::class)->addTag('test_tag');
+        $container->register('baz', TestDefinition2::class)->addTag('test_tag');
+
+        $container->register('foo', ServiceLocator::class)
+            ->setArguments([new TaggedIteratorArgument('test_tag', null, null, true)])
+            ->addTag('container.service_locator')
+        ;
+
+        (new ServiceLocatorTagPass())->process($container);
+
+        /** @var ServiceLocator $locator */
+        $locator = $container->get('foo');
+
+        $this->assertSame(TestDefinition1::class, \get_class($locator('bar')));
+        $this->assertSame(TestDefinition2::class, \get_class($locator('baz')));
+    }
+
     public function testIndexedByServiceIdWithDecoration()
     {
         $container = new ContainerBuilder();