Repository: twigphp/Twig (php)

Issue:
[3.9] Cannot assign the output of a custom tag to a variable using `set`
After Twig 3.9.0, it's not possible to capture the output of a custom tag (that uses `echo`) and assign it to a variable using `set`.

```twig
A {% set variable %}{% trans 'value' %}{% endset %} B {{ variable }} C
```

With Twig 3.9.0+, that will print:
```
A value B  C
```

Instead of the expected:
```
A  B value C
```

It works as expected if the custom tag output is not the only content, for example:
```twig
A {% set variable %}some {% trans 'value' %}{% endset %} B {{ variable }} C
```
Outputs:
```
A  B some value C
```

Pull request description:
Fix capturing output from extensions that still use echo
Fix #4039

Patch:
diff --git a/src/Extension/CoreExtension.php b/src/Extension/CoreExtension.php
index e6a72560..c7c6d14f 100644
--- a/src/Extension/CoreExtension.php
+++ b/src/Extension/CoreExtension.php
@@ -329,7 +329,7 @@ final class CoreExtension extends AbstractExtension
         }
 
         if (!\count($values)) {
-            throw new RuntimeError('The "cycle" function does not work on empty arrays.');
+            throw new RuntimeError('The "cycle" function does not work on empty arrays');
         }
 
         return $values[$position % \count($values)];
@@ -478,7 +478,7 @@ final class CoreExtension extends AbstractExtension
      * @param \DateTimeInterface|string|null  $date     A date or null to use the current time
      * @param \DateTimeZone|string|false|null $timezone The target timezone, null to use the default, false to leave unchanged
      *
-     * @return \DateTime|\DateTimeImmutable
+     * @return \DateTimeInterface
      *
      * @internal
      */
@@ -498,7 +498,7 @@ final class CoreExtension extends AbstractExtension
             return false !== $timezone ? $date->setTimezone($timezone) : $date;
         }
 
-        if ($date instanceof \DateTime) {
+        if ($date instanceof \DateTimeInterface) {
             $date = clone $date;
             if (false !== $timezone) {
                 $date->setTimezone($timezone);
diff --git a/src/Node/CaptureNode.php b/src/Node/CaptureNode.php
index 7c187727..534f1d01 100644
--- a/src/Node/CaptureNode.php
+++ b/src/Node/CaptureNode.php
@@ -43,8 +43,9 @@ class CaptureNode extends Node
         $compiler
             ->indent()
             ->subcompile($this->getNode('body'))
+            ->write("return; yield '';\n")
             ->outdent()
-            ->write("})() ?? new \EmptyIterator()")
+            ->write('})()')
         ;
         if ($useYield) {
             $compiler->raw(', false))');
diff --git a/src/Node/CheckSecurityNode.php b/src/Node/CheckSecurityNode.php
index 6e591aad..3525c9f9 100644
--- a/src/Node/CheckSecurityNode.php
+++ b/src/Node/CheckSecurityNode.php
@@ -24,16 +24,11 @@ class CheckSecurityNode extends Node
     private $usedTags;
     private $usedFunctions;
 
-    /**
-     * @param array<string, int> $usedFilters
-     * @param array<string, int> $usedTags
-     * @param array<string, int> $usedFunctions
-     */
     public function __construct(array $usedFilters, array $usedTags, array $usedFunctions)
     {
-        $this->usedFilters = $usedFilters;
-        $this->usedTags = $usedTags;
-        $this->usedFunctions = $usedFunctions;
+        $this->usedFilters = $this->collect($usedFilters);
+        $this->usedTags = $this->collect($usedTags);
+        $this->usedFunctions = $this->collect($usedFunctions);
 
         parent::__construct();
     }
@@ -82,4 +77,14 @@ class CheckSecurityNode extends Node
             ->write("}\n")
         ;
     }
+
+    private function collect(array $used)
+    {
+        $collected = [];
+        foreach ($used as $name => $node) {
+            $collected[$name] = $node instanceof Node ? $node->getTemplateLine() : null;
+        }
+
+        return $collected;
+    }
 }
diff --git a/src/NodeVisitor/SandboxNodeVisitor.php b/src/NodeVisitor/SandboxNodeVisitor.php
index d1108394..1446cee6 100644
--- a/src/NodeVisitor/SandboxNodeVisitor.php
+++ b/src/NodeVisitor/SandboxNodeVisitor.php
@@ -34,11 +34,8 @@ use Twig\Node\SetNode;
 final class SandboxNodeVisitor implements NodeVisitorInterface
 {
     private $inAModule = false;
-    /** @var array<string, int> */
     private $tags;
-    /** @var array<string, int> */
     private $filters;
-    /** @var array<string, int> */
     private $functions;
     private $needsToStringWrap = false;
 
@@ -54,22 +51,22 @@ final class SandboxNodeVisitor implements NodeVisitorInterface
         } elseif ($this->inAModule) {
             // look for tags
             if ($node->getNodeTag() && !isset($this->tags[$node->getNodeTag()])) {
-                $this->tags[$node->getNodeTag()] = $node->getTemplateLine();
+                $this->tags[$node->getNodeTag()] = $node;
             }
 
             // look for filters
             if ($node instanceof FilterExpression && !isset($this->filters[$node->getNode('filter')->getAttribute('value')])) {
-                $this->filters[$node->getNode('filter')->getAttribute('value')] = $node->getTemplateLine();
+                $this->filters[$node->getNode('filter')->getAttribute('value')] = $node;
             }
 
             // look for functions
             if ($node instanceof FunctionExpression && !isset($this->functions[$node->getAttribute('name')])) {
-                $this->functions[$node->getAttribute('name')] = $node->getTemplateLine();
+                $this->functions[$node->getAttribute('name')] = $node;
             }
 
             // the .. operator is equivalent to the range() function
             if ($node instanceof RangeBinary && !isset($this->functions['range'])) {
-                $this->functions['range'] = $node->getTemplateLine();
+                $this->functions['range'] = $node;
             }
 
             if ($node instanceof PrintNode) {


Test patch:
diff --git a/tests/EnvironmentTest.php b/tests/EnvironmentTest.php
index e7eb26e4..57a3bf34 100644
--- a/tests/EnvironmentTest.php
+++ b/tests/EnvironmentTest.php
@@ -414,7 +414,7 @@ class EnvironmentTest extends TestCase
      */
     public function testLegacyEchoingNode()
     {
-        $loader = new ArrayLoader(['echo_bar' => 'A{% set v %}B{% test %}C{% endset %}D{% test %}E{{ v }}F']);
+        $loader = new ArrayLoader(['echo_bar' => 'A{% set v %}B{% test %}C{% endset %}D{% test %}E{{ v }}F{% set w %}{% test %}{% endset %}G{{ w }}H']);
 
         $twig = new Environment($loader);
         $twig->addExtension(new EnvironmentTest_Extension());
@@ -430,7 +430,7 @@ EOF
             );
         }
 
-        $this->assertSame('ADbarEBbarCF', $twig->render('echo_bar'));
+        $this->assertSame('ADbarEBbarCFGbarH', $twig->render('echo_bar'));
     }
 
     protected function getMockLoader($templateName, $templateContent)
diff --git a/tests/Fixtures/functions/cycle_empty.test b/tests/Fixtures/functions/cycle_empty.test
index bb338452..1eaffc8e 100644
--- a/tests/Fixtures/functions/cycle_empty.test
+++ b/tests/Fixtures/functions/cycle_empty.test
@@ -5,4 +5,4 @@
 --DATA--
 return []
 --EXCEPTION--
-Twig\Error\RuntimeError: The "cycle" function does not work on empty arrays in "index.twig" at line 2.
+Twig\Error\RuntimeError: The "cycle" function does not work on empty arrays in "index.twig" at line 2
diff --git a/tests/Node/MacroTest.php b/tests/Node/MacroTest.php
index 88ce9b29..ce51187b 100644
--- a/tests/Node/MacroTest.php
+++ b/tests/Node/MacroTest.php
@@ -60,7 +60,8 @@ public function macro_foo(\$__foo__ = null, \$__bar__ = "Foo", ...\$__varargs__)
 
     return new Markup(implode('', iterator_to_array((function () use (\$context, \$macros, \$blocks) {
         yield "foo";
-    })() ?? new \EmptyIterator(), false)), \$this->env->getCharset());
+        return; yield '';
+    })(), false)), \$this->env->getCharset());
 }
 EOF
             , new Environment(new ArrayLoader()),
diff --git a/tests/Node/SetTest.php b/tests/Node/SetTest.php
index ed6ee9eb..f250b80e 100644
--- a/tests/Node/SetTest.php
+++ b/tests/Node/SetTest.php
@@ -57,7 +57,8 @@ EOF
 // line 1
 \$context["foo"] = ('' === \$tmp = implode('', iterator_to_array((function () use (&\$context, \$macros, \$blocks) {
     yield "foo";
-})() ?? new \EmptyIterator(), false))) ? '' : new Markup(\$tmp, \$this->env->getCharset());
+    return; yield '';
+})(), false))) ? '' : new Markup(\$tmp, \$this->env->getCharset());
 EOF
                 , new Environment(new ArrayLoader()),
             ];
@@ -66,7 +67,8 @@ EOF
 // line 1
 $context["foo"] = ('' === $tmp = \Twig\Extension\CoreExtension::captureOutput((function () use (&$context, $macros, $blocks) {
     yield "foo";
-})() ?? new \EmptyIterator())) ? '' : new Markup($tmp, $this->env->getCharset());
+    return; yield '';
+})())) ? '' : new Markup($tmp, $this->env->getCharset());
 EOF
                 , new Environment(new ArrayLoader()),
             ];