From f614afc0e054e7c728e32dc02d4b98cf0b0f3098 Mon Sep 17 00:00:00 2001
From: David Grudl <david@grudl.com>
Date: Mon, 16 Oct 2023 11:40:00 +0200
Subject: [PATCH] ServiceDefinition: prepends @self to setup immediately

---
 src/DI/Definitions/ServiceDefinition.php    | 23 ++++++++++++---------
 src/DI/Extensions/InjectExtension.php       |  8 +++----
 tests/DI/Definitions.ServiceDefinition.phpt |  6 +++---
 tests/DI/InjectExtension.basic.phpt         |  3 ++-
 4 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/src/DI/Definitions/ServiceDefinition.php b/src/DI/Definitions/ServiceDefinition.php
index b0866db27..ede56644b 100644
--- a/src/DI/Definitions/ServiceDefinition.php
+++ b/src/DI/Definitions/ServiceDefinition.php
@@ -110,10 +110,11 @@ public function setArgument($key, $value)
 	 */
 	public function setSetup(array $setup)
 	{
-		foreach ($setup as $v) {
-			if (!$v instanceof Statement) {
+		foreach ($setup as &$entity) {
+			if (!$entity instanceof Statement) {
 				throw new Nette\InvalidArgumentException('Argument must be Nette\DI\Definitions\Statement[].');
 			}
+			$entity = $this->prependSelf($entity);
 		}
 
 		$this->setup = $setup;
@@ -134,9 +135,10 @@ public function getSetup(): array
 	 */
 	public function addSetup($entity, array $args = [])
 	{
-		$this->setup[] = $entity instanceof Statement
+		$entity = $entity instanceof Statement
 			? $entity
 			: new Statement($entity, $args);
+		$this->setup[] = $this->prependSelf($entity);
 		return $this;
 	}
 
@@ -178,18 +180,19 @@ public function complete(Nette\DI\Resolver $resolver): void
 		$this->creator = $resolver->completeStatement($this->creator);
 
 		foreach ($this->setup as &$setup) {
-			if (
-				is_string($setup->getEntity())
-				&& strpbrk($setup->getEntity(), ':@?\\') === false
-			) { // auto-prepend @self
-				$setup = new Statement([new Reference(Reference::Self), $setup->getEntity()], $setup->arguments);
-			}
-
 			$setup = $resolver->completeStatement($setup, true);
 		}
 	}
 
 
+	private function prependSelf(Statement $setup): Statement
+	{
+		return is_string($setup->getEntity()) && strpbrk($setup->getEntity(), ':@?\\') === false
+			? new Statement([new Reference(Reference::Self), $setup->getEntity()], $setup->arguments)
+			: $setup;
+	}
+
+
 	public function generateMethod(Nette\PhpGenerator\Method $method, Nette\DI\PhpGenerator $generator): void
 	{
 		$code = $generator->formatStatement($this->creator) . ";\n";
diff --git a/src/DI/Extensions/InjectExtension.php b/src/DI/Extensions/InjectExtension.php
index a27ebd4d8..d32082c46 100644
--- a/src/DI/Extensions/InjectExtension.php
+++ b/src/DI/Extensions/InjectExtension.php
@@ -57,9 +57,9 @@ private function updateDefinition(Definitions\ServiceDefinition $def): void
 
 		foreach (self::getInjectProperties($class) as $property => $type) {
 			$builder = $this->getContainerBuilder();
-			$inject = new Definitions\Statement('$' . $property, [Definitions\Reference::fromType((string) $type)]);
+			$inject = new Definitions\Statement(['@self', '$' . $property], [Definitions\Reference::fromType((string) $type)]);
 			foreach ($setups as $key => $setup) {
-				if ($setup->getEntity() === $inject->getEntity()) {
+				if ($setup->getEntity() == $inject->getEntity()) { // intentionally ==
 					$inject = $setup;
 					$builder = null;
 					unset($setups[$key]);
@@ -70,9 +70,9 @@ private function updateDefinition(Definitions\ServiceDefinition $def): void
 		}
 
 		foreach (array_reverse(self::getInjectMethods($class)) as $method) {
-			$inject = new Definitions\Statement($method);
+			$inject = new Definitions\Statement(['@self', $method]);
 			foreach ($setups as $key => $setup) {
-				if ($setup->getEntity() === $inject->getEntity()) {
+				if ($setup->getEntity() == $inject->getEntity()) { // intentionally ==
 					$inject = $setup;
 					unset($setups[$key]);
 				}
diff --git a/tests/DI/Definitions.ServiceDefinition.phpt b/tests/DI/Definitions.ServiceDefinition.phpt
index 06dcb545b..fd933b356 100644
--- a/tests/DI/Definitions.ServiceDefinition.phpt
+++ b/tests/DI/Definitions.ServiceDefinition.phpt
@@ -85,9 +85,9 @@ test('Setting setups with arguments/statements', function () {
 	$def->addSetup(new Statement(stdClass::class, [1, 2]));
 	$def->addSetup(new Statement(stdClass::class, [1, 2]), [99]); // 99 is ignored
 	Assert::equal([
-		new Statement(stdClass::class, [1, 2]),
-		new Statement(stdClass::class, [1, 2]),
-		new Statement(stdClass::class, [1, 2]),
+		new Statement(['@self', stdClass::class], [1, 2]),
+		new Statement(['@self', stdClass::class], [1, 2]),
+		new Statement(['@self', stdClass::class], [1, 2]),
 	], $def->getSetup());
 });
 
diff --git a/tests/DI/InjectExtension.basic.phpt b/tests/DI/InjectExtension.basic.phpt
index 38de3ce1d..fb8172036 100644
--- a/tests/DI/InjectExtension.basic.phpt
+++ b/tests/DI/InjectExtension.basic.phpt
@@ -99,6 +99,7 @@ services:
 		inject: true
 		setup:
 		- injectB(1)
+		- @self::injectC(1)
 		- $e(@\ConcreteDependencyB)
 ');
 
@@ -128,7 +129,7 @@ Assert::equal([
 Assert::equal([
 	new Statement([new Reference('self'), 'injectA']),
 	new Statement([new Reference('self'), 'injectB'], [1]),
-	new Statement([new Reference('self'), 'injectC']),
+	new Statement([new Reference('self'), 'injectC'], [1]),
 	new Statement([new Reference('self'), 'injectD']),
 	new Statement([new Reference('self'), '$e'], [new Reference('b')]),
 	new Statement([new Reference('self'), '$c'], [new Reference('std')]),